feat: refact agent shell output (#39190)

Co-authored-by: autofix-ci[bot] <114827586+autofix-ci[bot]@users.noreply.github.com>

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__ = [

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,

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,
             ),

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,
     )

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:

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):

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())

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()

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=