Session state update not working properly

[shankha98]

The session state update behaviour is confusing.

In this example the session state value of topic should be updated from 'a brave kitten exploring a haunted house' to 'a boy in australia' when the final state is logged.

To Reproduce
Steps to reproduce the behaviour:

1. Run the script (a minimal version of the custom agent example in the docs

# Full runnable code for the StoryFlowAgent example
import logging
import json
from typing import AsyncGenerator
from typing_extensions import override

from google.adk.agents import LlmAgent, BaseAgent
from google.adk.agents.invocation_context import InvocationContext
from google.genai import types
from google.adk.sessions import InMemorySessionService
from google.adk.runners import Runner
from google.adk.events import Event
from pydantic import BaseModel, Field
from dotenv import load_dotenv

# Load environment variables from .env file
load_dotenv()

# --- Constants ---
APP_NAME = "story_app"
USER_ID = "12345"
SESSION_ID = "123344"
GEMINI_2_FLASH = "gemini-2.0-flash-lite"

# --- Configure Logging ---
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)


class DummyAgent1Output(BaseModel):
    test_value: str = Field(
        ..., description="The value of the key 'test_key' in the state session"
    )


class StoryOutput(BaseModel):
    story: str = Field(..., description="The story generated by the agent")


# --- Custom Orchestrator Agent ---
class StoryFlowAgent(BaseAgent):
    """
    Custom agent for a story generation and refinement workflow.

    This agent orchestrates a sequence of LLM agents to generate a story,
    critique it, revise it, check grammar and tone, and potentially
    regenerate the story if the tone is negative.
    """

    # --- Field Declarations for Pydantic ---
    # Declare the agents passed during initialization as class attributes with type hints
    story_generator: LlmAgent

    # model_config allows setting Pydantic configurations if needed, e.g., arbitrary_types_allowed
    model_config = {"arbitrary_types_allowed": True}

    def __init__(
        self,
        name: str,
        story_generator: LlmAgent,
    ):
        """
        Initializes the StoryFlowAgent.

        Args:
            name: The name of the agent.
            story_generator: An LlmAgent to generate the initial story.
            critic: An LlmAgent to critique the story.
            reviser: An LlmAgent to revise the story based on criticism.
            grammar_check: An LlmAgent to check the grammar.
            tone_check: An LlmAgent to analyze the tone.
        """
        # Define the sub_agents list for the framework
        sub_agents_list = [
            story_generator,
        ]

        # Pydantic will validate and assign them based on the class annotations.
        super().__init__(
            name=name,
            story_generator=story_generator,
            sub_agents=sub_agents_list,  # Pass the sub_agents list directly
        )

    @override
    async def _run_async_impl(
        self, ctx: InvocationContext
    ) -> AsyncGenerator[Event, None]:
        """
        Implements the custom orchestration logic for the story workflow.
        Uses the instance attributes assigned by Pydantic (e.g., self.story_generator).
        """
        logger.info(f"[{self.name}] Starting story generation workflow.")

        # 1. Initial Story Generation
        logger.info(f"[{self.name}] Running StoryGenerator...")
        async for event in self.story_generator.run_async(ctx):
            # logger.info(
            #     f"[{self.name}] Event from StoryGenerator: {event.model_dump_json(indent=2, exclude_none=True)}"
            # )
            yield event

        # Check if story was generated before proceeding
        if (
            "current_story" not in ctx.session.state
            or not ctx.session.state["current_story"]
        ):
            logger.error(
                f"[{self.name}] Failed to generate initial story. Aborting workflow."
            )
            return  # Stop processing if initial story failed

        # logger.info(
        #     f"[{self.name}] Story state after generator: {ctx.session.state.get('current_story')}"
        # )

        logger.info(f"[{self.name}] Workflow finished.")


