From 6e5fc1081bb8739adc4e8078210a86d80bd6953a Mon Sep 17 00:00:00 2001 From: Yunlu Wen Date: Fri, 17 Jul 2026 16:27:00 +0800 Subject: [PATCH] feat: refact agent shell output (#39190) Co-authored-by: autofix-ci[bot] <114827586+autofix-ci[bot]@users.noreply.github.com> --- .../src/dify_agent/layers/shell/configs.py | 1 + .../src/dify_agent/layers/shell/layer.py | 36 +++- .../dify_agent/runtime/compositor_factory.py | 2 + dify-agent/src/dify_agent/server/app.py | 1 + dify-agent/src/dify_agent/server/settings.py | 13 ++ .../dify_agent/layers/shell/test_configs.py | 1 + .../dify_agent/layers/shell/test_layer.py | 158 ++++++++++++++++++ .../local/dify_agent/server/test_settings.py | 37 ++++ .../envs/core-services/dify-agent.env.example | 5 + 9 files changed, 250 insertions(+), 4 deletions(-) diff --git a/dify-agent/src/dify_agent/layers/shell/configs.py b/dify-agent/src/dify_agent/layers/shell/configs.py index 3a195826744..63f1faf5d5e 100644 --- a/dify-agent/src/dify_agent/layers/shell/configs.py +++ b/dify-agent/src/dify_agent/layers/shell/configs.py @@ -94,6 +94,7 @@ class DifyShellLayerConfig(LayerConfig): env: list[DifyShellEnvVarConfig] = Field(default_factory=list) secret_refs: list[DifyShellSecretRefConfig] = Field(default_factory=list) sandbox: DifyShellSandboxConfig | None = None + redact_patterns: list[str] = Field(default_factory=list) __all__ = [ diff --git a/dify-agent/src/dify_agent/layers/shell/layer.py b/dify-agent/src/dify_agent/layers/shell/layer.py index 902d5ddc0fa..b38a151ddee 100644 --- a/dify-agent/src/dify_agent/layers/shell/layer.py +++ b/dify-agent/src/dify_agent/layers/shell/layer.py @@ -33,7 +33,7 @@ import logging import re import secrets import time -from dataclasses import dataclass +from dataclasses import dataclass, field from typing import ClassVar, Literal, NotRequired, Protocol, TypedDict, runtime_checkable from pydantic import BaseModel, ConfigDict, Field, NonNegativeInt, field_validator, model_validator @@ -57,6 +57,7 @@ from dify_agent.adapters.shell.protocols import ( ShellProviderProtocol, ShellResourceProtocol, ) +from dify_agent.agent_stub.protocol import AGENT_STUB_AUTH_JWE_ENV_VAR from dify_agent.agent_stub.shell_env import ShellAgentStubTokenFactory, build_shell_agent_stub_env from dify_agent.layers.execution_context import DifyExecutionContextLayerConfig from dify_agent.layers.shell.configs import DIFY_SHELL_LAYER_TYPE_ID, DifyShellLayerConfig @@ -251,6 +252,7 @@ class DifyShellLayer(PydanticAILayer[DifyShellLayerDeps, object, DifyShellLayerC config: DifyShellLayerConfig shell_provider: ShellProviderProtocol shell_home_root: str = "/home" + shell_redact_patterns: list[str] = field(default_factory=list) agent_stub_api_base_url: str | None = None agent_stub_token_factory: ShellAgentStubTokenFactory | None = None _shell_resource: ShellResourceProtocol | None = None @@ -269,6 +271,7 @@ class DifyShellLayer(PydanticAILayer[DifyShellLayerDeps, object, DifyShellLayerC *, shell_provider: ShellProviderProtocol | None, shell_home_root: str = "/home", + shell_redact_patterns: list[str] | None = None, agent_stub_api_base_url: str | None = None, agent_stub_token_factory: ShellAgentStubTokenFactory | None = None, ) -> Self: @@ -278,6 +281,7 @@ class DifyShellLayer(PydanticAILayer[DifyShellLayerDeps, object, DifyShellLayerC config=config, shell_provider=shell_provider, shell_home_root=_normalize_shell_home_root(shell_home_root), + shell_redact_patterns=shell_redact_patterns or [], agent_stub_api_base_url=agent_stub_api_base_url, agent_stub_token_factory=agent_stub_token_factory, ) @@ -417,7 +421,7 @@ class DifyShellLayer(PydanticAILayer[DifyShellLayerDeps, object, DifyShellLayerC exit_code=result.exit_code, output_path=observation.output_path, ), - observation.text, + self._redact_output(observation.text), ) except (RuntimeError, ValueError) as exc: return _tool_error_from_exception(exc) @@ -443,7 +447,7 @@ class DifyShellLayer(PydanticAILayer[DifyShellLayerDeps, object, DifyShellLayerC exit_code=result.exit_code, output_path=observation.output_path, ), - observation.text, + self._redact_output(observation.text), ) except (RuntimeError, ValueError) as exc: return _tool_error_from_exception(exc, job_id=job_id) @@ -469,7 +473,7 @@ class DifyShellLayer(PydanticAILayer[DifyShellLayerDeps, object, DifyShellLayerC exit_code=result.exit_code, output_path=observation.output_path, ), - observation.text, + self._redact_output(observation.text), ) except (RuntimeError, ValueError) as exc: return _tool_error_from_exception(exc, job_id=job_id) @@ -718,6 +722,30 @@ class DifyShellLayer(PydanticAILayer[DifyShellLayerDeps, object, DifyShellLayerC env.update(agent_stub_env) return env + def _redact_output(self, text: str) -> str: + """Redact sensitive content from shell output before the model sees it. + + Two layers of redaction are applied: + + 1. **Built-in token redaction** — the actual Agent Stub JWE token value + is always replaced with ``***``. This is unconditional and cannot be + disabled. + 2. **Pattern redaction** — regex patterns from both server-level + ``shell_redact_patterns`` and per-agent ``config.redact_patterns`` + are applied via ``re.sub`` to mask additional secrets. + """ + if not text: + return text + # Built-in: always redact the JWE token value. + env = self._build_shell_command_env(include_agent_stub_env=True) + jwe_value = env.get(AGENT_STUB_AUTH_JWE_ENV_VAR) + if jwe_value and len(jwe_value) > 8: + text = text.replace(jwe_value, "***") + # Server-level + per-agent regex patterns. + for pattern in (*self.shell_redact_patterns, *self.config.redact_patterns): + text = re.sub(pattern, "***", text) + return text + async def execute_complete_with_commands( commands: ShellCommandProtocol, diff --git a/dify-agent/src/dify_agent/runtime/compositor_factory.py b/dify-agent/src/dify_agent/runtime/compositor_factory.py index a341195bdff..b81cf737478 100644 --- a/dify-agent/src/dify_agent/runtime/compositor_factory.py +++ b/dify-agent/src/dify_agent/runtime/compositor_factory.py @@ -69,6 +69,7 @@ def create_default_layer_providers( inner_api_key: str = "", shell_provider: ShellProviderProtocol | None = None, shell_home_root: str = "/home", + shell_redact_patterns: list[str] | None = None, agent_stub_api_base_url: str | None = None, agent_stub_token_factory: ShellAgentStubTokenFactory | None = None, ) -> tuple[DifyAgentLayerProvider, ...]: @@ -94,6 +95,7 @@ def create_default_layer_providers( DifyShellLayerConfig.model_validate(config), shell_provider=shell_provider, shell_home_root=shell_home_root, + shell_redact_patterns=shell_redact_patterns or [], agent_stub_api_base_url=agent_stub_api_base_url, agent_stub_token_factory=agent_stub_token_factory, ), diff --git a/dify-agent/src/dify_agent/server/app.py b/dify-agent/src/dify_agent/server/app.py index 918a9292416..1013cf89bb0 100644 --- a/dify-agent/src/dify_agent/server/app.py +++ b/dify-agent/src/dify_agent/server/app.py @@ -67,6 +67,7 @@ def create_app(settings: ServerSettings | None = None) -> FastAPI: inner_api_key=resolved_settings.inner_api_key or "", shell_provider=shell_provider, shell_home_root=resolved_settings.shell_home_root, + shell_redact_patterns=resolved_settings.get_shell_redact_patterns(), agent_stub_api_base_url=resolved_settings.agent_stub_api_base_url, agent_stub_token_factory=agent_stub_token_factory, ) diff --git a/dify-agent/src/dify_agent/server/settings.py b/dify-agent/src/dify_agent/server/settings.py index 05174cb2637..d6e8dfbcb4f 100644 --- a/dify-agent/src/dify_agent/server/settings.py +++ b/dify-agent/src/dify_agent/server/settings.py @@ -52,6 +52,7 @@ class ServerSettings(BaseSettings): agent_stub_api_base_url: str | None = Field(default=None, validation_alias="DIFY_AGENT_STUB_API_BASE_URL") agent_stub_grpc_bind_address: str | None = Field(default=None, validation_alias="DIFY_AGENT_STUB_GRPC_BIND_ADDRESS") server_secret_key: str | None = None + shell_redact_patterns: str = "" outbound_http_connect_timeout: float = Field(default=10.0, ge=0) outbound_http_read_timeout: float = Field(default=600.0, ge=0) outbound_http_write_timeout: float = Field(default=30.0, ge=0) @@ -126,6 +127,18 @@ class ServerSettings(BaseSettings): stripped = value.strip() return stripped or None + def get_shell_redact_patterns(self) -> list[str]: + """Parse the JSON array from shell_redact_patterns; empty/blank → empty list.""" + stripped = self.shell_redact_patterns.strip() + if not stripped: + return [] + import json as _json + + parsed = _json.loads(stripped) + if not isinstance(parsed, list): + raise ValueError("DIFY_AGENT_SHELL_REDACT_PATTERNS must be a JSON array of strings") + return [str(p) for p in parsed] + @field_validator("shell_home_root") @classmethod def normalize_shell_home_root(cls, value: str) -> str: diff --git a/dify-agent/tests/local/dify_agent/layers/shell/test_configs.py b/dify-agent/tests/local/dify_agent/layers/shell/test_configs.py index ec429f94339..b6ad74965b9 100644 --- a/dify-agent/tests/local/dify_agent/layers/shell/test_configs.py +++ b/dify-agent/tests/local/dify_agent/layers/shell/test_configs.py @@ -34,6 +34,7 @@ def test_shell_layer_config_defaults_and_forbids_unknown_fields() -> None: "env": [], "secret_refs": [], "sandbox": None, + "redact_patterns": [], } with pytest.raises(ValidationError): diff --git a/dify-agent/tests/local/dify_agent/layers/shell/test_layer.py b/dify-agent/tests/local/dify_agent/layers/shell/test_layer.py index 274e779c2c4..3d1e8f292cd 100644 --- a/dify-agent/tests/local/dify_agent/layers/shell/test_layer.py +++ b/dify-agent/tests/local/dify_agent/layers/shell/test_layer.py @@ -1378,3 +1378,161 @@ def test_resource_context_reraises_non_expired_attach_error() -> None: with pytest.raises(ShellProviderError, match="some other error"): asyncio.run(scenario()) + + +# --------------------------------------------------------------------------- +# Output redaction tests +# --------------------------------------------------------------------------- + + +def _layer_with_redaction( + *, + commands: FakeCommands, + config: DifyShellLayerConfig | None = None, + shell_redact_patterns: list[str] | None = None, + token_value: str = "eyJhbGciOiJkaXIiLCJlbmMiOiJBMjU2R0NNIn0.fake-long-jwe-token-value", +) -> tuple[DifyShellLayer, FakeProvider]: + """Create a layer with agent_stub env injection and optional redaction patterns.""" + provider = FakeProvider(resource=FakeResource(commands=commands)) + layer = DifyShellLayer.from_config_with_settings( + config or DifyShellLayerConfig(), + shell_provider=provider, + shell_home_root="/home", + shell_redact_patterns=shell_redact_patterns, + agent_stub_api_base_url="http://localhost:5050/agent-stub", + agent_stub_token_factory=lambda execution_context, session_id: token_value, + ) + return layer, provider + + +def test_redact_output_replaces_jwe_token_value() -> None: + """The JWE token value should always be redacted from shell output.""" + token = "eyJhbGciOiJkaXIiLCJlbmMiOiJBMjU2R0NNIn0.super-secret-token-12345" + + def run_handler(script: str, cwd: str | None, env: Mapping[str, str] | None, timeout: float) -> ShellCommandResult: + return _command_result( + "job-1", + status="exited", + done=True, + exit_code=0, + output=f"DIFY_AGENT_STUB_AUTH_JWE={token}\n", + offset=100, + ) + + commands = FakeCommands( + run_handler=run_handler, + tail_handler=lambda _: _command_result("job-1", done=True, exit_code=0, status="exited", offset=100), + ) + layer, _provider = _layer_with_redaction(commands=commands, token_value=token) + _bind_execution_context(layer) + layer.runtime_state = _runtime_state() + tools = {tool.name: tool for tool in layer.tools} + + async def scenario() -> None: + async with layer.resource_context(): + result = await tools["shell_run"].function_schema.call({"script": "env"}, None) # pyright: ignore[reportArgumentType] + _, output = _parse_tagged_observation(result) + assert token not in output + assert "***" in output + + asyncio.run(scenario()) + + +def test_redact_output_applies_server_level_patterns() -> None: + """Server-level regex patterns from env var should redact matching content.""" + + def run_handler(script: str, cwd: str | None, env: Mapping[str, str] | None, timeout: float) -> ShellCommandResult: + return _command_result( + "job-1", + status="exited", + done=True, + exit_code=0, + output="api_key=sk-proj-abc123xyz\n", + offset=30, + ) + + commands = FakeCommands( + run_handler=run_handler, + tail_handler=lambda _: _command_result("job-1", done=True, exit_code=0, status="exited", offset=30), + ) + layer, _provider = _layer_with_redaction( + commands=commands, + shell_redact_patterns=[r"sk-proj-[A-Za-z0-9]+"], + ) + _bind_execution_context(layer) + layer.runtime_state = _runtime_state() + tools = {tool.name: tool for tool in layer.tools} + + async def scenario() -> None: + async with layer.resource_context(): + result = await tools["shell_run"].function_schema.call({"script": "cat .env"}, None) # pyright: ignore[reportArgumentType] + _, output = _parse_tagged_observation(result) + assert "sk-proj-abc123xyz" not in output + assert "api_key=***" in output + + asyncio.run(scenario()) + + +def test_redact_output_applies_per_agent_config_patterns() -> None: + """Per-agent redact_patterns from DifyShellLayerConfig should also apply.""" + + def run_handler(script: str, cwd: str | None, env: Mapping[str, str] | None, timeout: float) -> ShellCommandResult: + return _command_result( + "job-1", + status="exited", + done=True, + exit_code=0, + output="token: ghp_aBcDeFgHiJkLmNoPqRsTuVwXyZ1234567890\n", + offset=50, + ) + + commands = FakeCommands( + run_handler=run_handler, + tail_handler=lambda _: _command_result("job-1", done=True, exit_code=0, status="exited", offset=50), + ) + config = DifyShellLayerConfig(redact_patterns=[r"ghp_[A-Za-z0-9]{36}"]) + layer, _provider = _layer_with_redaction(commands=commands, config=config) + _bind_execution_context(layer) + layer.runtime_state = _runtime_state() + tools = {tool.name: tool for tool in layer.tools} + + async def scenario() -> None: + async with layer.resource_context(): + result = await tools["shell_run"].function_schema.call({"script": "echo $TOKEN"}, None) # pyright: ignore[reportArgumentType] + _, output = _parse_tagged_observation(result) + assert "ghp_" not in output + assert "token: ***" in output + + asyncio.run(scenario()) + + +def test_redact_output_skips_short_jwe_values() -> None: + """JWE values ≤8 chars should not be redacted to avoid false positives.""" + + def run_handler(script: str, cwd: str | None, env: Mapping[str, str] | None, timeout: float) -> ShellCommandResult: + return _command_result( + "job-1", + status="exited", + done=True, + exit_code=0, + output="short\n", + offset=6, + ) + + commands = FakeCommands( + run_handler=run_handler, + tail_handler=lambda _: _command_result("job-1", done=True, exit_code=0, status="exited", offset=6), + ) + # Token value is short — should NOT be redacted even if it appears in output. + layer, _provider = _layer_with_redaction(commands=commands, token_value="short") + _bind_execution_context(layer) + layer.runtime_state = _runtime_state() + tools = {tool.name: tool for tool in layer.tools} + + async def scenario() -> None: + async with layer.resource_context(): + result = await tools["shell_run"].function_schema.call({"script": "echo hi"}, None) # pyright: ignore[reportArgumentType] + _, output = _parse_tagged_observation(result) + assert "short" in output + + asyncio.run(scenario()) diff --git a/dify-agent/tests/local/dify_agent/server/test_settings.py b/dify-agent/tests/local/dify_agent/server/test_settings.py index 3233ab32cb5..25ae693963b 100644 --- a/dify-agent/tests/local/dify_agent/server/test_settings.py +++ b/dify-agent/tests/local/dify_agent/server/test_settings.py @@ -306,3 +306,40 @@ def test_build_shell_provider_returns_none_when_enterprise_endpoint_is_unset( def test_build_shell_provider_rejects_blank_shellctl_entrypoint() -> None: with pytest.raises(ValidationError, match="shellctl_entrypoint is required"): _ = ServerSettings(shell_provider="shellctl", shellctl_entrypoint=" ").build_shell_provider() + + +def test_server_settings_parses_shell_redact_patterns_json_array(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setenv("DIFY_AGENT_SHELL_REDACT_PATTERNS", '["sk-[A-Za-z0-9]+","ghp_[A-Za-z0-9]{36}"]') + + settings = ServerSettings() + + assert settings.get_shell_redact_patterns() == ["sk-[A-Za-z0-9]+", "ghp_[A-Za-z0-9]{36}"] + + +def test_server_settings_shell_redact_patterns_empty_string_yields_empty_list(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setenv("DIFY_AGENT_SHELL_REDACT_PATTERNS", "") + + settings = ServerSettings() + + assert settings.get_shell_redact_patterns() == [] + + +def test_server_settings_shell_redact_patterns_defaults_to_empty_list( + monkeypatch: pytest.MonkeyPatch, + tmp_path: Path, +) -> None: + monkeypatch.delenv("DIFY_AGENT_SHELL_REDACT_PATTERNS", raising=False) + monkeypatch.chdir(tmp_path) + + settings = ServerSettings() + + assert settings.get_shell_redact_patterns() == [] + + +def test_server_settings_rejects_non_array_shell_redact_patterns(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setenv("DIFY_AGENT_SHELL_REDACT_PATTERNS", '{"key": "value"}') + + settings = ServerSettings() + + with pytest.raises(ValueError, match="must be a JSON array"): + settings.get_shell_redact_patterns() diff --git a/docker/envs/core-services/dify-agent.env.example b/docker/envs/core-services/dify-agent.env.example index 52ac6bc0990..f2d7155f521 100644 --- a/docker/envs/core-services/dify-agent.env.example +++ b/docker/envs/core-services/dify-agent.env.example @@ -29,3 +29,8 @@ DIFY_AGENT_STUB_API_BASE_URL=http://agent_backend:5050/agent-stub # Replace this development default in production. # Generate one with: python -c 'import secrets; print(secrets.token_urlsafe(32))' DIFY_AGENT_SERVER_SECRET_KEY=MDEyMzQ1Njc4OWFiY2RlZjAxMjM0NTY3ODlhYmNkZWY + +# JSON array of regex patterns to redact from shell output shown to the agent. +# The JWE token value is always redacted regardless of this setting. +# Example: DIFY_AGENT_SHELL_REDACT_PATTERNS=["sk-[A-Za-z0-9]+","ghp_[A-Za-z0-9]{36}"] +DIFY_AGENT_SHELL_REDACT_PATTERNS=