Skip to content
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
44 changes: 37 additions & 7 deletions peak_assistant/streamlit/util/helpers.py
Original file line number Diff line number Diff line change
Expand Up @@ -23,15 +23,14 @@
from typing import List, Dict, Any, Optional, Tuple
from autogen_agentchat.messages import TextMessage, UserMessage
import streamlit as st
import hashlib
import html
import json
import os
import secrets
import tempfile
import time
import logging
from urllib.parse import urlparse, urljoin
from urllib.parse import urlparse
from pathlib import Path

# Import MCP configuration classes from centralized location
Expand All @@ -54,6 +53,39 @@
"token",
}

MCP_SUBPROCESS_DEFAULT_ENV_VARS = (
"PATH",
"HOME",
"USER",
"LOGNAME",
"SHELL",
"TMPDIR",
"TEMP",
"TMP",
"SystemRoot",
"COMSPEC",
"PATHEXT",
)


def build_mcp_subprocess_env(server_env: Optional[Dict[str, str]] = None) -> Dict[str, str]:
"""Build a minimal environment for stdio MCP subprocesses.

Streamlit loads .env secrets into the process environment, so copying all of
os.environ into test subprocesses can expose unrelated LLM API keys, OAuth
secrets, and other PEAK credentials to untrusted MCP packages. Keep only the
small set of platform defaults needed to locate and run local commands, then
add the explicit per-server environment from mcp_servers.json.
"""
env = {
key: os.environ[key]
for key in MCP_SUBPROCESS_DEFAULT_ENV_VARS
if key in os.environ
}
if server_env:
env.update(server_env)
return env


def validate_and_escape_oauth_url(url: str) -> Optional[str]:
"""
Expand Down Expand Up @@ -725,10 +757,8 @@ async def test_mcp_connection(server_name: str, server_config: MCPServerConfig)
try:
from autogen_ext.tools.mcp import McpWorkbench, StdioServerParams

# Build complete subprocess environment (mirrors CLI in mcp_config.py)
env = os.environ.copy()
if server_config.env:
env.update(server_config.env)
# Build a minimal subprocess environment plus explicit server env.
env = build_mcp_subprocess_env(server_config.env)

server_params = StdioServerParams(
command=server_config.command,
Expand Down Expand Up @@ -1147,7 +1177,7 @@ def get_agent_config_data() -> List[Dict[str, str]]:
try:
provider_config = loader.get_provider_config(provider_name)
provider_type = provider_config["type"]
except:
except Exception:
provider_type = "unknown"

agent_data.append({
Expand Down
15 changes: 10 additions & 5 deletions tests/unit_tests/test_streamlit_helpers_mcp_bugs.py
Original file line number Diff line number Diff line change
Expand Up @@ -24,7 +24,7 @@
Tests for Streamlit helpers MCP configuration bugs.

Bug 1: load_mcp_server_configs() does not interpolate ${ENV_VAR} patterns.
Bug 2: test_mcp_connection() does not include system env vars in subprocess.
Bug 2: test_mcp_connection() leaks full process env to subprocesses.
"""

import json
Expand Down Expand Up @@ -139,12 +139,13 @@ def test_missing_env_var_without_default_returns_empty(self, monkeypatch):


class TestMcpConnectionSubprocessEnv:
"""Bug 2: test_mcp_connection() should pass full system env to subprocess"""
"""Bug 2: test_mcp_connection() should avoid leaking full process env"""

@pytest.mark.asyncio
async def test_subprocess_env_includes_system_path(self, monkeypatch):
"""StdioServerParams.env should contain both custom and system env vars"""
"""StdioServerParams.env should contain custom env and safe default vars"""
monkeypatch.setenv("PATH", "/usr/bin:/usr/local/bin")
monkeypatch.setenv("UNRELATED_OAUTH_CLIENT_SECRET", "oauth-secret-not-for-mcp")

config = MCPServerConfig(
name="test-server",
Expand Down Expand Up @@ -175,11 +176,13 @@ def capture_workbench(server_params):
assert params.env["CUSTOM_KEY"] == "custom_val"
assert "PATH" in params.env
assert params.env["PATH"] == "/usr/bin:/usr/local/bin"
assert "UNRELATED_OAUTH_CLIENT_SECRET" not in params.env

@pytest.mark.asyncio
async def test_subprocess_env_with_no_config_env_gets_system_env(self, monkeypatch):
"""Even with env=None in config, subprocess should get system env"""
async def test_subprocess_env_with_no_config_env_gets_safe_default_env(self, monkeypatch):
"""Even with env=None in config, subprocess should only get safe defaults"""
monkeypatch.setenv("PATH", "/usr/bin:/usr/local/bin")
monkeypatch.setenv("PEAK_GLOBAL_SECRET", "global-secret-not-for-mcp")

config = MCPServerConfig(
name="test-server",
Expand All @@ -206,3 +209,5 @@ def capture_workbench(server_params):
assert success is True
params = captured_params["server_params"]
assert "PATH" in params.env
assert params.env["PATH"] == "/usr/bin:/usr/local/bin"
assert "PEAK_GLOBAL_SECRET" not in params.env
Loading