# --- Define the individual LLM agents ---
story_generator = LlmAgent(
    name="StoryGenerator",
    model=GEMINI_2_FLASH,
    instruction="""You are a story writer. Write a short story (around 100 words) about a cat,
    based on the topic provided in session state with key 'topic'""",
    input_schema=None,
    # output_schema=StoryOutput,
    output_key="current_story",  # Key for storing output in session state
)
state_var_picker = LlmAgent(
    name="StoryGenerator",
    model=GEMINI_2_FLASH,
    instruction="""You are a story writer. Write a short story (around 100 words) about a cat,
    based on the topic provided in session state with key 'topic'""",
    input_schema=None,
    # output_schema=StoryOutput,
    output_key="current_story",  # Key for storing output in session state
)

# --- Create the custom agent instance ---
story_flow_agent = StoryFlowAgent(
    name="StoryFlowAgent",
    story_generator=story_generator,
)

# --- Setup Runner and Session ---
session_service = InMemorySessionService()
initial_state = {"topic": "a brave kitten exploring a haunted house"}
session = session_service.create_session(
    app_name=APP_NAME,
    user_id=USER_ID,
    session_id=SESSION_ID,
    state=initial_state,  # Pass initial state here
)

runner = Runner(
    agent=story_flow_agent,  # Pass the custom orchestrator agent
    app_name=APP_NAME,
    session_service=session_service,
)


# --- Function to Interact with the Agent ---
def call_dummy_agent(user_input_topic: str):
    """
    Sends a new topic to the agent (overwriting the initial one if needed)
    and runs the workflow.
    """
    logger.info(f"Initial session state: {session.state}")

    current_session = session_service.get_session(
        app_name=APP_NAME, user_id=USER_ID, session_id=SESSION_ID
    )
    if not current_session:
        logger.error("Session not found!")
        return

    current_session.state.update({"topic": user_input_topic})

    updated_topic = current_session.state.get("topic")

    # Shows updated value
    logger.info(f"Updated session state topic to: {updated_topic}")

    content = types.Content(
        role="user",
        parts=[types.Part(text=f"Generate a story about: {user_input_topic}")],
    )
    events = runner.run(user_id=USER_ID, session_id=SESSION_ID, new_message=content)

    final_response = "No final response captured."
    for event in events:
        if event.is_final_response() and event.content and event.content.parts:
            # logger.info(
            #     f"Potential final response from [{event.author}]: {event.content.parts[0].text}"
            # )
            final_response = event.content.parts[0].text

    print("\n--- Agent Interaction Result ---")
    print("Agent Final Response: ", final_response)

    final_session = session_service.get_session(
        app_name=APP_NAME, user_id=USER_ID, session_id=SESSION_ID
    )
    print("Final Session State:")

    print(json.dumps(final_session.state, indent=2))
    print("-------------------------------\n")

    # Shows initial value
    logger.info("updated topic in session state: %s", final_session.state.get("topic"))


call_dummy_agent("a boy in australia")

The output log

INFO:__main__:Initial session state: {'topic': 'a brave kitten exploring a haunted house'}
INFO:__main__:Updated session state topic to: a boy in australia
INFO:__main__:[StoryFlowAgent] Starting story generation workflow.
INFO:__main__:[StoryFlowAgent] Running StoryGenerator...
INFO:google.adk.models.google_llm:Sending out request, model: gemini-2.0-flash-lite, backend: ml_dev, stream: False
INFO:google.adk.models.google_llm:
LLM Request:
-----------------------------------------------------------
System Instruction:
You are a story writer. Write a short story (around 100 words) about a cat,
        based on the topic provided in session state with key 'topic'

You are an agent. Your internal name is "StoryGenerator".
-----------------------------------------------------------
Contents:
{"parts":[{"text":"Generate a story about: a boy in australia"}],"role":"user"}
-----------------------------------------------------------
Functions:

-----------------------------------------------------------

INFO:google_genai.models:AFC is enabled with max remote calls: 10.
INFO:httpx:HTTP Request: POST https://generativelanguage.googleapis.com/v1beta/models/gemini-2.0-flash-lite:generateContent "HTTP/1.1 200 OK"
INFO:google.adk.models.google_llm:
LLM Response:
-----------------------------------------------------------
Text:
StoryGenerator:
The Australian sun beat down on ten-year-old Finn as he raced across the red dirt. His bare feet kicked up a cloud of ochre dust, chasing after a rogue cricket ball. He lived on a sprawling cattle station, miles from anything, and his best friends were his kelpie, Bluey, and the endless outback. Today, however, he felt a pang of loneliness. He missed the city, the shops, the other kids. Suddenly, a flash of emerald green caught his eye. Nestled under a gum tree, a sleek, black cat with jade eyes stared back at him. A new friend, perhaps? Finn grinned. The outback, he realized, always had a surprise or two.

