feat: add AWS_USE_IMDS support for ambient IAM credential resolution on Bedrock

docs/docs/installation/github.md

@@ -280,6 +280,35 @@ To use local models via Ollama:
 
 **Note:** For local models, you'll need to use a self-hosted runner with Ollama installed, as GitHub Actions hosted runners cannot access localhost services.
 
+##### Using Amazon Bedrock
+
+To use Amazon Bedrock models with static IAM credentials:
+
+```yaml
+      env:
+        GITHUB_TOKEN: ${{ secrets.GITHUB_TOKEN }}
+        config.model: "bedrock/anthropic.claude-3-5-sonnet-20240620-v1:0"
+        config.fallback_models: '["bedrock/anthropic.claude-3-5-sonnet-20240620-v1:0"]'
+        aws.AWS_ACCESS_KEY_ID: ${{ secrets.AWS_ACCESS_KEY_ID }}
+        aws.AWS_SECRET_ACCESS_KEY: ${{ secrets.AWS_SECRET_ACCESS_KEY }}
+        aws.AWS_REGION_NAME: "us-east-1"
+```
+
+**Recommended: IAM Role Credentials on AWS Compute**
+
+When the GitHub Actions runner is on AWS infrastructure (EC2, ECS, EKS), use the instance/task IAM role directly — no secrets required:
+
+```yaml
+      env:
+        GITHUB_TOKEN: ${{ secrets.GITHUB_TOKEN }}
+        config.model: "bedrock/anthropic.claude-3-5-sonnet-20240620-v1:0"
+        config.fallback_models: '["bedrock/anthropic.claude-3-5-sonnet-20240620-v1:0"]'
+        AWS_USE_IMDS: "true"
+        # AWS_REGION_NAME: us-east-1  # optional if instance metadata provides the region
+```
+
+The IAM role must have `bedrock:InvokeModel` on the target model ARN. See [Bedrock model configuration](../usage-guide/changing_a_model.md#amazon-bedrock) for the full IAM policy example and supported models.
+
 #### Advanced Configuration Options
 
 ##### Custom Review Instructions
@@ -732,4 +761,4 @@ After you set up AWS CodeCommit using the instructions above, here is an example
 PYTHONPATH="/PATH/TO/PROJECTS/pr-agent" python pr_agent/cli.py \
   --pr_url https://us-east-1.console.aws.amazon.com/codesuite/codecommit/repositories/MY_REPO_NAME/pull-requests/321 \
   review
-```
+```

docs/docs/usage-guide/changing_a_model.md

@@ -111,7 +111,7 @@ Please note that the `custom_model_max_tokens` setting should be configured in a
     Commercial models such as GPT-5, Claude Sonnet, and Gemini have demonstrated robust capabilities in generating structured output for code analysis tasks with large input. In contrast, most open-source models currently available (as of January 2025) face challenges with these complex tasks.
 
     Based on our testing, local open-source models are suitable for experimentation and learning purposes (mainly for the `ask` command), but they are not suitable for production-level code analysis tasks.
-    
+
     Hence, for production workflows and real-world usage, we recommend using commercial models.
 
 ### Hugging Face
@@ -251,6 +251,43 @@ model="bedrock/us.meta.llama4-scout-17b-instruct-v1:0"
 fallback_models=["bedrock/us.meta.llama4-maverick-17b-instruct-v1:0"]
 ```
 
+#### Using IAM Role Credentials (Recommended on AWS Compute)
+
+When running PR-Agent on AWS infrastructure (EC2, ECS/Fargate, EKS with IRSA, Lambda, or any self-hosted GitHub Actions runner on AWS), the instance or task already has an IAM role attached. You can use those ambient credentials directly instead of storing long-lived static keys.
+
+Set `AWS_USE_IMDS=true` in the environment. PR-Agent will resolve credentials via boto3's standard provider chain, which handles all AWS compute contexts transparently:
+
+| Compute context | Mechanism |
+|---|---|
+| EC2 instance with IAM role | IMDSv2 (169.254.169.254) |
+| ECS / Fargate task role | Task metadata endpoint |
+| EKS pod with IRSA | Web identity token + STS |
+| Lambda function | Runtime-injected credentials |
+
+Minimal GitHub Actions workflow (no AWS secret keys required):
+
+```yaml
+- uses: Codium-ai/pr-agent@main
+  env:
+    GITHUB_TOKEN: ${{ secrets.GITHUB_TOKEN }}
+    AWS_USE_IMDS: "true"
+    # AWS_REGION_NAME: us-east-1  # optional if the instance metadata provides it
+  with:
+    command: review
+```
+
+The IAM role must have `bedrock:InvokeModel` permission on the target model ARN, for example:
+
+```json
+{
+  "Effect": "Allow",
+  "Action": "bedrock:InvokeModel",
+  "Resource": "arn:aws:bedrock:us-east-1::foundation-model/anthropic.claude-3-5-sonnet-20240620-v1:0"
+}
+```
+
+If you also configure static keys in `[aws]`, they serve as an automatic fallback: if the ambient credentials fail a Bedrock call (e.g., the role lacks `bedrock:InvokeModel`), PR-Agent retries with the static keys and logs a warning.
+
 #### Custom Inference Profiles
 
 To use a custom inference profile with Amazon Bedrock (for cost allocation tags and other configuration settings), add the `model_id` parameter to your configuration:
@@ -339,7 +376,7 @@ key = "..." # your Codestral api key
 To use model from Openrouter, for example, set:
 
 ```toml
-[config] # in configuration.toml 
+[config] # in configuration.toml
 model="openrouter/anthropic/claude-3.7-sonnet"
 fallback_models=["openrouter/deepseek/deepseek-chat"]
 custom_model_max_tokens=20000

pr_agent/algo/__init__.py

@@ -159,12 +159,22 @@
     'bedrock/anthropic.claude-sonnet-4-20250514-v1:0': 200000,
     'bedrock/anthropic.claude-sonnet-4-5-20250929-v1:0': 200000,
     'bedrock/anthropic.claude-sonnet-4-6': 200000,
+    'bedrock/anthropic.claude-sonnet-4-6-v1:0': 200000,
+    'bedrock/anthropic.claude-opus-4-5-20251101-v1:0': 200000,
     "bedrock/us.anthropic.claude-opus-4-20250514-v1:0": 200000,
     "bedrock/us.anthropic.claude-opus-4-1-20250805-v1:0": 200000,
     "bedrock/us.anthropic.claude-opus-4-6-20260120-v1:0": 200000,
     "bedrock/global.anthropic.claude-opus-4-5-20251101-v1:0": 200000,
+    "bedrock/eu.anthropic.claude-opus-4-5-20251101-v1:0": 200000,
+    "bedrock/au.anthropic.claude-opus-4-5-20251101-v1:0": 200000,
+    "bedrock/jp.anthropic.claude-opus-4-5-20251101-v1:0": 200000,
+    "bedrock/apac.anthropic.claude-opus-4-5-20251101-v1:0": 200000,
     "bedrock/us.anthropic.claude-opus-4-5-20251101-v1:0": 200000,
     "bedrock/global.anthropic.claude-opus-4-6-v1:0": 200000,
+    "bedrock/eu.anthropic.claude-opus-4-6-v1:0": 200000,
+    "bedrock/au.anthropic.claude-opus-4-6-v1:0": 200000,
+    "bedrock/jp.anthropic.claude-opus-4-6-v1:0": 200000,
+    "bedrock/apac.anthropic.claude-opus-4-6-v1:0": 200000,
     "bedrock/us.anthropic.claude-opus-4-6-v1:0": 200000,
     "bedrock/us.anthropic.claude-3-5-sonnet-20241022-v2:0": 100000,
     "bedrock/us.anthropic.claude-haiku-4-5-20251001-v1:0": 200000,
@@ -179,16 +189,23 @@
     "bedrock/us.anthropic.claude-sonnet-4-5-20250929-v1:0": 200000,
     "bedrock/au.anthropic.claude-sonnet-4-5-20250929-v1:0": 200000,
     "bedrock/us.anthropic.claude-sonnet-4-6": 200000,
+    "bedrock/us.anthropic.claude-sonnet-4-6-v1:0": 200000,
     "bedrock/au.anthropic.claude-sonnet-4-6": 200000,
+    "bedrock/au.anthropic.claude-sonnet-4-6-v1:0": 200000,
     "bedrock/apac.anthropic.claude-3-5-sonnet-20241022-v2:0": 100000,
     "bedrock/apac.anthropic.claude-3-7-sonnet-20250219-v1:0": 200000,
     "bedrock/apac.anthropic.claude-sonnet-4-20250514-v1:0": 200000,
     "bedrock/eu.anthropic.claude-sonnet-4-5-20250929-v1:0": 200000,
     "bedrock/eu.anthropic.claude-sonnet-4-6": 200000,
+    "bedrock/eu.anthropic.claude-sonnet-4-6-v1:0": 200000,
     "bedrock/jp.anthropic.claude-sonnet-4-5-20250929-v1:0": 200000,
     "bedrock/jp.anthropic.claude-sonnet-4-6": 200000,
+    "bedrock/jp.anthropic.claude-sonnet-4-6-v1:0": 200000,
+    "bedrock/apac.anthropic.claude-sonnet-4-6": 200000,
+    "bedrock/apac.anthropic.claude-sonnet-4-6-v1:0": 200000,
     "bedrock/global.anthropic.claude-sonnet-4-5-20250929-v1:0": 200000,
     "bedrock/global.anthropic.claude-sonnet-4-6": 200000,
+    "bedrock/global.anthropic.claude-sonnet-4-6-v1:0": 200000,
     'claude-3-5-sonnet': 100000,
     'bedrock/us.meta.llama4-scout-17b-instruct-v1:0': 128000,
     'bedrock/us.meta.llama4-maverick-17b-instruct-v1:0': 128000,

pr_agent/algo/ai_handlers/litellm_ai_handler.py

@@ -1,18 +1,25 @@
+import json
 import os
+
 import litellm
 import openai
 import requests
 from litellm import acompletion
-from tenacity import retry, retry_if_exception_type, retry_if_not_exception_type, stop_after_attempt
-
-from pr_agent.algo import CLAUDE_EXTENDED_THINKING_MODELS, NO_SUPPORT_TEMPERATURE_MODELS, SUPPORT_REASONING_EFFORT_MODELS, USER_MESSAGE_ONLY_MODELS, STREAMING_REQUIRED_MODELS
+from tenacity import (retry, retry_if_exception_type,
+                      retry_if_not_exception_type, stop_after_attempt)
+
+from pr_agent.algo import (CLAUDE_EXTENDED_THINKING_MODELS,
+                           NO_SUPPORT_TEMPERATURE_MODELS,
+                           STREAMING_REQUIRED_MODELS,
+                           SUPPORT_REASONING_EFFORT_MODELS,
+                           USER_MESSAGE_ONLY_MODELS)
 from pr_agent.algo.ai_handlers.base_ai_handler import BaseAiHandler
-from pr_agent.algo.ai_handlers.litellm_helpers import _handle_streaming_response, MockResponse, _get_azure_ad_token, \
-    _process_litellm_extra_body
+from pr_agent.algo.ai_handlers.litellm_helpers import (
+    MockResponse, _get_azure_ad_token, _handle_streaming_response,
+    _process_litellm_extra_body)
 from pr_agent.algo.utils import ReasoningEffort, get_version
 from pr_agent.config_loader import get_settings
 from pr_agent.log import get_logger
-import json
 
 MODEL_RETRIES = 2
 DUMMY_LITELLM_API_KEY = "dummy_key"  # placeholder set when no OpenAI key is configured
@@ -33,6 +40,9 @@ def __init__(self):
         self.azure = False
         self.api_base = None
         self.repetition_penalty = None
+        self._aws_imds_mode = False
+        self._aws_static_creds = None
+        self._aws_imds_fell_back = False
 
         if get_settings().get("LITELLM.DISABLE_AIOHTTP", False):
             litellm.disable_aiohttp_transport = True
@@ -41,7 +51,46 @@ def __init__(self):
             litellm.openai_key = get_settings().openai.key
         elif 'OPENAI_API_KEY' not in os.environ:
             litellm.api_key = DUMMY_LITELLM_API_KEY
-        if get_settings().get("aws.AWS_ACCESS_KEY_ID"):
+        if os.environ.get("AWS_USE_IMDS", "").strip().lower() in ("1", "true", "yes"):
+            import boto3
+            session = boto3.Session()
+            try:
+                creds = session.get_credentials()
+                if creds:
+                    frozen = creds.get_frozen_credentials()
+                    os.environ["AWS_ACCESS_KEY_ID"] = frozen.access_key
+                    os.environ["AWS_SECRET_ACCESS_KEY"] = frozen.secret_key
+                    if frozen.token:
+                        os.environ["AWS_SESSION_TOKEN"] = frozen.token
+                    elif "AWS_SESSION_TOKEN" in os.environ:
+                        del os.environ["AWS_SESSION_TOKEN"]
+                    self._aws_imds_mode = True
+                    get_logger().info("Using ambient AWS credentials from IMDS/task-role/IRSA")
+                else:
+                    get_logger().warning("AWS_USE_IMDS is set but boto3 found no credentials; falling through to static keys")
+            except Exception as e:
+                get_logger().error(f"AWS_USE_IMDS: failed to resolve credentials via boto3: {e}; falling through to static keys")
+            if not os.environ.get("AWS_REGION_NAME") and not get_settings().get("aws.AWS_REGION_NAME"):
+                try:
+                    region = session.region_name
+                    if region:
+                        os.environ["AWS_REGION_NAME"] = region
+                        get_logger().info(f"AWS region resolved from environment: {region}")
+                    else:
+                        get_logger().warning("AWS_USE_IMDS: could not determine AWS region; set AWS_REGION_NAME explicitly")
+                except Exception as e:
+                    get_logger().warning(f"AWS_USE_IMDS: failed to resolve region via boto3: {e}")
+            if get_settings().get("aws.AWS_ACCESS_KEY_ID"):
+                if get_settings().aws.AWS_SECRET_ACCESS_KEY and get_settings().aws.AWS_REGION_NAME:
+                    self._aws_static_creds = {
+                        "AWS_ACCESS_KEY_ID": get_settings().aws.AWS_ACCESS_KEY_ID,
+                        "AWS_SECRET_ACCESS_KEY": get_settings().aws.AWS_SECRET_ACCESS_KEY,
+                        "AWS_REGION_NAME": get_settings().aws.AWS_REGION_NAME,
+                    }
+                    static_token = get_settings().get("aws.AWS_SESSION_TOKEN", None)
+                    if static_token:
+                        self._aws_static_creds["AWS_SESSION_TOKEN"] = static_token
+        elif get_settings().get("aws.AWS_ACCESS_KEY_ID"):
             assert get_settings().aws.AWS_SECRET_ACCESS_KEY and get_settings().aws.AWS_REGION_NAME, "AWS credentials are incomplete"
             os.environ["AWS_ACCESS_KEY_ID"] = get_settings().aws.AWS_ACCESS_KEY_ID
             os.environ["AWS_SECRET_ACCESS_KEY"] = get_settings().aws.AWS_SECRET_ACCESS_KEY
@@ -108,7 +157,7 @@ def __init__(self):
         # Support mistral models
         if get_settings().get("MISTRAL.KEY", None):
             os.environ["MISTRAL_API_KEY"] = get_settings().get("MISTRAL.KEY")
-        
+
         # Support codestral models
         if get_settings().get("CODESTRAL.KEY", None):
             os.environ["CODESTRAL_API_KEY"] = get_settings().get("CODESTRAL.KEY")
@@ -120,7 +169,7 @@ def __init__(self):
             access_token = _get_azure_ad_token()
             litellm.api_key = access_token
             openai.api_key = access_token
-            
+
             # Set API base from settings
             self.api_base = get_settings().azure_ad.api_base
             litellm.api_base = self.api_base
@@ -153,6 +202,37 @@ def __init__(self):
         # Models that require streaming
         self.streaming_required_models = STREAMING_REQUIRED_MODELS
 
+    def _refresh_imds_credentials(self):
+        """Refresh ambient AWS credentials from boto3 provider chain. Called before each Bedrock call
+        to avoid serving stale credentials from long-lived processes (EC2 roles rotate every ~6h)."""
+        try:
+            import boto3
+            creds = boto3.Session().get_credentials()
+            if creds:
+                frozen = creds.get_frozen_credentials()
+                os.environ["AWS_ACCESS_KEY_ID"] = frozen.access_key
+                os.environ["AWS_SECRET_ACCESS_KEY"] = frozen.secret_key
+                if frozen.token:
+                    os.environ["AWS_SESSION_TOKEN"] = frozen.token
+                elif "AWS_SESSION_TOKEN" in os.environ:
+                    del os.environ["AWS_SESSION_TOKEN"]
+            else:
+                get_logger().warning("IMDS credential refresh: boto3 returned no credentials")
+        except Exception as e:
+            get_logger().warning(f"IMDS credential refresh failed: {e}")
+
+    def _activate_static_aws_fallback(self):
+        """Swap process env to static credentials for Bedrock fallback after IMDS failure."""
+        os.environ["AWS_ACCESS_KEY_ID"] = self._aws_static_creds["AWS_ACCESS_KEY_ID"]
+        os.environ["AWS_SECRET_ACCESS_KEY"] = self._aws_static_creds["AWS_SECRET_ACCESS_KEY"]
+        os.environ["AWS_REGION_NAME"] = self._aws_static_creds["AWS_REGION_NAME"]
+        if "AWS_SESSION_TOKEN" in self._aws_static_creds:
+            os.environ["AWS_SESSION_TOKEN"] = self._aws_static_creds["AWS_SESSION_TOKEN"]
+        elif "AWS_SESSION_TOKEN" in os.environ:
+            del os.environ["AWS_SESSION_TOKEN"]
+        self._aws_imds_fell_back = True
+        get_logger().warning("Bedrock call failed with ambient (IMDS) credentials; retrying with static credentials")
+
     def prepare_logs(self, response, system, user, resp, finish_reason):
         response_log = response.dict().copy()
         response_log['system'] = system
@@ -268,6 +348,8 @@ def deployment_id(self):
         stop=stop_after_attempt(MODEL_RETRIES),
     )
     async def chat_completion(self, model: str, system: str, user: str, temperature: float = 0.2, img_path: str = None):
