fix: correct variable names, CallbackContext import, and add async session handling

docs/sessions/state.md

@@ -115,7 +115,7 @@ This is the simplest method for saving an agent's final text response directly i
           print(f"Agent responded.") # Response text is also in event.content
 
     # --- Check Updated State ---
-    updated_session = await session_service.get_session(app_name=APP_NAME, user_id=USER_ID, session_id=session_id)
+    updated_session = await session_service.get_session(app_name=app_name, user_id=user_id, session_id=session_id)
     print(f"State after agent run: {updated_session.state}")
     # Expected output might include: {'last_greeting': 'Hello there! How can I help you today?'}
     ```
@@ -207,7 +207,7 @@ For more comprehensive details on context objects, refer to the [Context documen
 
     ```python
     # In an agent callback or tool function
-    from google.adk.agents import CallbackContext # or ToolContext
+    from google.adk.agents.callback_context import CallbackContext # or ToolContext
 
     def my_callback_or_tool_function(context: CallbackContext, # Or ToolContext
                                      # ... other parameters ...

examples/python/snippets/agents/llm-agent/capital_agent.py

@@ -1,3 +1,17 @@
+# 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.
+
 # --- Full example code demonstrating LlmAgent with Tools vs. Output Schema ---
 import json # Needed for pretty printing dicts
 
@@ -79,10 +93,6 @@ def get_capital_city(country: str) -> str:
 # --- 5. Set up Session Management and Runners ---
 session_service = InMemorySessionService()
 
-# Create separate sessions for clarity, though not strictly necessary if context is managed
-session_service.create_session(app_name=APP_NAME, user_id=USER_ID, session_id=SESSION_ID_TOOL_AGENT)
-session_service.create_session(app_name=APP_NAME, user_id=USER_ID, session_id=SESSION_ID_SCHEMA_AGENT)
-
 # Create a runner for EACH agent
 capital_runner = Runner(
     agent=capital_agent_with_tool,
@@ -116,7 +126,7 @@ async def call_agent_and_print(
 
     print(f"<<< Agent '{agent_instance.name}' Response: {final_response_content}")
 
-    current_session = session_service.get_session(app_name=APP_NAME,
+    await current_session = session_service.get_session(app_name=APP_NAME,
                                                   user_id=USER_ID,
                                                   session_id=session_id)
     stored_output = current_session.state.get(agent_instance.output_key)
@@ -135,6 +145,10 @@ async def call_agent_and_print(
 
 # --- 7. Run Interactions ---
 async def main():
+    # Create separate sessions for clarity, though not strictly necessary if context is managed
+    await session_service.create_session(app_name=APP_NAME, user_id=USER_ID, session_id=SESSION_ID_TOOL_AGENT)
+    await session_service.create_session(app_name=APP_NAME, user_id=USER_ID, session_id=SESSION_ID_SCHEMA_AGENT)
+
     print("--- Testing Agent with Tool ---")
     await call_agent_and_print(capital_runner, capital_agent_with_tool, SESSION_ID_TOOL_AGENT, '{"country": "France"}')
     await call_agent_and_print(capital_runner, capital_agent_with_tool, SESSION_ID_TOOL_AGENT, '{"country": "Canada"}')
@@ -143,5 +157,16 @@ async def main():
     await call_agent_and_print(structured_runner, structured_info_agent_schema, SESSION_ID_SCHEMA_AGENT, '{"country": "France"}')
     await call_agent_and_print(structured_runner, structured_info_agent_schema, SESSION_ID_SCHEMA_AGENT, '{"country": "Japan"}')
 
+# Run directly if using Jupyter notebooks or IPython shell that support top-level await
 if __name__ == "__main__":
     await main()
+
+# --- OR ---
+
+# Uncomment the following lines if running as a standard Python script (.py file):
+# import asyncio
+# if __name__ == "__main__":
+#     try:
+#         asyncio.run(main())
+#     except Exception as e:
+#         print(f"An error occurred while running the agent demo: {e}")