-----------------------------------------------------------
Function calls:

-----------------------------------------------------------
Raw response:
{"candidates":[{"content":{"parts":[{"text":"StoryGenerator:\nThe Australian sun beat down on ten-year-old Finn as he raced across the red dirt. His bare feet kicked up a cloud of ochre dust, chasing after a rogue cricket ball. He lived on a sprawling cattle station, miles from anything, and his best friends were his kelpie, Bluey, and the endless outback. Today, however, he felt a pang of loneliness. He missed the city, the shops, the other kids. Suddenly, a flash of emerald green caught his eye. Nestled under a gum tree, a sleek, black cat with jade eyes stared back at him. A new friend, perhaps? Finn grinned. The outback, he realized, always had a surprise or two.\n"}],"role":"model"},"finish_reason":"STOP","avg_logprobs":-0.6154808869233003}],"model_version":"gemini-2.0-flash-lite","usage_metadata":{"candidates_token_count":148,"candidates_tokens_details":[{"modality":"TEXT","token_count":148}],"prompt_token_count":60,"prompt_tokens_details":[{"modality":"TEXT","token_count":60}],"total_token_count":208},"automatic_function_calling_history":[]}
-----------------------------------------------------------

INFO:__main__:[StoryFlowAgent] Workflow finished.

--- Agent Interaction Result ---
Agent Final Response:  StoryGenerator:
The Australian sun beat down on ten-year-old Finn as he raced across the red dirt. His bare feet kicked up a cloud of ochre dust, chasing after a rogue cricket ball. He lived on a sprawling cattle station, miles from anything, and his best friends were his kelpie, Bluey, and the endless outback. Today, however, he felt a pang of loneliness. He missed the city, the shops, the other kids. Suddenly, a flash of emerald green caught his eye. Nestled under a gum tree, a sleek, black cat with jade eyes stared back at him. A new friend, perhaps? Finn grinned. The outback, he realized, always had a surprise or two.

Final Session State:
{
  "topic": "a brave kitten exploring a haunted house",
  "current_story": "StoryGenerator:\nThe Australian sun beat down on ten-year-old Finn as he raced across the red dirt. His bare feet kicked up a cloud of ochre dust, chasing after a rogue cricket ball. He lived on a sprawling cattle station, miles from anything, and his best friends were his kelpie, Bluey, and the endless outback. Today, however, he felt a pang of loneliness. He missed the city, the shops, the other kids. Suddenly, a flash of emerald green caught his eye. Nestled under a gum tree, a sleek, black cat with jade eyes stared back at him. A new friend, perhaps? Finn grinned. The outback, he realized, always had a surprise or two.\n"
}
-------------------------------

INFO:__main__:updated topic in session state: a brave kitten exploring a haunted house

Expected behavior
The final state log should reflect the updated value of the story topic but is not doing so. But the updated topic is being used by the Agent to produce the story

Desktop

• OS: debian 12
• Python version(python -V): 3.13.2
• ADK version(pip show google-adk): 0.2.0

Additional context
Use the following env vars to run

GOOGLE_GENAI_USE_VERTEXAI=FALSE
GOOGLE_API_KEY=<gemini_api_key>

[wietsevenema]

I think you might want to use session_service.create_session to have your user_input_topic as initial state, would that work?

This is what I think happens in your example:

1. get_session returns a deep copy: When your call_dummy_agent calls session_service.get_session(...), it receives a copy of the session data, not a direct reference to the session stored within the service.
2. Modifying the copy: The line current_session.state.update({"topic": user_input_topic}) modifies this copy, not the original session data held by session_service.
3. runner.run gets a fresh copy: When runner.run(...) is called, it internally calls session_service.get_session(...) again. This fetches another copy of the original, unmodified session data from the service (because the changes made in step 2 were only to a temporary copy).