+        if self._aws_imds_mode and not self._aws_imds_fell_back and 'bedrock/' in model:
+            self._refresh_imds_credentials()
         try:
             resp, finish_reason = None, None
             deployment_id = self.deployment_id
@@ -419,6 +501,11 @@ async def chat_completion(self, model: str, system: str, user: str, temperature:
             get_logger().error(f"Rate limit error during LLM inference: {e}")
             raise
         except openai.APIError as e:
+            if (self._aws_imds_mode
+                    and not self._aws_imds_fell_back
+                    and self._aws_static_creds
+                    and 'bedrock/' in model):
+                self._activate_static_aws_fallback()
             get_logger().warning(f"Error during LLM inference: {e}")
             raise
         except Exception as e:

pr_agent/settings/.secrets_template.toml

@@ -130,9 +130,14 @@ key = ""
 api_base = ""
 
 [aws]
+# When running on AWS compute (EC2, ECS, EKS, Lambda) with an IAM role attached,
+# set AWS_USE_IMDS=true in the environment instead of providing static keys here.
+# These keys are only needed as an optional fallback when AWS_USE_IMDS=true, or
+# when running outside AWS compute without ambient credentials.
 AWS_ACCESS_KEY_ID = ""
 AWS_SECRET_ACCESS_KEY = ""
 AWS_REGION_NAME = ""
+AWS_SESSION_TOKEN = ""  # optional: only needed for STS-derived or temporary credentials
 
 [aws_secrets_manager]
 secret_arn = ""         # The ARN of the AWS Secrets Manager secret containing PR-Agent configuration