docs: Adds a sample agent to illustrate state usage via callbacks.

PiperOrigin-RevId: 764981675

contributing/samples/session_state_agent/README.md

@@ -0,0 +1,66 @@
+# Sample Agent to demo session state persistence.
+
+## Lifecycle of session state
+
+After assigning a state using the context object (e.g.
+`tool_context.state['log_query_var'] = 'log_query_var_value'`):
+
+* The state is available for use in a later callback.
+* Once the resulting event is processed by the runner and appneded in the
+  session, the state will be also persisted in the session.
+
+This sample agent is for demonstrating the aforementioned behavior.
+
+## Run the agent
+
+Run below command:
+
+```bash
+$ adk run contributing/samples/session_state_agent --replay contributing/samples/session_state_agent/input.json
+```
+
+And you should see below output:
+
+```bash
+[user]: hello world!
+===================== In before_agent_callback ==============================
+** Asserting keys are cached in context: ['before_agent_callback_state_key'] pass ✅
+** Asserting keys are already persisted in session: [] pass ✅
+** Asserting keys are not persisted in session yet: ['before_agent_callback_state_key'] pass ✅
+============================================================
+===================== In before_model_callback ==============================
+** Asserting keys are cached in context: ['before_agent_callback_state_key', 'before_model_callback_state_key'] pass ✅
+** Asserting keys are already persisted in session: ['before_agent_callback_state_key'] pass ✅
+** Asserting keys are not persisted in session yet: ['before_model_callback_state_key'] pass ✅
+============================================================
+===================== In after_model_callback ==============================
+** Asserting keys are cached in context: ['before_agent_callback_state_key', 'before_model_callback_state_key', 'after_model_callback_state_key'] pass ✅
+** Asserting keys are already persisted in session: ['before_agent_callback_state_key'] pass ✅
+** Asserting keys are not persisted in session yet: ['before_model_callback_state_key', 'after_model_callback_state_key'] pass ✅
+============================================================
+[root_agent]: Hello! How can I help you verify something today?
+
+===================== In after_agent_callback ==============================
+** Asserting keys are cached in context: ['before_agent_callback_state_key', 'before_model_callback_state_key', 'after_model_callback_state_key', 'after_agent_callback_state_key'] pass ✅
+** Asserting keys are already persisted in session: ['before_agent_callback_state_key', 'before_model_callback_state_key', 'after_model_callback_state_key'] pass ✅
+** Asserting keys are not persisted in session yet: ['after_agent_callback_state_key'] pass ✅
+============================================================
+```
+
+## Detailed Explanation
+
+As rule of thumb, to read and write session state, user should assume the
+state is available after writing via the context object
+(`tool_context`, `callback_context` or `readonly_context`).
+
+### Current Behavior
+
+The current behavior of pesisting states are:
+
+* for `before_agent_callback`: state delta will be persisted after all callbacks are processed.
+* for `before_model_callback`: state delta will be persisted with the final LlmResponse,
+  aka. after `after_model_callback` is processed.
+* for `after_model_callback`: state delta will be persisted together with the event of LlmResponse.
+* for `after_agent_callback`: state delta will be persisted after all callbacks are processed.
+
+**NOTE**: the current behavior is considered implementation detail and may be changed later. **DO NOT** rely on it.

contributing/samples/session_state_agent/__init__.py

@@ -0,0 +1 @@
+from . import agent

contributing/samples/session_state_agent/agent.py