[shankha98]

Turns out even though the example in the docs show state updates directly, it is advised to not explicitly update the state but use the append_event() method of the session service instance. The link in the doc is here.

The code sample works like intended if I swap out the state update part with this bit

    state_changes = {"topic": user_input_topic, "updated": "true", "update_count": 1}

    # --- Create Event with Actions ---
    actions_with_update = EventActions(state_delta=state_changes)
    # This event might represent an internal system action, not just an agent response
    system_event = Event(
        invocation_id="inv_login_update",
        author="system",  # Or 'agent', 'tool' etc.
        actions=actions_with_update,
        timestamp=time.time(),
        # content might be None or represent the action taken
    )

    # --- Append the Event (This updates the state) ---
    session_service.append_event(session, system_event)
    print("`append_event` called with explicit state delta.")

    updated_session = session_service.get_session(
        app_name=APP_NAME, user_id=USER_ID, session_id=SESSION_ID
    )

    print("Updated Session State:")
    print(json.dumps(updated_session.state, indent=2))
    print("-------------------------------\n")

The logs from the updated code -

INFO:__main__:Initial session state: {'topic': 'a brave kitten exploring a haunted house'}
`append_event` called with explicit state delta.
Updated Session State:
{
  "topic": "a boy in australia",
  "updated": "true",
  "update_count": 1
}
-------------------------------

WARNING:google.adk.runners:Event from an unknown agent: system, event id: Uu1v3too
INFO:__main__:[StoryFlowAgent] Starting story generation workflow.
INFO:__main__:[StoryFlowAgent] Running StoryGenerator...
INFO:google.adk.models.google_llm:Sending out request, model: gemini-2.0-flash-lite, backend: ml_dev, stream: False
INFO:google.adk.models.google_llm:
LLM Request:
-----------------------------------------------------------
System Instruction:
You are a story writer. Write a short story (around 100 words) about a cat,
    based on the topic provided in session state with key 'topic'

You are an agent. Your internal name is "StoryGenerator".
-----------------------------------------------------------
Contents:
{"parts":[{"text":"Generate a story about: a boy in australia"}],"role":"user"}
-----------------------------------------------------------
Functions:

-----------------------------------------------------------

INFO:google_genai.models:AFC is enabled with max remote calls: 10.
INFO:httpx:HTTP Request: POST https://generativelanguage.googleapis.com/v1beta/models/gemini-2.0-flash-lite:generateContent "HTTP/1.1 200 OK"
INFO:google.adk.models.google_llm:
LLM Response:
-----------------------------------------------------------
Text:
The midday sun beat down on Leo's neck as he trudged along the red dirt track. He was on the hunt. Not for a kangaroo, though he'd seen plenty bounding across the outback. No, Leo was searching for the perfect piece of spinifex grass to tickle his cat, Matilda.

Back home in the corrugated iron shack, Matilda, a ginger tabby, would be napping in a sunbeam. Leo knew her routine: a brief stretch, a dramatic yawn, then back to slumber. He pictured her now, tail twitching, waiting for her daily dose of silliness. He found it - a perfect, prickly blade. "Here comes trouble, Matilda," he grinned, heading home.

-----------------------------------------------------------
Function calls:

-----------------------------------------------------------
Raw response:
{"candidates":[{"content":{"parts":[{"text":"The midday sun beat down on Leo's neck as he trudged along the red dirt track. He was on the hunt. Not for a kangaroo, though he'd seen plenty bounding across the outback. No, Leo was searching for the perfect piece of spinifex grass to tickle his cat, Matilda.\n\nBack home in the corrugated iron shack, Matilda, a ginger tabby, would be napping in a sunbeam. Leo knew her routine: a brief stretch, a dramatic yawn, then back to slumber. He pictured her now, tail twitching, waiting for her daily dose of silliness. He found it - a perfect, prickly blade. \"Here comes trouble, Matilda,\" he grinned, heading home.\n"}],"role":"model"},"finish_reason":"STOP","avg_logprobs":-0.7460601702839339}],"model_version":"gemini-2.0-flash-lite","usage_metadata":{"candidates_token_count":147,"candidates_tokens_details":[{"modality":"TEXT","token_count":147}],"prompt_token_count":60,"prompt_tokens_details":[{"modality":"TEXT","token_count":60}],"total_token_count":207},"automatic_function_calling_history":[]}
-----------------------------------------------------------

