From c694fb01af92299b9685790af55697b4315675b3 Mon Sep 17 00:00:00 2001 From: openhands Date: Tue, 24 Mar 2026 12:27:06 +0000 Subject: [PATCH] fix: support AWS IAM credentials for Bedrock models without LLM_API_KEY When using --override-with-envs, AWS Bedrock/SageMaker models no longer require LLM_API_KEY. These models use AWS IAM credentials instead, which can be provided via: - Environment variables: AWS_ACCESS_KEY_ID, AWS_SECRET_ACCESS_KEY, AWS_REGION_NAME - Default AWS credential chain (boto3): ~/.aws/credentials, IAM roles, etc. Changes: - Add AWS credential support to LLMEnvOverrides (aws_access_key_id, aws_secret_access_key, aws_region_name) - Add is_aws_auth_model() helper to detect bedrock/, bedrock_converse/, sagemaker/ model prefixes - Update require_for_headless() to skip LLM_API_KEY validation for AWS models - Update _ensure_agent() to build LLM kwargs dynamically, allowing api_key to be None for AWS-authenticated models - Improve error messages with AWS-specific guidance - Add comprehensive tests for AWS authentication path Fixes #611 Co-authored-by: openhands --- openhands_cli/stores/__init__.py | 14 ++ openhands_cli/stores/agent_store.py | 158 +++++++++++++++--- tests/stores/test_env_llm_overrides.py | 221 +++++++++++++++++++++++++ 3 files changed, 372 insertions(+), 21 deletions(-) diff --git a/openhands_cli/stores/__init__.py b/openhands_cli/stores/__init__.py index 6d052389e..8f27f5b94 100644 --- a/openhands_cli/stores/__init__.py +++ b/openhands_cli/stores/__init__.py @@ -1,7 +1,14 @@ from openhands_cli.stores.agent_store import ( + ENV_AWS_ACCESS_KEY_ID, + ENV_AWS_REGION_NAME, + ENV_AWS_SECRET_ACCESS_KEY, + ENV_LLM_API_KEY, + ENV_LLM_BASE_URL, + ENV_LLM_MODEL, AgentStore, MissingEnvironmentVariablesError, check_and_warn_env_vars, + is_aws_auth_model, ) from openhands_cli.stores.cli_settings import ( DEFAULT_MAX_REFINEMENT_ITERATIONS, @@ -15,6 +22,13 @@ "CliSettings", "CriticSettings", "DEFAULT_MAX_REFINEMENT_ITERATIONS", + "ENV_AWS_ACCESS_KEY_ID", + "ENV_AWS_REGION_NAME", + "ENV_AWS_SECRET_ACCESS_KEY", + "ENV_LLM_API_KEY", + "ENV_LLM_BASE_URL", + "ENV_LLM_MODEL", "MissingEnvironmentVariablesError", "check_and_warn_env_vars", + "is_aws_auth_model", ] diff --git a/openhands_cli/stores/agent_store.py b/openhands_cli/stores/agent_store.py index 2cf657691..26b3ff93b 100644 --- a/openhands_cli/stores/agent_store.py +++ b/openhands_cli/stores/agent_store.py @@ -126,23 +126,73 @@ def get_default_critic(llm: LLM, *, enable_critic: bool = True) -> CriticBase | ENV_LLM_BASE_URL = "LLM_BASE_URL" ENV_LLM_MODEL = "LLM_MODEL" +# AWS credential environment variables (standard AWS SDK names) +ENV_AWS_ACCESS_KEY_ID = "AWS_ACCESS_KEY_ID" +ENV_AWS_SECRET_ACCESS_KEY = "AWS_SECRET_ACCESS_KEY" +ENV_AWS_REGION_NAME = "AWS_REGION_NAME" + +# Model prefixes that use AWS IAM authentication instead of API key +AWS_AUTH_MODEL_PREFIXES = ("bedrock/", "bedrock_converse/", "sagemaker/") + + +def is_aws_auth_model(model: str | None) -> bool: + """Check if the model uses AWS IAM authentication instead of API key. + + AWS Bedrock and SageMaker models use IAM credentials from: + - Explicit env vars: AWS_ACCESS_KEY_ID, AWS_SECRET_ACCESS_KEY, AWS_REGION_NAME + - Default AWS credential chain (boto3): ~/.aws/credentials, IAM roles, etc. + + Args: + model: The model string (e.g., "bedrock/anthropic.claude-3-sonnet") + + Returns: + True if the model uses AWS authentication, False otherwise. + """ + if model is None: + return False + return model.startswith(AWS_AUTH_MODEL_PREFIXES) + class MissingEnvironmentVariablesError(Exception): """Raised when required environment variables are missing for headless mode. This exception is raised when --override-with-envs is enabled but required - environment variables (LLM_API_KEY and LLM_MODEL) are not set. + environment variables (LLM_API_KEY and LLM_MODEL, or AWS credentials for + AWS-authenticated models) are not set. """ - def __init__(self, missing_vars: list[str]) -> None: + def __init__(self, missing_vars: list[str], *, is_aws_model: bool = False) -> None: self.missing_vars = missing_vars + self.is_aws_model = is_aws_model vars_str = ", ".join(missing_vars) - super().__init__( - f"Missing required environment variable(s): {vars_str}\n" - f"When using --override-with-envs, you must set:\n" - f" - {ENV_LLM_API_KEY}: Your LLM API key\n" - f" - {ENV_LLM_MODEL}: The model to use (e.g., claude-sonnet-4-5-20250929)" - ) + + if is_aws_model: + # AWS model - API key not required, but model is + super().__init__( + f"Missing required environment variable(s): {vars_str}\n" + "When using --override-with-envs with AWS Bedrock/SageMaker:\n" + f" - {ENV_LLM_MODEL}: The model to use " + "(e.g., bedrock/anthropic.claude-3-sonnet)\n" + "\n" + "AWS credentials are obtained from the standard credential chain:\n" + " - Environment variables: AWS_ACCESS_KEY_ID, AWS_SECRET_ACCESS_KEY,\n" + " AWS_REGION_NAME\n" + " - AWS credentials file: ~/.aws/credentials\n" + " - IAM roles (for EC2, ECS, Lambda, etc.)" + ) + else: + # Standard model - API key required + super().__init__( + f"Missing required environment variable(s): {vars_str}\n" + "When using --override-with-envs, you must set:\n" + f" - {ENV_LLM_API_KEY}: Your LLM API key\n" + f" - {ENV_LLM_MODEL}: The model to use " + "(e.g., claude-sonnet-4-5-20250929)\n" + "\n" + "For AWS Bedrock/SageMaker models, LLM_API_KEY is not required.\n" + "Use a model prefix like 'bedrock/' " + "(e.g., bedrock/anthropic.claude-3-sonnet)." + ) def check_and_warn_env_vars() -> None: @@ -179,11 +229,19 @@ class LLMEnvOverrides(BaseModel): Use the `from_env()` class method to load values from environment variables when env overrides are enabled. + + For AWS Bedrock/SageMaker models: + - api_key is not required + - AWS credentials can be provided via env vars or the default AWS credential chain """ api_key: SecretStr | None = None base_url: str | None = None model: str | None = None + # AWS credentials - optional, boto3 uses default credential chain if not set + aws_access_key_id: SecretStr | None = None + aws_secret_access_key: SecretStr | None = None + aws_region_name: str | None = None @classmethod def from_env(cls, enabled: bool = False) -> LLMEnvOverrides: @@ -214,20 +272,56 @@ def from_env(cls, enabled: bool = False) -> LLMEnvOverrides: if model: result["model"] = model + # AWS credentials (optional - boto3 uses default credential chain if not set) + aws_access_key_id = os.environ.get(ENV_AWS_ACCESS_KEY_ID) or None + if aws_access_key_id: + result["aws_access_key_id"] = SecretStr(aws_access_key_id) + + aws_secret_access_key = os.environ.get(ENV_AWS_SECRET_ACCESS_KEY) or None + if aws_secret_access_key: + result["aws_secret_access_key"] = SecretStr(aws_secret_access_key) + + aws_region_name = os.environ.get(ENV_AWS_REGION_NAME) or None + if aws_region_name: + result["aws_region_name"] = aws_region_name + return cls(**result) def require_for_headless(self) -> None: + """Validate required environment variables for headless mode. + + For AWS-authenticated models (bedrock/, sagemaker/): + - Only LLM_MODEL is required + - LLM_API_KEY is NOT required (uses AWS IAM credentials) + - AWS credentials can come from env vars or default credential chain + + For standard models: + - Both LLM_MODEL and LLM_API_KEY are required + """ missing: list[str] = [] - if self.api_key is None: - missing.append(ENV_LLM_API_KEY) + + # Model is always required if self.model is None: missing.append(ENV_LLM_MODEL) + + # API key is only required for non-AWS models + uses_aws_auth = is_aws_auth_model(self.model) + if not uses_aws_auth and self.api_key is None: + missing.append(ENV_LLM_API_KEY) + if missing: - raise MissingEnvironmentVariablesError(missing) + raise MissingEnvironmentVariablesError(missing, is_aws_model=uses_aws_auth) def has_overrides(self) -> bool: """Check if any overrides are set.""" - return any([self.api_key, self.base_url, self.model]) + return any([ + self.api_key, + self.base_url, + self.model, + self.aws_access_key_id, + self.aws_secret_access_key, + self.aws_region_name, + ]) def apply_llm_overrides(llm: LLM, overrides: LLMEnvOverrides) -> LLM: @@ -280,15 +374,35 @@ def _ensure_agent(self, agent: Agent | None, overrides: LLMEnvOverrides) -> Agen # In env override mode, require enough info to create an agent. overrides.require_for_headless() - assert overrides.api_key is not None assert overrides.model is not None - llm = LLM( - model=overrides.model, - api_key=overrides.api_key.get_secret_value(), - base_url=overrides.base_url, - usage_id="agent", - ) + # Build LLM kwargs - api_key is optional for AWS-authenticated models + llm_kwargs: dict[str, Any] = { + "model": overrides.model, + "usage_id": "agent", + } + + if overrides.api_key is not None: + llm_kwargs["api_key"] = overrides.api_key.get_secret_value() + + if overrides.base_url is not None: + llm_kwargs["base_url"] = overrides.base_url + + # Add AWS credentials if provided (boto3 uses default chain if not set) + if overrides.aws_access_key_id is not None: + llm_kwargs["aws_access_key_id"] = ( + overrides.aws_access_key_id.get_secret_value() + ) + + if overrides.aws_secret_access_key is not None: + llm_kwargs["aws_secret_access_key"] = ( + overrides.aws_secret_access_key.get_secret_value() + ) + + if overrides.aws_region_name is not None: + llm_kwargs["aws_region_name"] = overrides.aws_region_name + + llm = LLM(**llm_kwargs) return get_default_cli_agent(llm) def _apply_env_overrides(self, agent: Agent, overrides: LLMEnvOverrides) -> Agent: @@ -318,8 +432,10 @@ def load_or_create( * Load it from disk. * Apply any env overrides that are present (even partial). - If no persisted agent exists: - * Require a full env spec (LLM_API_KEY + LLM_MODEL) to create - a default Agent. + * For standard models: Require LLM_API_KEY + LLM_MODEL to create + a default Agent. + * For AWS Bedrock/SageMaker models: Only LLM_MODEL is required; + AWS credentials come from env vars or the default credential chain. * Otherwise, raise an error. Runtime configuration (tools, context, MCP, metadata, critic) is diff --git a/tests/stores/test_env_llm_overrides.py b/tests/stores/test_env_llm_overrides.py index 3bae2ea2e..f8ae7db45 100644 --- a/tests/stores/test_env_llm_overrides.py +++ b/tests/stores/test_env_llm_overrides.py @@ -8,6 +8,9 @@ from openhands.sdk import LLM from openhands_cli.stores.agent_store import ( + ENV_AWS_ACCESS_KEY_ID, + ENV_AWS_REGION_NAME, + ENV_AWS_SECRET_ACCESS_KEY, ENV_LLM_API_KEY, ENV_LLM_BASE_URL, ENV_LLM_MODEL, @@ -15,6 +18,7 @@ MissingEnvironmentVariablesError, apply_llm_overrides, check_and_warn_env_vars, + is_aws_auth_model, ) @@ -704,3 +708,220 @@ def test_single_missing_var(self) -> None: assert ENV_LLM_MODEL in error_str assert "Missing required environment variable(s)" in error_str + + def test_aws_model_error_message(self) -> None: + """Error for AWS model should mention AWS credential chain.""" + error = MissingEnvironmentVariablesError([ENV_LLM_MODEL], is_aws_model=True) + error_str = str(error) + + assert ENV_LLM_MODEL in error_str + assert "AWS credentials" in error_str + assert "credential chain" in error_str + # Should NOT mention LLM_API_KEY as required for AWS models + assert "LLM_API_KEY: Your LLM API key" not in error_str + + +class TestIsAwsAuthModel: + """Tests for is_aws_auth_model helper function.""" + + def test_bedrock_model_detected(self) -> None: + """bedrock/ prefix should be detected as AWS auth model.""" + assert is_aws_auth_model("bedrock/anthropic.claude-3-sonnet") is True + assert is_aws_auth_model("bedrock/meta.llama2-70b") is True + + def test_bedrock_converse_model_detected(self) -> None: + """bedrock_converse/ prefix should be detected as AWS auth model.""" + assert is_aws_auth_model("bedrock_converse/anthropic.claude-3") is True + + def test_sagemaker_model_detected(self) -> None: + """sagemaker/ prefix should be detected as AWS auth model.""" + assert is_aws_auth_model("sagemaker/my-endpoint") is True + + def test_standard_models_not_detected(self) -> None: + """Non-AWS models should not be detected.""" + assert is_aws_auth_model("claude-sonnet-4-5-20250929") is False + assert is_aws_auth_model("gpt-4") is False + assert is_aws_auth_model("anthropic/claude-3") is False + assert is_aws_auth_model("openai/gpt-4") is False + + def test_none_model(self) -> None: + """None model should return False.""" + assert is_aws_auth_model(None) is False + + def test_bedrock_arn_model(self) -> None: + """Bedrock ARN model should be detected.""" + model = ( + "bedrock/arn:aws-us-gov:bedrock:us-gov-west-1:123456:inference-profile/" + "us-gov.anthropic.claude-sonnet-4-5-20250929-v1:0" + ) + assert is_aws_auth_model(model) is True + + +class TestLLMEnvOverridesWithAwsCredentials: + """Tests for LLMEnvOverrides with AWS credentials.""" + + def test_aws_credentials_loaded_from_env(self) -> None: + """AWS credentials should be loaded from environment variables.""" + env_vars = { + ENV_LLM_MODEL: "bedrock/anthropic.claude-3", + ENV_AWS_ACCESS_KEY_ID: "test-access-key", + ENV_AWS_SECRET_ACCESS_KEY: "test-secret-key", + ENV_AWS_REGION_NAME: "us-west-2", + } + with patch.dict(os.environ, env_vars, clear=False): + overrides = LLMEnvOverrides.from_env(enabled=True) + assert overrides.model == "bedrock/anthropic.claude-3" + assert overrides.aws_access_key_id is not None + assert overrides.aws_access_key_id.get_secret_value() == "test-access-key" + assert overrides.aws_secret_access_key is not None + assert ( + overrides.aws_secret_access_key.get_secret_value() == "test-secret-key" + ) + assert overrides.aws_region_name == "us-west-2" + + def test_aws_model_without_api_key_passes_validation(self) -> None: + """AWS model should not require LLM_API_KEY.""" + overrides = LLMEnvOverrides(model="bedrock/anthropic.claude-3") + # Should not raise - AWS models don't require api_key + overrides.require_for_headless() + + def test_aws_model_without_model_fails_validation(self) -> None: + """AWS model still requires LLM_MODEL to be set.""" + overrides = LLMEnvOverrides() # No model set + with pytest.raises(MissingEnvironmentVariablesError) as exc_info: + overrides.require_for_headless() + assert ENV_LLM_MODEL in exc_info.value.missing_vars + + def test_standard_model_without_api_key_fails_validation(self) -> None: + """Standard model should require LLM_API_KEY.""" + overrides = LLMEnvOverrides(model="claude-sonnet-4-5-20250929") + with pytest.raises(MissingEnvironmentVariablesError) as exc_info: + overrides.require_for_headless() + assert ENV_LLM_API_KEY in exc_info.value.missing_vars + + def test_has_overrides_with_aws_credentials(self) -> None: + """has_overrides should return True when AWS credentials are set.""" + overrides = LLMEnvOverrides( + aws_access_key_id=SecretStr("key"), + ) + assert overrides.has_overrides() is True + + overrides2 = LLMEnvOverrides( + aws_region_name="us-west-2", + ) + assert overrides2.has_overrides() is True + + +class TestAgentStoreWithAwsModels: + """Tests for AgentStore with AWS Bedrock/SageMaker models.""" + + def test_agent_created_for_bedrock_model_without_api_key(self, tmp_path) -> None: + """Agent should be created for Bedrock model without LLM_API_KEY.""" + from openhands_cli.stores import AgentStore + + conversations_dir = tmp_path / "conversations" + conversations_dir.mkdir(exist_ok=True) + + # Set bedrock model but not API key + env_vars = { + ENV_LLM_MODEL: "bedrock/anthropic.claude-3-sonnet", + } + + with ( + patch( + "openhands_cli.stores.agent_store.get_persistence_dir", + return_value=str(tmp_path), + ), + patch( + "openhands_cli.stores.agent_store.get_conversations_dir", + return_value=str(conversations_dir), + ), + patch.dict(os.environ, env_vars, clear=False), + ): + # Ensure LLM_API_KEY is not set + os.environ.pop(ENV_LLM_API_KEY, None) + + store = AgentStore() + agent = store.load_or_create(env_overrides_enabled=True) + + assert agent is not None + assert agent.llm.model == "bedrock/anthropic.claude-3-sonnet" + # api_key should be None for bedrock model + assert agent.llm.api_key is None + + def test_agent_created_with_aws_credentials(self, tmp_path) -> None: + """Agent should be created with explicit AWS credentials.""" + from openhands_cli.stores import AgentStore + + conversations_dir = tmp_path / "conversations" + conversations_dir.mkdir(exist_ok=True) + + env_vars = { + ENV_LLM_MODEL: "bedrock/anthropic.claude-3-sonnet", + ENV_AWS_ACCESS_KEY_ID: "test-access-key", + ENV_AWS_SECRET_ACCESS_KEY: "test-secret-key", + ENV_AWS_REGION_NAME: "us-west-2", + } + + with ( + patch( + "openhands_cli.stores.agent_store.get_persistence_dir", + return_value=str(tmp_path), + ), + patch( + "openhands_cli.stores.agent_store.get_conversations_dir", + return_value=str(conversations_dir), + ), + patch.dict(os.environ, env_vars, clear=False), + ): + os.environ.pop(ENV_LLM_API_KEY, None) + + store = AgentStore() + agent = store.load_or_create(env_overrides_enabled=True) + + assert agent is not None + assert agent.llm.model == "bedrock/anthropic.claude-3-sonnet" + # AWS credentials are stored as SecretStr in the SDK's LLM class + assert agent.llm.aws_access_key_id is not None + assert isinstance(agent.llm.aws_access_key_id, SecretStr) + assert agent.llm.aws_access_key_id.get_secret_value() == "test-access-key" + assert agent.llm.aws_secret_access_key is not None + assert isinstance(agent.llm.aws_secret_access_key, SecretStr) + assert ( + agent.llm.aws_secret_access_key.get_secret_value() == "test-secret-key" + ) + assert agent.llm.aws_region_name == "us-west-2" + + def test_bedrock_arn_model_works(self, tmp_path) -> None: + """Bedrock model with ARN should work without API key.""" + from openhands_cli.stores import AgentStore + + conversations_dir = tmp_path / "conversations" + conversations_dir.mkdir(exist_ok=True) + + model = ( + "bedrock/arn:aws-us-gov:bedrock:us-gov-west-1:123456:inference-profile/" + "us-gov.anthropic.claude-sonnet-4-5-20250929-v1:0" + ) + env_vars = { + ENV_LLM_MODEL: model, + } + + with ( + patch( + "openhands_cli.stores.agent_store.get_persistence_dir", + return_value=str(tmp_path), + ), + patch( + "openhands_cli.stores.agent_store.get_conversations_dir", + return_value=str(conversations_dir), + ), + patch.dict(os.environ, env_vars, clear=False), + ): + os.environ.pop(ENV_LLM_API_KEY, None) + + store = AgentStore() + agent = store.load_or_create(env_overrides_enabled=True) + + assert agent is not None + assert agent.llm.model == model