docs(callbacks): correct Python callback signatures and return types

docs/callbacks/design-patterns-and-best-practices.md

@@ -76,7 +76,7 @@ Alter data just before it's sent to the LLM/tool or just after it's received.
 - **`before_model_callback`:** Modify `llm_request` (e.g., add system instructions based on `state`)
 - **`after_model_callback`:** Modify the returned `LlmResponse` (e.g., format text, filter content)
 - **`before_tool_callback`:** Modify the tool `args` dictionary (or Map in Java)
-- **`after_tool_callback`:** Modify the `tool_response` dictionary (or Map in Java)
+- **`after_tool_callback`:** Modify the `tool_response` (the tool's raw return value in Python, or a Map in Java)
 
 **Example Use Case:**
 `before_model_callback` appends "User language preference: Spanish" to `llm_request.config.system_instruction` if `context.state['lang'] == 'es'`.
@@ -116,15 +116,15 @@ A `before_tool_callback` for a secure API checks for an auth token in state; if
 Save or load session-related files or large data blobs during the agent lifecycle.
 
 **Implementation:**
-- **Saving:** Use `callback_context.save_artifact` / `await tool_context.save_artifact` to store data:
+- **Saving:** Use `await callback_context.save_artifact` / `await tool_context.save_artifact` to store data:
   - Generated reports
   - Logs
   - Intermediate data
-- **Loading:** Use `load_artifact` to retrieve previously stored artifacts
+- **Loading:** Use `await callback_context.load_artifact` / `await tool_context.load_artifact` to retrieve previously stored artifacts
 - **Tracking:** Changes are tracked via `Event.actions.artifact_delta`
 
 **Example Use Case:**
-An `after_tool_callback` for a "generate_report" tool saves the output file using `await tool_context.save_artifact("report.pdf", report_part)`. A `before_agent_callback` might load a configuration artifact using `callback_context.load_artifact("agent_config.json")`.
+An `after_tool_callback` for a "generate_report" tool saves the output file using `await tool_context.save_artifact("report.pdf", report_part)`. A `before_agent_callback` might load a configuration artifact using `await callback_context.load_artifact("agent_config.json")`.
 
 ## Best Practices for Callbacks
 
@@ -134,7 +134,7 @@ An `after_tool_callback` for a "generate_report" tool saves the output file usin
 Design each callback for a single, well-defined purpose (e.g., just logging, just validation). Avoid monolithic callbacks.
 
 **Mind Performance:**
-Callbacks execute synchronously within the agent's processing loop. Avoid long-running or blocking operations (network calls, heavy computation). Offload if necessary, but be aware this adds complexity.
+Callbacks execute inline within the agent's processing loop, and the loop waits for each one to finish. In Python, declare a callback `async def` when it must `await` something itself, such as `save_artifact`; ADK awaits it. Either way, avoid long-running or blocking operations (network calls, heavy computation). Offload if necessary, but be aware this adds complexity.
 
 ### Error Handling
 
@@ -168,6 +168,6 @@ If a callback performs actions with external side effects (e.g., incrementing an
 - Add clear docstrings explaining their purpose, when they run, and any side effects (especially state modifications)
 
 **Use Correct Context Type:**
-Always use the specific context type provided (`CallbackContext` for agent/model, `ToolContext` for tools) to ensure access to the appropriate methods and properties.
+Use the context type your SDK documents for the hook you are implementing. In Python, `CallbackContext` and `ToolContext` are aliases of the same `Context` class and expose identical methods, so prefer `Context` for new code; the older names are kept for backward compatibility and you will still meet them in existing code. See [Context objects](../context/index.md) for the full set of context types.
 
 By applying these patterns and best practices, you can effectively use callbacks to create more robust, observable, and customized agent behaviors in ADK.

docs/callbacks/index.md

@@ -79,18 +79,18 @@ When the ADK framework encounters a point where a callback can run (e.g., just b
 
     * The specific return type can vary depending on the language. In Java, the equivalent return type is `Optional.empty()`. In Kotlin, it is `CallbackChoice.Continue(value)` (for `before_*` callbacks) or returning the original object (for `after_*` callbacks). Refer to the API documentation for language specific guidance.
     * This is the standard way to signal that your callback has finished its work (e.g., logging, inspection, minor modifications to input arguments) and that the ADK agent should **proceed with its normal operation**.
-    * For `before_*` callbacks (`before_agent`, `before_model`, `before_tool`), returning `CallbackChoice.Continue(...)` means the next step in the sequence (running the agent logic, calling the LLM, executing the tool) will occur.
-    * For `after_*` callbacks (`after_agent`, `after_model`, `after_tool`), returning the result just produced (the agent's output, the LLM's response, the tool's result) as is means the framework will continue processing.
+    * For `before_*` callbacks (`before_agent`, `before_model`, `before_tool`), returning `None` means the next step in the sequence (running the agent logic, calling the LLM, executing the tool) will occur.
+    * For `after_*` callbacks (`after_agent`, `after_model`, `after_tool`), returning `None` leaves the result just produced (the agent's output, the LLM's response, the tool's result) untouched. In Python, returning that result back instead of `None` is not equivalent for `after_agent_callback`: it emits a second event carrying the same content.
 
 2. **`return <Specific Object>` (Override Default Behavior):**
 
     * Returning a *specific type of object* (instead of signaling "Continue") is how you **override** the ADK agent's default behavior. In Kotlin, this is achieved by returning `CallbackChoice.Break(value)` (for `before_*` callbacks) or a replacement object (for `after_*` callbacks). The framework will use the object you return and *skip* the step that would normally follow or *replace* the result that was just generated.
-    * **`before_agent_callback` → `CallbackChoice.Break(Content)`**: Skips the agent's main execution logic. The returned `Content` object is immediately treated as the agent's final output for this turn. Useful for handling simple requests directly or enforcing access control.
-    * **`before_model_callback` → `CallbackChoice.Break(LlmResponse)`**: Skips the call to the external Large Language Model. The returned `LlmResponse` object is processed as if it were the actual response from the LLM. Ideal for implementing input guardrails, prompt validation, or serving cached responses.
-    * **`before_tool_callback` → `CallbackChoice.Break(Map<String, Any>)`**: Skips the execution of the actual tool function (or sub-agent). The returned `Map` is used as the result of the tool call, which is then typically passed back to the LLM. Perfect for validating tool arguments, applying policy restrictions, or returning mocked/cached tool results.
-    * **`after_agent_callback` → `Content`**: *Replaces* the `Content` that the agent's run logic just produced.
+    * **`before_agent_callback` → `Content`**: Skips the agent's main execution logic. The returned `Content` object is immediately treated as the agent's final output for this turn. Useful for handling simple requests directly or enforcing access control.
+    * **`before_model_callback` → `LlmResponse`**: Skips the call to the external Large Language Model. The returned `LlmResponse` object is processed as if it were the actual response from the LLM. Ideal for implementing input guardrails, prompt validation, or serving cached responses.
+    * **`before_tool_callback` → `dict`**: Skips the execution of the actual tool function (or sub-agent). The returned `dict` is used as the result of the tool call, which is then typically passed back to the LLM. Perfect for validating tool arguments, applying policy restrictions, or returning mocked/cached tool results.
+    * **`after_agent_callback` → `Content`**: *Appends* the returned `Content` as an additional event after the output the agent's run logic already produced. It does not replace that output; use `after_model_callback` to change a model response.
     * **`after_model_callback` → `LlmResponse`**: *Replaces* the `LlmResponse` received from the LLM. Useful for sanitizing outputs, adding standard disclaimers, or modifying the LLM's response structure.
-    * **`after_tool_callback` → `Map<String, Any>`**: *Replaces* the `Map` result returned by the tool. Allows for post-processing or standardization of tool outputs before they are sent back to the LLM.
+    * **`after_tool_callback` → `dict`**: *Replaces* the result returned by the tool. Allows for post-processing or standardization of tool outputs before they are sent back to the LLM.
 
 **Conceptual Code Example (Guardrail):**
 

docs/callbacks/types-of-callbacks.md

@@ -37,8 +37,38 @@ These callbacks are available on *any* agent that inherits from `BaseAgent` (inc
     | `after_agent_callback` | `callback_context` |
     | `before_model_callback` | `callback_context`, `llm_request` |
     | `after_model_callback` | `callback_context`, `llm_response` |
+    | `on_model_error_callback` | `callback_context`, `llm_request`, `error` |
     | `before_tool_callback` | `tool`, `args`, `tool_context` |
     | `after_tool_callback` | `tool`, `args`, `tool_context`, `tool_response` |
+    | `on_tool_error_callback` | `tool`, `args`, `tool_context`, `error` |
+
+    Only `before_agent_callback` and `after_agent_callback` are fields on every
+    `BaseAgent`. The six model and tool callbacks in this table are fields on
+    `LlmAgent` only.
+
+??? note "Python: `async` callbacks and lists of callbacks"
+
+    In Python, a callback may be a plain `def` or an `async def`. ADK awaits
+    the result either way.
+
+    Every callback field also accepts a list of functions instead of a single
+    function. ADK invokes them in the order listed and stops at the first one
+    that returns a result: that value becomes the callback result, and the
+    remaining callbacks are skipped. What counts as a result differs by family.
+    The six `before_`/`after_` agent, model and tool hooks stop only on a
+    *truthy* value, so a callback returning `None`, or another falsy value such
+    as an empty `dict`, lets the next one run. `on_model_error_callback` and
+    `on_tool_error_callback` stop on any value that is not `None`, so an empty
+    `dict` from an `on_tool_error_callback` ends the chain, suppresses the
+    exception, and becomes the tool result.
+
+    ```python
+    root_agent = LlmAgent(
+        name="my_agent",
+        model="gemini-flash-latest",
+        before_model_callback=[check_policy, log_request],
+    )
+    ```
 
 ### Before Agent Callback
 
@@ -89,7 +119,7 @@ These callbacks are available on *any* agent that inherits from `BaseAgent` (inc
 
 ### After Agent Callback
 
-**When:** Called *immediately after* the agent's `_run_async_impl` (or `_run_live_impl`) method successfully completes. It does *not* run if the agent was skipped due to `before_agent_callback` returning content or if `end_invocation` was set during the agent's run.
+**When:** Called *immediately after* the agent's `_run_async_impl` (or `_run_live_impl`) method successfully completes. It does *not* run if the agent was skipped due to `before_agent_callback` returning content. In Python, setting `end_invocation` during the agent's run also skips it, but only on the `run_async` path; `run_live` does not re-check `end_invocation` after `_run_live_impl` finishes, so the callback still runs there.
 
 **Purpose:** Useful for cleanup tasks, post-execution validation, logging the completion of an agent's activity, or modifying final state.
 
@@ -134,11 +164,11 @@ These callbacks are available on *any* agent that inherits from `BaseAgent` (inc
 *   **Expected Outcome:** You'll see two scenarios:
     1. In the session *without* the `add_concluding_note: True` state, the callback allows the agent's original output ("Processing complete!") to be used.
     2. In the session *with* that state flag, the callback intercepts the agent's original output and appends it with its own message ("Concluding note added...").
-* **Understanding Callbacks:** This highlights how `after_` callbacks allow **post-processing** or **modification**. You can inspect the result of a step (the agent's run) and decide whether to let it pass through, change it, or completely replace it based on your logic.
+* **Understanding Callbacks:** This highlights how `after_` callbacks allow **post-processing**. You can inspect the result of a step and decide whether to let it pass through or add to it. As the note above says, an `after_agent_callback` cannot replace the agent's output: the content it returns is emitted as an *additional* event after the agent's own events. `after_model_callback` and `after_tool_callback` do replace the value they are given.
 
 ## LLM Interaction Callbacks
 
-These callbacks are specific to `LlmAgent` and provide hooks around the interaction with the Large Language Model.
+These callbacks are specific to `LlmAgent` and provide hooks around the interaction with the Large Language Model. In Python, `LlmAgent` also accepts an `on_model_error_callback`, which runs when the model call raises an exception. If it returns an `LlmResponse`, the exception is suppressed and that response is used instead.
 
 ### Before Model Callback
 
@@ -217,7 +247,7 @@ If the callback returns `None` (or a `Maybe.empty()` object in Java), the LLM co
 
 ## Tool Execution Callbacks
 
-These callbacks are also specific to `LlmAgent` and trigger around the execution of tools (including `FunctionTool`, `AgentTool`, etc.) that the LLM might request.
+These callbacks are also specific to `LlmAgent` and trigger around the execution of tools (including `FunctionTool`, `AgentTool`, etc.) that the LLM might request. In Python, `LlmAgent` also accepts an `on_tool_error_callback`, which runs when the tool raises an exception. If it returns a `dict`, the exception is suppressed and that dict is used as the tool result.
 
 ### Before Tool Callback
 
@@ -230,6 +260,14 @@ These callbacks are also specific to `LlmAgent` and trigger around the execution
 1. If the callback returns `None` (or a `Maybe.empty()` object in Java), the tool's `run_async` method is executed with the (potentially modified) `args`.
 2. If a dictionary (or `Map` in Java) is returned, the tool's `run_async` method is **skipped**. The returned dictionary is used directly as the result of the tool call. This is useful for caching or overriding tool behavior.
 
+!!! note "Python: only `None` lets the tool run"
+
+    ADK compares the returned value against `None`, so an empty `dict` counts
+    as an override: the tool is skipped and `{}` becomes the tool result.
+    Return `None`, not `{}`, when you want the tool to execute. With a list of
+    callbacks this applies to the last value produced, because an empty `dict`
+    does not stop the chain and is discarded if a later callback returns
+    something else.
 
 ??? "Code"
     === "Python"
@@ -269,6 +307,18 @@ These callbacks are also specific to `LlmAgent` and trigger around the execution
 1. If the callback returns `None` (or a `Maybe.empty()` object in Java), the original `tool_response` is used.
 2. If a new dictionary is returned, it **replaces** the original `tool_response`. This allows modifying or filtering the result seen by the LLM.
 
+!!! note "Python: `tool_response` is the tool's raw return value"
+
+    ADK wraps a non-`dict` result into `{"result": <value>}` only *after* the
+    callback has run, so a tool annotated `-> str` hands your
+    `after_tool_callback` a `str`, not a `dict`. Check the type before calling
+    dictionary methods on it.
+
+!!! note "Python: an empty `dict` still replaces the response"
+
+    ADK compares the returned value against `None`, so returning `{}` replaces
+    the tool response with `{}`. Return `None` to keep the original.
+
 ??? "Code"
     === "Python"
 

docs/safety/index.md

@@ -78,7 +78,7 @@ For example, a query tool can be designed to expect a policy to be read from the
 
     ```py
     # Conceptual example: Setting policy data intended for tool context
-    # In a real ADK app, this might be set in InvocationContext.session.state
+    # In a real ADK app, this might be written to session state from a callback,
     # or passed during tool initialization, then retrieved via ToolContext.
 
     policy = {} # Assuming policy is a dictionary
@@ -87,8 +87,10 @@ For example, a query tool can be designed to expect a policy to be read from the
 
     # Conceptual: Storing policy where the tool can access it via ToolContext later.
     # This specific line might look different in practice.
-    # For example, storing in session state:
-    invocation_context.session.state["query_tool_policy"] = policy
+    # For example, storing in session state from a callback. Write through the
+    # context's `state`, not `session.state` directly: only the context records
+    # the change in the event's state_delta so the SessionService persists it.
+    callback_context.state["query_tool_policy"] = policy
 
     # Or maybe passing during tool init:
     query_tool = QueryTool(policy=policy)
@@ -168,10 +170,10 @@ During the tool execution, [**`Tool Context`**](../tools-custom/index.md#tool-co
     ```py
     def query(query: str, tool_context: ToolContext) -> str | dict:
       # Assume 'policy' is retrieved from context, e.g., via session state:
-      # policy = tool_context.invocation_context.session.state.get('query_tool_policy', {})
+      # policy = tool_context.state.get('query_tool_policy', {})
 
       # --- Placeholder Policy Enforcement ---
-      policy = tool_context.invocation_context.session.state.get('query_tool_policy', {}) # Example retrieval
+      policy = tool_context.state.get('query_tool_policy', {}) # Example retrieval
       actual_tables = explainQuery(query) # Hypothetical function call
 
       if not set(actual_tables).issubset(set(policy.get('tables', []))):
@@ -383,7 +385,6 @@ When modifications to the tools to add guardrails aren't possible, the [**`Befor
     ```py
     # Hypothetical callback function
     def validate_tool_params(
-        callback_context: CallbackContext, # Correct context type
         tool: BaseTool,
         args: Dict[str, Any],
         tool_context: ToolContext
@@ -392,7 +393,7 @@ When modifications to the tools to add guardrails aren't possible, the [**`Befor
       print(f"Callback triggered for tool: {tool.name}, args: {args}")
 
       # Example validation: Check if a required user ID from state matches an arg
-      expected_user_id = callback_context.state.get("session_user_id")
+      expected_user_id = tool_context.state.get("session_user_id")
       actual_user_id_in_args = args.get("user_id_param") # Assuming tool takes 'user_id_param'
 
       if actual_user_id_in_args != expected_user_id: