feat: add independently configured AgentSentry V3 moderation

Author: yuweihong.666

tests/test_agentsentry_public_compatibility.py

@@ -0,0 +1,106 @@
+# Copyright (c) 2025 Beijing Volcano Engine Technology Co., Ltd. and/or its affiliates.
+#
+# 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.
+
+import builtins
+
+import json
+
+from types import SimpleNamespace
+
+from unittest.mock import Mock
+
+import pytest
+
+from google.genai import types
+
+from google.adk.models import LlmResponse
+
+
+@pytest.fixture
+def public_env(monkeypatch):
+    for key in (
+        "AGENTSENTRY_URL",
+        "ENABLE_LLM_SHIELD",
+        "TOOL_LLM_SHIELD_REGION",
+        "TOOL_LLM_SHIELD_ACCOUNT_ID",
+        "AGENTSENTRY_AK",
+        "AGENTSENTRY_SK",
+        "AGENTSENTRY_ACCOUNT_ID",
+        "AGENTSENTRY_INSTANCE_ID",
+        "AGENTSENTRY_REGION",
+    ):
+        monkeypatch.delenv(key, raising=False)
+    monkeypatch.setenv("TOOL_LLM_SHIELD_APP_ID", "public-app")
+    monkeypatch.setenv("TOOL_LLM_SHIELD_API_KEY", "public-api-key")
+    monkeypatch.setenv("TOOL_LLM_SHIELD_URL", "https://public-custom.example")
+    # Public construction/calls must not even import the private adapter.
+    original = builtins.__import__
+
+    def guarded(name, *args, **kwargs):
+        if name.endswith("agentsentry_v3"):
+            raise AssertionError("Private adapter imported by public mode")
+        return original(name, *args, **kwargs)
+
+    monkeypatch.setattr(builtins, "__import__", guarded)
+
+
+def test_veadk_public_request_and_callbacks_unchanged(public_env, monkeypatch):
+    import veadk.tools.builtin_tools.llm_shield as module
+
+    p = module.LLMShieldPlugin()
+    assert p._sentry_v3 is None
+    post = Mock(
+        return_value=SimpleNamespace(
+            status_code=200, json=lambda: {"Result": {"Decision": {"DecisionType": 1}}}
+        )
+    )
+    monkeypatch.setattr(module.requests, "post", post)
+    assert p._request_llm_shield("text", "user") is None
+    args = post.call_args
+    assert args.args[0] == "https://public-custom.example/v2/moderate"
+    assert args.kwargs["headers"]["x-api-key"] == "public-api-key"
+    assert args.kwargs["params"] == {"Action": "Moderate", "Version": "2025-08-31"}
+    assert json.loads(args.kwargs["data"]) == {
+        "Message": {"Role": "user", "Content": "text", "ContentType": 1},
+        "Scene": "public-app",
+    }
+    p._request_llm_shield = Mock(return_value=None)
+    ctx = SimpleNamespace()
+    p.before_model_callback(
+        ctx,
+        SimpleNamespace(
+            contents=[
+                types.Content(
+                    role="user",
+                    parts=[types.Part(text="first"), types.Part(text="second")],
+                )
+            ]
+        ),
+    )
+    assert p._request_llm_shield.call_args.kwargs == {
+        "message": "first",
+        "role": "user",
+        "hook_name": "before_model",
+    }
+    p.after_model_callback(
+        ctx,
+        LlmResponse(
+            content=types.Content(role="model", parts=[types.Part(text="reply")])
+        ),
+    )
+    assert p._request_llm_shield.call_args.kwargs["hook_name"] == "after_model"
+    p.before_tool_callback(SimpleNamespace(name="tool"), {"x": "arg"}, ctx)
+    assert p._request_llm_shield.call_args.kwargs["role"] == "user"
+    p.after_tool_callback(None, {}, ctx, "result")
+    assert p._request_llm_shield.call_args.kwargs["hook_name"] == "after_tool_call"

tests/test_agentsentry_review_regressions.py

