diff --git a/.env.example b/.env.example index 18e2b51e..9372fcf2 100644 --- a/.env.example +++ b/.env.example @@ -54,9 +54,18 @@ PROXY_API_KEY="my-super-secret-password-123" # OPTIONAL # =========================================== -# AWS region (default: us-east-1) +# AWS SSO/auth region (default: us-east-1) +# This controls the OIDC token refresh endpoint: https://oidc.{region}.amazonaws.com/token +# Set this to your SSO provider's region (e.g., us-east-2 for organizations using that region). # KIRO_REGION="us-east-1" +# AWS Q Developer API region (default: us-east-1) +# This controls the Q API endpoint: https://q.{region}.amazonaws.com +# Note: q.amazonaws.com only exists in specific regions. If your SSO region (KIRO_REGION) +# is different from the Q API region (e.g., your SSO is us-east-2 but Q API is us-east-1), +# set KIRO_API_REGION to the correct Q API region. +# KIRO_API_REGION="us-east-1" + # =========================================== # SERVER SETTINGS # =========================================== diff --git a/kiro/__init__.py b/kiro/__init__.py index 42c86b1f..058a786e 100644 --- a/kiro/__init__.py +++ b/kiro/__init__.py @@ -54,6 +54,7 @@ from kiro.config import ( PROXY_API_KEY, REGION, + API_REGION, HIDDEN_MODELS, APP_VERSION, ) @@ -105,6 +106,7 @@ # Configuration "PROXY_API_KEY", "REGION", + "API_REGION", "HIDDEN_MODELS", "APP_VERSION", diff --git a/kiro/auth.py b/kiro/auth.py index 9ec39cf8..74eb9fc3 100644 --- a/kiro/auth.py +++ b/kiro/auth.py @@ -118,6 +118,7 @@ def __init__( refresh_token: Optional[str] = None, profile_arn: Optional[str] = None, region: str = "us-east-1", + api_region: str = "us-east-1", creds_file: Optional[str] = None, client_id: Optional[str] = None, client_secret: Optional[str] = None, @@ -129,7 +130,11 @@ def __init__( Args: refresh_token: Refresh token for obtaining access token profile_arn: AWS CodeWhisperer profile ARN - region: AWS region (default: us-east-1) + region: AWS SSO/auth region used for OIDC token refresh endpoint (default: us-east-1) + api_region: AWS Q Developer API region for endpoint URL construction (default: us-east-1). + Separate from ``region`` because q.amazonaws.com only exists in certain + regions. Set via ``KIRO_API_REGION`` env var when your SSO region + (``KIRO_REGION``) differs from the Q API region. creds_file: Path to JSON file with credentials (optional) client_id: OAuth client ID (for AWS SSO OIDC, optional) client_secret: OAuth client secret (for AWS SSO OIDC, optional) @@ -161,13 +166,15 @@ def __init__( # Auth type will be determined after loading credentials self._auth_type: AuthType = AuthType.KIRO_DESKTOP - # Dynamic URLs based on region + # Dynamic URLs based on regions + # SSO/auth region: used for OIDC token refresh endpoint self._refresh_url = get_kiro_refresh_url(region) - self._api_host = get_kiro_api_host(region) - self._q_host = get_kiro_q_host(region) + # API region: used for Q Developer API endpoint (separate from SSO region) + self._api_host = get_kiro_api_host(api_region) + self._q_host = get_kiro_q_host(api_region) # Log initialized endpoints for diagnostics (helps with DNS issues like #58) - logger.info(f"Auth manager initialized: region={region}, api_host={self._api_host}, q_host={self._q_host}") + logger.info(f"Auth manager initialized: sso_region={region}, api_region={api_region}, api_host={self._api_host}, q_host={self._q_host}") # Fingerprint for User-Agent self._fingerprint = get_machine_fingerprint() diff --git a/kiro/config.py b/kiro/config.py index 9f1f6ce0..3434a2be 100644 --- a/kiro/config.py +++ b/kiro/config.py @@ -130,9 +130,18 @@ def _get_raw_env_value(var_name: str, env_file: str = ".env") -> Optional[str]: # Profile ARN for AWS CodeWhisperer PROFILE_ARN: str = os.getenv("PROFILE_ARN", "") -# AWS region (default us-east-1) +# AWS SSO/auth region (default us-east-1) +# Used for OIDC token refresh endpoint (e.g., https://oidc.{region}.amazonaws.com/token) +# May differ from API region — set this to match your SSO provider's region REGION: str = os.getenv("KIRO_REGION", "us-east-1") +# AWS Q API region (default us-east-1) +# Used for the Q Developer API endpoint (e.g., https://q.{region}.amazonaws.com) +# Note: q.amazonaws.com endpoints only exist in specific regions. If your SSO region +# (KIRO_REGION) differs from your Q API region, set KIRO_API_REGION explicitly. +# See supported regions: https://docs.aws.amazon.com/amazonq/latest/qdeveloper-ug/regions.html +API_REGION: str = os.getenv("KIRO_API_REGION", "us-east-1") + # Path to credentials file (optional, alternative to .env) # Read directly from .env to avoid escape sequence issues on Windows # (e.g., \a in path D:\Projects\adolf is interpreted as bell character) diff --git a/main.py b/main.py index 9f9cb410..1dbb0c1a 100644 --- a/main.py +++ b/main.py @@ -59,6 +59,7 @@ REFRESH_TOKEN, PROFILE_ARN, REGION, + API_REGION, KIRO_CREDS_FILE, KIRO_CLI_DB_FILE, PROXY_API_KEY, @@ -341,6 +342,7 @@ async def lifespan(app: FastAPI): refresh_token=REFRESH_TOKEN, profile_arn=PROFILE_ARN, region=REGION, + api_region=API_REGION, creds_file=KIRO_CREDS_FILE if KIRO_CREDS_FILE else None, sqlite_db=KIRO_CLI_DB_FILE if KIRO_CLI_DB_FILE else None, ) diff --git a/tests/unit/test_auth_manager.py b/tests/unit/test_auth_manager.py index 5ba36dec..9cfd47ab 100644 --- a/tests/unit/test_auth_manager.py +++ b/tests/unit/test_auth_manager.py @@ -47,24 +47,27 @@ def test_initialization_stores_credentials(self): def test_initialization_sets_correct_urls_for_region(self): """ - What it does: Verifies URL formation based on region. - Purpose: Ensure URLs are dynamically formed with the correct region. + What it does: Verifies URL formation based on region and api_region. + Purpose: Ensure SSO refresh URL uses region, and API hosts use api_region. """ - print("Setup: Creating KiroAuthManager with region eu-west-1...") + print("Setup: Creating KiroAuthManager with region=eu-west-1, api_region=eu-central-1...") manager = KiroAuthManager( refresh_token="test_token", - region="eu-west-1" + region="eu-west-1", + api_region="eu-central-1", ) - - print("Verification: URLs contain correct region...") + + print("Verification: refresh_url uses SSO region (eu-west-1)...") print(f"Comparing refresh_url: Expected 'eu-west-1' in URL, Got '{manager._refresh_url}'") assert "eu-west-1" in manager._refresh_url - - print(f"Comparing api_host: Expected 'eu-west-1' in URL, Got '{manager._api_host}'") - assert "eu-west-1" in manager._api_host - - print(f"Comparing q_host: Expected 'eu-west-1' in URL, Got '{manager._q_host}'") - assert "eu-west-1" in manager._q_host + + print("Verification: api_host uses api_region (eu-central-1)...") + print(f"Comparing api_host: Expected 'eu-central-1' in URL, Got '{manager._api_host}'") + assert "eu-central-1" in manager._api_host + + print("Verification: q_host uses api_region (eu-central-1)...") + print(f"Comparing q_host: Expected 'eu-central-1' in URL, Got '{manager._q_host}'") + assert "eu-central-1" in manager._q_host def test_initialization_generates_fingerprint(self): """ @@ -1686,6 +1689,89 @@ async def test_refresh_token_aws_sso_oidc_no_retry_without_sqlite_db( assert mock_client.post.call_count == 1 +# ============================================================================= +# Tests for KIRO_API_REGION separation from KIRO_REGION (Issue: q.{region}.amazonaws.com +# only exists in specific regions, while SSO region can be anything) +# ============================================================================= + +class TestKiroAuthManagerApiRegionSeparation: + """Tests for KIRO_API_REGION separation from KIRO_REGION. + + Background: q.amazonaws.com endpoints only exist in specific regions (e.g., us-east-1). + Users with SSO in regions like us-east-2 would get DNS failures if we used their + SSO region for the Q API endpoint. KIRO_API_REGION lets them configure these separately. + """ + + def test_api_host_uses_api_region_not_sso_region(self): + """ + What it does: Verifies api_host uses api_region, not the SSO region. + Purpose: Ensure Q API calls go to a valid endpoint even when SSO is in a + region that has no q.amazonaws.com endpoint (e.g., us-east-2). + """ + print("Setup: Creating KiroAuthManager with sso_region=us-east-2, api_region=us-east-1...") + manager = KiroAuthManager( + refresh_token="test_token", + region="us-east-2", # SSO region (e.g., org SSO is in us-east-2) + api_region="us-east-1", # Q API region (q.us-east-1.amazonaws.com exists) + ) + + print("Verification: api_host uses api_region (us-east-1)...") + print(f"api_host: {manager._api_host}") + assert "us-east-1" in manager._api_host + assert "us-east-2" not in manager._api_host + + print("Verification: q_host uses api_region (us-east-1)...") + print(f"q_host: {manager._q_host}") + assert "us-east-1" in manager._q_host + assert "us-east-2" not in manager._q_host + + def test_refresh_url_uses_sso_region_not_api_region(self): + """ + What it does: Verifies refresh_url uses the SSO region, not api_region. + Purpose: Ensure Kiro Desktop token refresh goes to the correct regional endpoint. + """ + print("Setup: Creating KiroAuthManager with sso_region=us-east-2, api_region=us-east-1...") + manager = KiroAuthManager( + refresh_token="test_token", + region="us-east-2", + api_region="us-east-1", + ) + + print("Verification: _refresh_url uses SSO region (us-east-2)...") + print(f"_refresh_url: {manager._refresh_url}") + assert "us-east-2" in manager._refresh_url + assert "us-east-1" not in manager._refresh_url + + def test_api_region_defaults_to_us_east_1(self): + """ + What it does: Verifies api_region defaults to us-east-1 when not specified. + Purpose: Ensure backward compatibility — existing users need no config change. + """ + print("Setup: Creating KiroAuthManager without api_region param...") + manager = KiroAuthManager(refresh_token="test_token") + + print("Verification: api_host defaults to us-east-1...") + assert "us-east-1" in manager._api_host + assert "us-east-1" in manager._q_host + + def test_api_region_can_be_set_independently(self): + """ + What it does: Verifies api_region can be any valid Q API region. + Purpose: Support users whose Q API region differs from us-east-1. + """ + print("Setup: Creating KiroAuthManager with api_region=eu-central-1...") + manager = KiroAuthManager( + refresh_token="test_token", + region="us-east-1", + api_region="eu-central-1", + ) + + print("Verification: api_host uses eu-central-1...") + assert "eu-central-1" in manager._api_host + assert "eu-central-1" in manager._q_host + assert "us-east-1" not in manager._api_host + + # ============================================================================= # Tests for is_token_expired() method # =============================================================================