@@ -0,0 +1,180 @@
+# Copyright 2025 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""The agent to demo the session state lifecycle.
+
+This agent illustrate how session state will be cached in context and persisted
+in session state.
+"""
+
+
+import logging
+from typing import Optional
+
+from google.adk.agents.callback_context import CallbackContext
+from google.adk.agents.llm_agent import Agent
+from google.adk.models.llm_request import LlmRequest
+from google.adk.models.llm_response import LlmResponse
+from google.genai import types
+
+logger = logging.getLogger('google_adk.' + __name__)
+
+
+async def assert_session_values(
+    ctx: CallbackContext,
+    title: str,
+    *,
+    keys_in_ctx_session: Optional[list[str]] = None,
+    keys_in_service_session: Optional[list[str]] = None,
+    keys_not_in_service_session: Optional[list[str]] = None,
+):
+  session_in_ctx = ctx._invocation_context.session
+  session_in_service = (
+      await ctx._invocation_context.session_service.get_session(
+          app_name=session_in_ctx.app_name,
+          user_id=session_in_ctx.user_id,
+          session_id=session_in_ctx.id,
+      )
+  )
+  assert session_in_service is not None
+
+  print(f'===================== {title} ==============================')
+  print(
+      f'** Asserting keys are cached in context: {keys_in_ctx_session}', end=' '
+  )
+  for key in keys_in_ctx_session or []:
+    assert key in session_in_ctx.state
+  print('\033[92mpass ✅\033[0m')
+
+  print(
+      '** Asserting keys are already persisted in session:'
+      f' {keys_in_service_session}',
+      end=' ',
+  )
+  for key in keys_in_service_session or []:
+    assert key in session_in_service.state
+  print('\033[92mpass ✅\033[0m')
+
+  print(
+      '** Asserting keys are not persisted in session yet:'
+      f' {keys_not_in_service_session}',
+      end=' ',
+  )
+  for key in keys_not_in_service_session or []:
+    assert key not in session_in_service.state
+  print('\033[92mpass ✅\033[0m')
+  print('============================================================')
+
+
+async def before_agent_callback(
+    callback_context: CallbackContext,
+) -> Optional[types.Content]:
+  if 'before_agent_callback_state_key' in callback_context.state:
+    return types.ModelContent('Sorry, I can only reply once.')
+
+  callback_context.state['before_agent_callback_state_key'] = (
+      'before_agent_callback_state_value'
+  )
+
+  await assert_session_values(
+      callback_context,
+      'In before_agent_callback',
+      keys_in_ctx_session=['before_agent_callback_state_key'],
+      keys_in_service_session=[],
+      keys_not_in_service_session=['before_agent_callback_state_key'],
+  )
+
+
+async def before_model_callback(
+    callback_context: CallbackContext, llm_request: LlmRequest
+):
+  callback_context.state['before_model_callback_state_key'] = (
+      'before_model_callback_state_value'
+  )
+
+  await assert_session_values(
+      callback_context,
+      'In before_model_callback',
+      keys_in_ctx_session=[
+          'before_agent_callback_state_key',
+          'before_model_callback_state_key',
+      ],
+      keys_in_service_session=['before_agent_callback_state_key'],
+      keys_not_in_service_session=['before_model_callback_state_key'],
+  )
+
+
+async def after_model_callback(
+    callback_context: CallbackContext, llm_response: LlmResponse
+):
+  callback_context.state['after_model_callback_state_key'] = (
+      'after_model_callback_state_value'
+  )
+
+  await assert_session_values(
+      callback_context,
+      'In after_model_callback',
+      keys_in_ctx_session=[
+          'before_agent_callback_state_key',
+          'before_model_callback_state_key',
+          'after_model_callback_state_key',
+      ],
+      keys_in_service_session=[
+          'before_agent_callback_state_key',
+      ],
+      keys_not_in_service_session=[
+          'before_model_callback_state_key',
+          'after_model_callback_state_key',
+      ],
+  )
+
+
+async def after_agent_callback(callback_context: CallbackContext):
+  callback_context.state['after_agent_callback_state_key'] = (
+      'after_agent_callback_state_value'
+  )
+
+  await assert_session_values(
+      callback_context,
+      'In after_agent_callback',
+      keys_in_ctx_session=[
+          'before_agent_callback_state_key',
+          'before_model_callback_state_key',
+          'after_model_callback_state_key',
+          'after_agent_callback_state_key',
+      ],
+      keys_in_service_session=[
+          'before_agent_callback_state_key',
+          'before_model_callback_state_key',
+          'after_model_callback_state_key',
+      ],
+      keys_not_in_service_session=[
+          'after_agent_callback_state_key',
+      ],
+  )
+
+
+root_agent = Agent(
+    name='root_agent',
+    description='a verification agent.',
+    instruction=(
+        'Log all users query with `log_query` tool. Must always remind user you'
+        ' cannot answer second query because your setup.'
+    ),
+    model='gemini-2.0-flash-001',
+    before_agent_callback=before_agent_callback,
+    before_model_callback=before_model_callback,
+    after_model_callback=after_model_callback,
+    after_agent_callback=after_agent_callback,
+)

contributing/samples/session_state_agent/input.json

@@ -0,0 +1,4 @@
+{
+  "state": {},
+  "queries": ["hello world!"]
+}