@@ -0,0 +1,86 @@
+# Copyright (c) 2025 Beijing Volcano Engine Technology Co., Ltd. and/or its affiliates.
+#
+# 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.
+
+from types import SimpleNamespace as NS
+from unittest.mock import Mock
+
+from google.genai import types
+from google.adk.models import LlmResponse
+
+
+def shield(monkeypatch):
+    for key, value in {
+        "AGENTSENTRY_URL": "http://sentry/OpenTOP/V3/Lumen/Moderate",
+        "AGENTSENTRY_ACCOUNT_ID": "account",
+        "AGENTSENTRY_INSTANCE_ID": "instance",
+        "AGENTSENTRY_AK": "ak",
+        "AGENTSENTRY_SK": "sk",
+    }.items():
+        monkeypatch.setenv(key, value)
+    from veadk.tools.builtin_tools.llm_shield import LLMShieldPlugin
+
+    plugin = LLMShieldPlugin()
+    plugin._sentry_v3.moderate = Mock(return_value=None)
+    return plugin
+
+
+def test_hidden_thought_is_never_a_plan(monkeypatch):
+    plugin = shield(monkeypatch)
+    ctx = NS(agent_states={})
+    plugin.after_model_callback(
+        ctx,
+        LlmResponse(
+            content=types.Content(
+                role="model",
+                parts=[
+                    types.Part(text="SECRET", thought=True),
+                    types.Part(text="visible"),
+                    types.Part(
+                        function_call=types.FunctionCall(name="record", args={})
+                    ),
+                ],
+            )
+        ),
+    )
+    plugin.before_tool_callback(NS(name="record"), {}, ctx)
+    assert plugin._sentry_v3.moderate.call_args.args[0] == "visible"
+
+
+def test_dispatch_once_per_invocation_and_no_old_text_for_image(monkeypatch):
+    plugin = shield(monkeypatch)
+    user = types.Content(role="user", parts=[types.Part(text="hello")])
+    ctx = NS(agent_states={}, user_content=user)
+    req = NS(contents=[user])
+    plugin.before_model_callback(ctx, req)
+    plugin.before_model_callback(ctx, req)
+    assert plugin._sentry_v3.moderate.call_count == 1
+    image = types.Content(
+        role="user",
+        parts=[
+            types.Part(inline_data=types.Blob(data=b"image", mime_type="image/png"))
+        ],
+    )
+    ctx = NS(agent_states={}, user_content=image)
+    response = plugin.before_model_callback(ctx, NS(contents=[user, image]))
+    assert response is not None
+    assert "unsupported" in response.content.parts[0].text.lower()
+    assert plugin._sentry_v3.moderate.call_count == 1
+
+
+def test_private_disable_requires_no_credentials(monkeypatch):
+    plugin = shield(monkeypatch)
+    monkeypatch.setenv("AGENTSENTRY_ENABLED", "false")
+    monkeypatch.delenv("AGENTSENTRY_SK")
+    disabled = type(plugin)()
+    assert disabled.before_model_callback(NS(), NS()) is None

tests/test_agentsentry_v3_callbacks.py

@@ -0,0 +1,101 @@
+# Copyright (c) 2025 Beijing Volcano Engine Technology Co., Ltd. and/or its affiliates.
+#
+# 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.
+
+from types import SimpleNamespace
+from unittest.mock import Mock
+import pytest
+from google.genai import types
+from google.adk.models import LlmResponse
+
+
+@pytest.fixture
+def config(monkeypatch):
+    for k, v in {
+        "AGENTSENTRY_URL": "http://sentry/OpenTOP/V3/Lumen/Moderate",
+        "AGENTSENTRY_ACCOUNT_ID": "account",
+        "AGENTSENTRY_INSTANCE_ID": "instance",
+        "AGENTSENTRY_AK": "ak",
+        "AGENTSENTRY_SK": "sk",
+    }.items():
+        monkeypatch.setenv(k, v)
+
+
+@pytest.fixture
+def veadk(config):
+    from veadk.tools.builtin_tools.llm_shield import LLMShieldPlugin
+
+    shield = LLMShieldPlugin()
+    shield._sentry_v3.moderate = Mock(return_value=None)
+    return shield
+
+
+def test_all_input_parts_and_hook(veadk):
+    context = SimpleNamespace(agent_states={})
+    request = SimpleNamespace(
+        contents=[
+            types.Content(
+                role="user", parts=[types.Part(text="first"), types.Part(text="second")]
+            )
+        ]
+    )
+    assert veadk.before_model_callback(context, request) is None
+    args = veadk._sentry_v3.moderate.call_args
+    assert args.args[0] == "first\nsecond"
+    assert args.kwargs["hook_name"] == "before_dispatch"
+
+
+def test_plan_and_tool_are_linked_and_blocked(veadk):
+    ctx = SimpleNamespace(agent_states={}, function_call_id="call")
+    content = types.Content(
+        role="model",
+        parts=[
+            types.Part(text="first"),
+            types.Part(text="risk in second"),
+            types.Part(
+                function_call=types.FunctionCall(name="record", args={"action": "test"})
+            ),
+        ],
+    )
+    veadk.after_model_callback(ctx, LlmResponse(content=content))
+    veadk._sentry_v3.moderate.return_value = "blocked"
+    assert veadk.before_tool_callback(
+        SimpleNamespace(name="record"), {"action": "test"}, ctx
+    ) == {"result": "blocked"}
+    call = veadk._sentry_v3.moderate.call_args
+    assert call.args[0] == "first\nrisk in second"
+    assert call.kwargs["role"] == "assistant"
+    assert call.kwargs["tool_calls"][0]["Function"]["Name"] == "record"
+    # A different invocation may not use the first invocation's plan.
+    other = SimpleNamespace(agent_states={})
+    assert (
+        veadk.before_tool_callback(
+            SimpleNamespace(name="record"), {"action": "test"}, other
+        )["reason"]
+        == "missing_visible_execution_plan"
+    )
+
+
+def test_unsupported_outputs_not_relabelled(veadk):
+    ctx = SimpleNamespace(agent_states={})
+    assert (
+        veadk.after_model_callback(
+            ctx,
+            LlmResponse(
+                content=types.Content(role="model", parts=[types.Part(text="final")])
+            ),
+        )
+        is None
+    )
+    assert veadk.after_tool_callback(None, {}, ctx, "result") is None
+    veadk._sentry_v3.moderate.assert_not_called()

