feat: add after_tool_callback to screen tool output

Input and model output results that report `MATCH_FOUND` are now blocked even when `invocation_result` is not `SUCCESS` and `block_on_screening_failure` is `False`.

Merge #6969

Fixes #6966

PiperOrigin-RevId: 994696801

Author: princemittalr

docs/guides/integrations/model_armor/index.md

@@ -1,18 +1,18 @@
 # Model Armor
 
-`ModelArmorPlugin` screens user input and model output against [Google Cloud Model Armor](https://cloud.google.com/security-command-center/docs/model-armor-overview) templates. When a filter matches, or when screening cannot complete, the offending content is replaced with a safe message before it reaches the model or the user.
+`ModelArmorPlugin` screens user input, model output, and tool output against [Google Cloud Model Armor](https://cloud.google.com/security-command-center/docs/model-armor-overview) templates. When a filter matches, or when screening cannot complete, the offending content is replaced with a safe message before it reaches the model or the user.
 
 ## Introduction
 
 Model Armor is a Google Cloud service that inspects text for prompt injection and jailbreak attempts, harmful content, and sensitive data. You define what to look for in a *template* — a named, server-side policy — and the service returns a verdict for each piece of text you send it.
 
-The integration is two public types: `ModelArmorPlugin`, a `BasePlugin` subclass driven by `PluginManager`, and `ModelArmorConfig`, which says which templates to screen against and what to do about a match. The plugin reads text off the `LlmRequest` and `LlmResponse`, calls Model Armor, and returns a replacement `LlmResponse` when content should be blocked.
+The integration is two public types: `ModelArmorPlugin`, a `BasePlugin` subclass driven by `PluginManager`, and `ModelArmorConfig`, which says which templates to screen against and what to do about a match. The plugin reads text off the `LlmRequest`, `LlmResponse`, and tool results, calls Model Armor, and returns a replacement when content should be blocked.
 
 Key features:
 
-- **Input and output screening**, each governed by its own template, and each optional.
+- **Input, output, and tool output screening**, each governed by its own template, and each optional.
 - **Block screening failures by default**: by default screening failures are blocked rather than delivered.
-- **Blocked responses are marked** with `custom_metadata['model_armor_blocked']` so your application can detect them.
+- **Blocked input and model output responses are marked** with `custom_metadata['model_armor_blocked']` so your application can detect them (blocked tool outputs return a plain `{'error': ...}` dict to the model and do not carry this marker).
 
 ## Get started
 
@@ -76,22 +76,34 @@ Credentials come from Application Default Credentials.
 3. The text is sent to Model Armor's `SanitizeModelResponse` method.
 4. It acts on the result (below).
 
+### `after_tool_callback` - screening tool output
+
+`after_tool_callback` runs after each tool call:
+
+1. If `tool_output_template_name` is unset, it returns immediately and nothing is screened.
+2. It extracts and serializes text from the tool result (skipping raw bytes and non-text media parts), splitting oversized text into overlapping 65,536-character chunks.
+3. Each chunk is sent to Model Armor's `SanitizeUserPrompt` method.
+4. It acts on each result (below).
+
 ### Acting on a result
 
-| `invocation_result` | Meaning | Plugin behavior |
-| :--- | :--- | :--- |
-| `SUCCESS` | Every filter ran. | Check `filter_match_state`. |
-| Anything else | Some or all filters were skipped, failed, or the field was unset. | Screening failure. |
+| `filter_match_state` | `invocation_result` | Meaning | Plugin behavior |
+| :--- | :--- | :--- | :--- |
+| `MATCH_FOUND` | Any | At least one filter tripped. | Block content. |
+| Anything else | `SUCCESS` | Every filter ran with no match. | Pass through untouched. |
+| Anything else | Anything else | Some or all filters were skipped, failed, or the field was unset. | Screening failure. |
 
-When screening completes successfully, a `filter_match_state` of `MATCH_FOUND` means at least one filter tripped, and the content is blocked. Anything else passes through untouched.
+`filter_match_state` is checked first, so a partial result that still reports `MATCH_FOUND` is always blocked.
 
 A screening failure is routed through `block_on_screening_failure` and blocked by default.
 
 ### The blocked response
 
-Blocking returns an `LlmResponse` carrying the message for the direction that
-was screened: `input_blocked_message` for user input, `output_blocked_message`
-for model output.
+Blocking user input or model output returns an `LlmResponse` carrying the
+message for the direction that was screened: `input_blocked_message` for user
+input, `output_blocked_message` for model output. Blocking tool output returns
+`{'error': tool_output_blocked_message}` to the model as the tool result instead
+of an `LlmResponse`.
 
 ### Template paths and regional endpoints
 
@@ -124,27 +136,30 @@ Options introduced by `ModelArmorPlugin` (those inherited from `BasePlugin` are
 | :--- | :--- | :--- | :--- |
 | `prompt_template_name` | `str \| None` | `None` | Template used to screen user input. Unset means input is not screened. |
 | `response_template_name` | `str \| None` | `None` | Template used to screen model output. Unset means output is not screened. |
+| `tool_output_template_name` | `str \| None` | `None` | Template used to screen tool output. Unset means tool output is not screened. |
 | `input_blocked_message` | `str` | `"I'm sorry, but I can't help with that request."` | Replacement text shown when user input is blocked. |
 | `output_blocked_message` | `str` | `"I'm sorry, but I can't help with that request."` | Replacement text shown when model output is blocked. |
+| `tool_output_blocked_message` | `str` | `"Tool output was blocked by Model Armor."` | Replacement error message returned to the model when tool output is blocked. |
 | `block_on_screening_failure` | `bool` | `True` | Whether to block content that could not be screened. |
 
-At least one of the two template names must be set.
+At least one of `prompt_template_name`, `response_template_name`, or `tool_output_template_name` must be set.
 
-#### `prompt_template_name` and `response_template_name`
+#### `prompt_template_name`, `response_template_name`, and `tool_output_template_name`
 
-Both fields require fully-qualified resource paths formatted as:
+All template fields require fully-qualified resource paths formatted as:
 `projects/{project}/locations/{location}/templates/{template}`
 
-You can configure either or both:
+You can configure any combination:
 
 - `prompt_template_name`: Screens user input prompts before forwarding to the model.
 - `response_template_name`: Screens model responses before delivering to the user.
+- `tool_output_template_name`: Screens tool outputs before returning to the model.
 
-If both are set they must reside in the same GCP location — see [Template paths and regional endpoints](#template-paths-and-regional-endpoints).
+If multiple templates are set they must reside in the same GCP location — see [Template paths and regional endpoints](#template-paths-and-regional-endpoints).
 
-#### `input_blocked_message` and `output_blocked_message`
+#### `input_blocked_message`, `output_blocked_message`, and `tool_output_blocked_message`
 
-Defines the replacement text returned to the user when a prompt or response is blocked. Screening failures reuse the message for the direction that failed.
+Defines the replacement text returned to the user when a prompt or response is blocked, or returned to the model in `{'error': tool_output_blocked_message}` when tool output is blocked. Screening failures reuse the message for the direction that failed.
 
 #### `block_on_screening_failure`
 
@@ -174,7 +189,7 @@ config = ModelArmorConfig(
 
 ### Detecting blocks in your application
 
-Blocked responses carry a marker, so a UI can render them differently from a real answer:
+Blocked input and model output responses carry a `custom_metadata['model_armor_blocked']` marker, so a UI can render them differently from a real answer (blocked tool outputs return a plain `{'error': ...}` dict to the model and do not carry this marker):
 
 ```python
 async for event in runner.run_async(...):
@@ -184,7 +199,11 @@ async for event in runner.run_async(...):
 
 ## Limitations
 
-- **Tool output is not screened.** Only the most recent `user` content with text parts is sent for screening. Tool results are added to the request as `user` content whose only part is a `function_response` and doesn't reach Model Armor.
+- **Live streaming and long-running tools are not screened.** Live streaming tools send yielded chunks directly to the live request queue and long-running tools deliver their final `function_response` on a later turn, so their outputs skip `after_tool_callback`.
+
+- **Media in tool results is not screened.** Raw `bytes` values and non-text `Part` objects (such as `inline_data` or `file_data` media) in tool results are skipped during text extraction and are not sent to Model Armor.
+
+- **Plugin registration order matters.** `PluginManager` stops at the first plugin that returns a non-`None` value from a callback. Register `ModelArmorPlugin` ahead of plugins that return a value from `after_tool_callback` (such as `MultimodalToolResultsPlugin`) so tool output screening is not bypassed.
 
 - **Enforcement mode is limited.** The Model Armor plugin is currently limited to logging detection results and blocking content. Future extensions could include replacing or redacting text.
 

src/google/adk/integrations/model_armor/_config.py

@@ -21,6 +21,7 @@
 from pydantic import model_validator
 
 _DEFAULT_BLOCKED_MESSAGE = "I'm sorry, but I can't help with that request."
+_DEFAULT_TOOL_OUTPUT_BLOCKED_MESSAGE = 'Tool output was blocked by Model Armor.'
 
 
 class ModelArmorConfig(BaseModel):
@@ -50,15 +51,30 @@ class ModelArmorConfig(BaseModel):
   output_blocked_message: str = _DEFAULT_BLOCKED_MESSAGE
   """The safe replacement text returned to the user when model output is blocked."""
 
+  tool_output_template_name: Optional[str] = None
+  """The Model Armor template used to screen tool output.
+
+  Should use the fully-qualified resource name:
+  ``projects/{project}/locations/{location}/templates/{template}``.
+  If unset, tool output screening is skipped.
+  """
+
+  tool_output_blocked_message: str = _DEFAULT_TOOL_OUTPUT_BLOCKED_MESSAGE
+  """The safe replacement text returned when tool output is blocked."""
+
   block_on_screening_failure: bool = True
   """Whether to block when Model Armor screening fails."""
 
   @model_validator(mode='after')
   def _validate_templates(self) -> ModelArmorConfig:
     """Ensure at least one template is configured."""
-    if not self.prompt_template_name and not self.response_template_name:
+    if (
+        not self.prompt_template_name
+        and not self.response_template_name
+        and not self.tool_output_template_name
+    ):
       raise ValueError(
-          'At least one of prompt_template_name or response_template_name'
-          ' must be set for ModelArmorConfig.'
+          'At least one of prompt_template_name, response_template_name, or'
+          ' tool_output_template_name must be set for ModelArmorConfig.'
       )
     return self

src/google/adk/integrations/model_armor/_plugin.py

@@ -14,18 +14,22 @@
 
 """Model Armor guardrail plugin.
 
-Screens user input and model output with Google Cloud Model Armor in both
-unary (``run_async``) and live (``run_live``) modes, through the ordinary
-``before_model_callback`` and ``after_model_callback`` seams.
+Screens user input, model output, and tool output with Google Cloud Model Armor
+in both unary (``run_async``) and live (``run_live``) modes, through the
+ordinary ``before_model_callback``, ``after_model_callback``, and
+``after_tool_callback`` seams.
 
 - input reaches ``before_model_callback`` as request content.
 - output reaches ``after_model_callback`` as content parts in unary mode, and
   as an output transcription in live mode.
+- tool output reaches ``after_tool_callback`` as the tool result.
 """
 
 from __future__ import annotations
 
+import json
 import logging
+from typing import Any
 from typing import Optional
 
 from google.api_core.client_options import ClientOptions
@@ -38,6 +42,8 @@
 from ...models.llm_request import LlmRequest
 from ...models.llm_response import LlmResponse
 from ...plugins.base_plugin import BasePlugin
+from ...tools.base_tool import BaseTool
+from ...tools.tool_context import ToolContext
 from ._config import ModelArmorConfig
 
 try:
@@ -57,9 +63,13 @@
 
 USER_AGENT = f'adk-model-armor-plugin google-adk/{version.__version__}'
 
+# Model Armor skips filters past 65,536 tokens (~1 token/char worst case).
+_MAX_TOOL_OUTPUT_CHARS = 65_536
+_TOOL_OUTPUT_CHUNK_OVERLAP_CHARS = 4_096
+
 
 class ModelArmorPlugin(BasePlugin):
-  """A plugin that screens input and output with Google Cloud Model Armor."""
+  """A plugin that screens input, output, and tool output with Model Armor."""
 
   def __init__(
       self,
@@ -87,7 +97,9 @@ def __init__(
     self._credentials = credentials
 
     self._location = _shared_template_location(
-        config.prompt_template_name, config.response_template_name
+        config.prompt_template_name,
+        config.response_template_name,
+        config.tool_output_template_name,
     )
 
   async def before_model_callback(
@@ -186,6 +198,42 @@ async def _sanitize_model_response(
     )
     return response.sanitization_result
 
+  async def after_tool_callback(
+      self,
+      *,
+      tool: BaseTool,
+      tool_args: dict[str, Any],
+      tool_context: ToolContext,
+      result: dict[str, Any],
+  ) -> Optional[dict[str, Any]]:
+    """Screens tool output text against the configured tool output template."""
+    if not self._config.tool_output_template_name:
+      return None
+
+    text = _extract_tool_result_text(result)
+    if not text:
+      return None
+
+    step = _MAX_TOOL_OUTPUT_CHARS - _TOOL_OUTPUT_CHUNK_OVERLAP_CHARS
+    for start in range(0, len(text), step):
+      chunk = text[start : start + _MAX_TOOL_OUTPUT_CHARS]
+      try:
+        sanitization_result = await self._sanitize_user_prompt(
+            chunk, self._config.tool_output_template_name
+        )
+      except Exception:  # pylint: disable=broad-except
+        logger.exception('Model Armor tool output screening call failed.')
+        if self._config.block_on_screening_failure:
+          return {'error': self._config.tool_output_blocked_message}
+      else:
+        if self._should_block(sanitization_result, direction='tool output'):
+          return {'error': self._config.tool_output_blocked_message}
+
+      if start + _MAX_TOOL_OUTPUT_CHARS >= len(text):
+        break
+
+    return None
+
   async def close(self) -> None:
     """Closes the underlying client."""
     if self._client:
@@ -199,19 +247,30 @@ def _handle_sanitization_result(
       blocked_message: str,
   ) -> Optional[LlmResponse]:
     """Handles a Model Armor sanitization result."""
+    if self._should_block(result, direction=direction):
+      return self._blocked_response(blocked_message)
+    return None
+
+  def _should_block(
+      self,
+      result: modelarmor_v1.SanitizationResult,
+      *,
+      direction: str,
+  ) -> bool:
+    """Returns whether a sanitization result should be blocked."""
+    if result.filter_match_state == modelarmor_v1.FilterMatchState.MATCH_FOUND:
+      logger.warning('Model Armor %s sanitization match found.', direction)
+      return True
+
     if result.invocation_result != modelarmor_v1.InvocationResult.SUCCESS:
       logger.error(
           'Model Armor %s sanitization did not succeed: invocation_result=%r',
           direction,
           result.invocation_result,
       )
-      return self._handle_screening_failure(blocked_message)
-
-    if result.filter_match_state == modelarmor_v1.FilterMatchState.MATCH_FOUND:
-      logger.warning('Model Armor %s sanitization match found.', direction)
-      return self._blocked_response(blocked_message)
+      return self._config.block_on_screening_failure
 
-    return None
+    return False
 
   def _handle_screening_failure(
       self, blocked_message: str
@@ -272,6 +331,54 @@ def _content_text(content: Optional[types.Content]) -> Optional[str]:
   return '\n'.join(texts)
 
 
+_SKIP = object()
+
+
+def _filter_screenable_value(value: object) -> object:
+  """Filters bytes and non-text media parts out of a tool result value."""
+  if isinstance(
+      value, (bytes, bytearray, memoryview, types.Blob, types.FileData)
+  ):
+    return _SKIP
+  if isinstance(value, types.Part):
+    if value.text and not value.thought and value.text.strip():
+      return value.text
+    return _SKIP
+  if isinstance(value, dict):
+    cleaned = {
+        k: filtered
+        for k, v in value.items()
+        if not isinstance(k, (bytes, bytearray, memoryview))
+        and (filtered := _filter_screenable_value(v)) is not _SKIP
+    }
+    return cleaned if cleaned or not value else _SKIP
+  if isinstance(value, (list, tuple)):
+    cleaned_list = [
+        filtered
+        for item in value
+        if (filtered := _filter_screenable_value(item)) is not _SKIP
+    ]
+    return cleaned_list if cleaned_list or not value else _SKIP
+  return value
+
+
+def _extract_tool_result_text(result: Any) -> Optional[str]:
+  """Extracts screenable text from a tool result."""
+  if result is None:
+    return None
+  filtered = _filter_screenable_value(result)
+  if filtered is _SKIP:
+    return None
+  if isinstance(filtered, str):
+    return filtered if filtered.strip() else None
+  try:
+    serialized = json.dumps(filtered, default=str, ensure_ascii=False)
+    return serialized if serialized not in ('{}', 'null', '[]', '""') else None
+  except Exception:  # pylint: disable=broad-except
+    text = str(filtered)
+    return text if text.strip() else None
+
+
 def _regional_endpoint(location: str) -> str:
   """Builds the Model Armor regional endpoint for a location."""
   return f'modelarmor.{location}.rep.googleapis.com'