Add support for custom env vars (#131)
Some checks failed
Lint / lint (push) Has been cancelled
Test / test (3.10) (push) Has been cancelled
Test / test (3.11) (push) Has been cancelled
Test / test (3.12) (push) Has been cancelled
Test / test (3.13) (push) Has been cancelled

## 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
This commit is contained in:
Suzanne Wang 2025-08-25 14:02:03 -07:00 committed by GitHub
parent 30df222bfc
commit f794e17e78
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
4 changed files with 78 additions and 16 deletions

View file

@ -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():
@ -325,7 +324,9 @@ async def example_bash_command():
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

View file

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

View file

@ -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

View file

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