tests/test_agentsentry_v3_configuration.py

@@ -0,0 +1,68 @@
+# Copyright (c) 2025 Beijing Volcano Engine Technology Co., Ltd. and/or its affiliates.
+#
+# 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.
+
+import pytest
+
+
+@pytest.fixture
+def independent(monkeypatch):
+    for key in (
+        "TOOL_LLM_SHIELD_URL",
+        "TOOL_LLM_SHIELD_APP_ID",
+        "TOOL_LLM_SHIELD_API_KEY",
+        "TOOL_LLM_SHIELD_REGION",
+        "TOOL_LLM_SHIELD_ACCOUNT_ID",
+        "ENABLE_LLM_SHIELD",
+    ):
+        monkeypatch.delenv(key, raising=False)
+    for key, value in {
+        "AGENTSENTRY_URL": "http://sentry/OpenTOP/V3/Lumen/Moderate",
+        "AGENTSENTRY_ACCOUNT_ID": "account",
+        "AGENTSENTRY_INSTANCE_ID": "instance",
+        "AGENTSENTRY_AK": "ak",
+        "AGENTSENTRY_SK": "sk",
+    }.items():
+        monkeypatch.setenv(key, value)
+    monkeypatch.delenv("AGENTSENTRY_REGION", raising=False)
+
+
+@pytest.mark.parametrize("polluted", [False, True])
+def test_independent_configuration(independent, monkeypatch, polluted):
+    if polluted:
+        for key, value in {
+            "TOOL_LLM_SHIELD_URL": "https://public.example",
+            "TOOL_LLM_SHIELD_APP_ID": "cloud-app",
+            "TOOL_LLM_SHIELD_API_KEY": "cloud-secret",
+            "TOOL_LLM_SHIELD_REGION": "cloud-region",
+            "ENABLE_LLM_SHIELD": "false",
+        }.items():
+            monkeypatch.setenv(key, value)
+    from veadk.tools.builtin_tools.llm_shield import LLMShieldPlugin
+
+    plugin = LLMShieldPlugin()
+    assert plugin._sentry_v3 is not None
+    assert plugin._sentry_v3.account_id == "account"
+    assert plugin._sentry_v3.access_key == "ak"
+    assert plugin.region == "cn-private"
+    assert plugin.api_key == ""
+
+
+def test_invalid_private_url_never_falls_back(independent, monkeypatch):
+    from veadk.tools.builtin_tools.llm_shield import LLMShieldPlugin
+
+    monkeypatch.setenv("AGENTSENTRY_URL", "http://sentry/wrong-path")
+
+    factory = LLMShieldPlugin
+    with pytest.raises(ValueError, match="V3 Moderate URL"):
+        factory()