From f794e17e785ebd1729259e7db365912da4c97686 Mon Sep 17 00:00:00 2001 From: Suzanne Wang Date: Mon, 25 Aug 2025 14:02:03 -0700 Subject: [PATCH] Add support for custom env vars (#131) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit ## Key changes - Adds env field to `ClaudeCodeOptions`, allowing custom env vars to cli - Updates tests and examples ## Motivation Bringing Python SDK to feature parity with TS SDK, which supports custom env vars ## Notes - Environment variables are merged in order: system env → user env → SDK required vars - This implementation seems slightly more robust than the TypeScript version, which can exclude OS envs vars if a user passes a minimal env object - Some linting changes seem to have been picked up --- examples/streaming_mode.py | 31 +++++------ .../_internal/transport/subprocess_cli.py | 9 +++- src/claude_code_sdk/types.py | 1 + tests/test_transport.py | 53 +++++++++++++++++++ 4 files changed, 78 insertions(+), 16 deletions(-) diff --git a/examples/streaming_mode.py b/examples/streaming_mode.py index da2bc1b..29ed9e3 100755 --- a/examples/streaming_mode.py +++ b/examples/streaming_mode.py @@ -188,9 +188,7 @@ async def example_manual_message_handling(): print("=== Manual Message Handling Example ===") async with ClaudeSDKClient() as client: - await client.query( - "List 5 programming languages and their main use cases" - ) + await client.query("List 5 programming languages and their main use cases") # Manually process messages with custom logic languages_found = [] @@ -231,13 +229,14 @@ async def example_with_options(): allowed_tools=["Read", "Write"], # Allow file operations max_thinking_tokens=10000, system_prompt="You are a helpful coding assistant.", + env={ + "ANTHROPIC_MODEL": "claude-3-7-sonnet-20250219", + }, ) async with ClaudeSDKClient(options=options) as client: print("User: Create a simple hello.txt file with a greeting message") - await client.query( - "Create a simple hello.txt file with a greeting message" - ) + await client.query("Create a simple hello.txt file with a greeting message") tool_uses = [] async for msg in client.receive_response(): @@ -308,25 +307,27 @@ async def example_async_iterable_prompt(): async def example_bash_command(): """Example showing tool use blocks when running bash commands.""" print("=== Bash Command Example ===") - + async with ClaudeSDKClient() as client: print("User: Run a bash echo command") await client.query("Run a bash echo command that says 'Hello from bash!'") - + # Track all message types received message_types = [] - + async for msg in client.receive_messages(): message_types.append(type(msg).__name__) - + if isinstance(msg, UserMessage): # User messages can contain tool results for block in msg.content: if isinstance(block, TextBlock): print(f"User: {block.text}") elif isinstance(block, ToolResultBlock): - print(f"Tool Result (id: {block.tool_use_id}): {block.content[:100] if block.content else 'None'}...") - + print( + f"Tool Result (id: {block.tool_use_id}): {block.content[:100] if block.content else 'None'}..." + ) + elif isinstance(msg, AssistantMessage): # Assistant messages can contain tool use blocks for block in msg.content: @@ -337,15 +338,15 @@ async def example_bash_command(): if block.name == "Bash": command = block.input.get("command", "") print(f" Command: {command}") - + elif isinstance(msg, ResultMessage): print("Result ended") if msg.total_cost_usd: print(f"Cost: ${msg.total_cost_usd:.4f}") break - + print(f"\nMessage types received: {', '.join(set(message_types))}") - + print("\n") diff --git a/src/claude_code_sdk/_internal/transport/subprocess_cli.py b/src/claude_code_sdk/_internal/transport/subprocess_cli.py index 66fb6eb..47bcad7 100644 --- a/src/claude_code_sdk/_internal/transport/subprocess_cli.py +++ b/src/claude_code_sdk/_internal/transport/subprocess_cli.py @@ -175,13 +175,20 @@ class SubprocessCLITransport(Transport): ) # Enable stdin pipe for both modes (but we'll close it for string mode) + # Merge environment variables: system -> user -> SDK required + process_env = { + **os.environ, + **self._options.env, # User-provided env vars + "CLAUDE_CODE_ENTRYPOINT": "sdk-py", + } + self._process = await anyio.open_process( cmd, stdin=PIPE, stdout=PIPE, stderr=self._stderr_file, cwd=self._cwd, - env={**os.environ, "CLAUDE_CODE_ENTRYPOINT": "sdk-py"}, + env=process_env, ) if self._process.stdout: diff --git a/src/claude_code_sdk/types.py b/src/claude_code_sdk/types.py index 98afdef..2c52907 100644 --- a/src/claude_code_sdk/types.py +++ b/src/claude_code_sdk/types.py @@ -137,6 +137,7 @@ class ClaudeCodeOptions: cwd: str | Path | None = None settings: str | None = None add_dirs: list[str | Path] = field(default_factory=list) + env: dict[str, str] = field(default_factory=dict) extra_args: dict[str, str | None] = field( default_factory=dict ) # Pass arbitrary CLI flags diff --git a/tests/test_transport.py b/tests/test_transport.py index 6ab1363..aa6a8e9 100644 --- a/tests/test_transport.py +++ b/tests/test_transport.py @@ -1,5 +1,7 @@ """Tests for Claude SDK transport layer.""" +import os +import uuid from unittest.mock import AsyncMock, MagicMock, patch import anyio @@ -299,3 +301,54 @@ class TestSubprocessCLITransport: assert "--mcp-config" in cmd mcp_idx = cmd.index("--mcp-config") assert cmd[mcp_idx + 1] == json_config + + def test_env_vars_passed_to_subprocess(self): + """Test that custom environment variables are passed to the subprocess.""" + + async def _test(): + test_value = f"test-{uuid.uuid4().hex[:8]}" + custom_env = { + "MY_TEST_VAR": test_value, + } + + options = ClaudeCodeOptions(env=custom_env) + + # Mock the subprocess to capture the env argument + with patch( + "anyio.open_process", new_callable=AsyncMock + ) as mock_open_process: + mock_process = MagicMock() + mock_process.stdout = MagicMock() + mock_stdin = MagicMock() + mock_stdin.aclose = AsyncMock() # Add async aclose method + mock_process.stdin = mock_stdin + mock_process.returncode = None + mock_open_process.return_value = mock_process + + transport = SubprocessCLITransport( + prompt="test", + options=options, + cli_path="/usr/bin/claude", + ) + + await transport.connect() + + # Verify open_process was called with correct env vars + mock_open_process.assert_called_once() + call_kwargs = mock_open_process.call_args.kwargs + assert "env" in call_kwargs + env_passed = call_kwargs["env"] + + # Check that custom env var was passed + assert env_passed["MY_TEST_VAR"] == test_value + + # Verify SDK identifier is present + assert "CLAUDE_CODE_ENTRYPOINT" in env_passed + assert env_passed["CLAUDE_CODE_ENTRYPOINT"] == "sdk-py" + + # Verify system env vars are also included with correct values + if "PATH" in os.environ: + assert "PATH" in env_passed + assert env_passed["PATH"] == os.environ["PATH"] + + anyio.run(_test)