Skip to content

feat(model_armor): screen tool output in ModelArmorPlugin (or companion plugin) #6966

Description

@saurabh4269

Is your feature request related to a specific problem?

Same problem as #6964: the documented limitation that tool output is not screened by ModelArmorPlugin. For agents that load untrusted text through tools, screening only user prompts and model output leaves a documented security gap.

https://github.com/google/adk-python/blob/main/docs/guides/integrations/model_armor/index.md#limitations

Describe the Solution You'd Like

One of (maintainer choice):

  1. Extend ModelArmorConfig with e.g. screen_tool_output: bool = False and implement screening in ModelArmorPlugin.after_tool_callback, reusing the prompt template; or
  2. Ship ToolOutputModelArmorPlugin in google.adk.integrations.model_armor as an optional companion with shared config/client.

Either approach should mirror fail-closed semantics of the existing plugin and include unit tests in tests/unittests/integrations/model_armor/.

Impact on your work

Production agent pipelines where tool output is the dominant untrusted input channel. Reference sample implementation: #6964 / #6965.

Willingness to contribute

Yes. Sample PR #6965 is open; happy to contribute a core plugin PR after design alignment.

Describe Alternatives You've Considered

  • Document-only workaround: insufficient; limitation is already documented but not addressed.
  • Per-app custom plugins forever: duplicates logic across every ADK 2.8 adopter with tool-heavy threat models.

Proposed API / Implementation

@dataclass
class ModelArmorConfig:
    ...
    screen_tool_output: bool = False
    tool_output_blocked_message: str = "Tool output was blocked by Model Armor."
async def after_tool_callback(self, *, tool, tool_args, tool_context, result, **kwargs):
    if not self._config.screen_tool_output:
        return None
    ...

Additional context

Activity

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Metadata

Metadata

Assignees

Labels

models[Component] This issue is related to model supportneeds review[Status] The PR/issue is awaiting review from the maintainer

Projects

No projects

    Milestone

    No milestone

    Relationships

    None yet

    Development

    No branches or pull requests

    Issue actions