INFO:__main__:[StoryFlowAgent] Workflow finished.

--- Agent Interaction Result ---
Agent Final Response:  The midday sun beat down on Leo's neck as he trudged along the red dirt track. He was on the hunt. Not for a kangaroo, though he'd seen plenty bounding across the outback. No, Leo was searching for the perfect piece of spinifex grass to tickle his cat, Matilda.

Back home in the corrugated iron shack, Matilda, a ginger tabby, would be napping in a sunbeam. Leo knew her routine: a brief stretch, a dramatic yawn, then back to slumber. He pictured her now, tail twitching, waiting for her daily dose of silliness. He found it - a perfect, prickly blade. "Here comes trouble, Matilda," he grinned, heading home.

Final Session State:
{
  "topic": "a boy in australia",
  "updated": "true",
  "update_count": 1,
  "current_story": "The midday sun beat down on Leo's neck as he trudged along the red dirt track. He was on the hunt. Not for a kangaroo, though he'd seen plenty bounding across the outback. No, Leo was searching for the perfect piece of spinifex grass to tickle his cat, Matilda.\n\nBack home in the corrugated iron shack, Matilda, a ginger tabby, would be napping in a sunbeam. Leo knew her routine: a brief stretch, a dramatic yawn, then back to slumber. He pictured her now, tail twitching, waiting for her daily dose of silliness. He found it - a perfect, prickly blade. \"Here comes trouble, Matilda,\" he grinned, heading home.\n"
}
-------------------------------

INFO:__main__:updated topic in session state: a boy in australia

Would be nice if they update the code example in the docs here with the right approach.

[MahorShekhar]

Just adding here, I have also observed that when multiple tool calls are executed, only last tool call's session state update is persisted. This is weird.

[smtmobl]

I'm facing the same issue with state update in these methods:

async def _run_async_impl(self, ctx: InvocationContext) -> AsyncGenerator[Event, None]:
after_agent_callback
in both, whatever states added/updated within the method will be discarded and will not be persisted into DB (in my case, I using mySQL).
_run_async_impl is behave a bit different, when invoke a sub agent inside _run_async_impl methhod, the sub agent can still get the state data out, but finally, they will be discarded.

[hammertoe]

I'm also observing the same behaviour, it seems like the state is intermittent. #1119 like some kind of timing or race condition, as sometimes it works, sometimes it doesn't.

FYI, Claude may have found the cause:

Based on my research of the ADK source code and documentation, I found the issue! Let me explain what's happening:
The Root Cause:
Looking at the ADK architecture, I found that agents implement the _run_async_impl method which yields events and the Runner manages the overall Event Loop, receives events yielded by the Execution Logic.
However, there's a key distinction between LlmAgent and custom BaseAgent implementations:
LlmAgent Flow:
LlmAgent has its own internal flow that calls _run_async_impl but also handles LLM interactions, tool calls, and model callbacks. When an LlmAgent runs, it may not always call your custom _run_async_impl - it depends on the flow control and whether the LLM decides to use tools or transfer to sub-agents.
The Issue:
The WriterAgent is an LlmAgent, not a BaseAgent. When the LLM receives the enriched turtle data in session state, it may:

✅ Sometimes use our custom _run_async_impl method
❌ Sometimes bypass it and use the default LlmAgent flow
🎲 Inconsistent behavior depending on the LLM's internal decision-making

The Solution:
Instead of overriding _run_async_impl on an LlmAgent (which is unreliable), we should use ADK's callback system. Looking at the documentation, callbacks are called at specific, predefined points and allow you to observe, customize, and even control the agent's behavior.
We should use a before_agent_callback that can reliably modify the enriched message before the LLM processes it: