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..4972afd9 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() @@ -310,7 +317,11 @@ def _load_credentials_from_file(self, file_path: str) -> None: - refreshToken: Refresh token - accessToken: Access token (if already available) - profileArn: Profile ARN - - region: AWS region + - region: Auth/SSO region — stored as _sso_region and used only for the auth refresh + URL (e.g., https://prod.{region}.auth.desktop.kiro.dev/refreshToken). The Kiro API + endpoints (_api_host, _q_host) are controlled exclusively by the KIRO_REGION env var + (default: us-east-1) and are NOT updated from this field. This prevents invalid + organization login regions (e.g., us-east-2) from breaking API calls. - expiresAt: Token expiration time (ISO 8601) Additional fields for AWS SSO OIDC (kiro-cli): @@ -342,12 +353,14 @@ def _load_credentials_from_file(self, file_path: str) -> None: if 'profileArn' in data: self._profile_arn = data['profileArn'] if 'region' in data: - self._region = data['region'] - # Update URLs for new region - self._refresh_url = get_kiro_refresh_url(self._region) - self._api_host = get_kiro_api_host(self._region) - self._q_host = get_kiro_q_host(self._region) - logger.info(f"Region updated from credentials file: region={self._region}, api_host={self._api_host}, q_host={self._q_host}") + # Store as SSO/auth region only — does NOT override _api_host/_q_host. + # See docstring for details on why API host is kept separate. + self._sso_region = data['region'] + self._refresh_url = get_kiro_refresh_url(self._sso_region) + logger.info( + f"Auth region from credentials file: sso_region={self._sso_region} " + f"(API stays at region={self._region}, api_host={self._api_host})" + ) # Load clientIdHash and device registration for Enterprise Kiro IDE if 'clientIdHash' in data: 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..c0c28fba 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): """ @@ -122,6 +125,91 @@ def test_load_credentials_file_not_found(self, tmp_path): print(f"Comparing refresh_token: Expected 'fallback_token', Got '{manager._refresh_token}'") assert manager._refresh_token == "fallback_token" + def test_api_host_not_overridden_by_credentials_file_region(self, tmp_path): + """ + What it does: Verifies API host stays at configured region when credentials contain a different region. + Purpose: Prevent invalid regions (e.g., us-east-2 from org login) from breaking API calls. + + Regression test for: credentials file with us-east-2 (org login region) causing gateway + to send API requests to https://q.us-east-2.amazonaws.com which does not exist. + """ + print("Setup: Creating credentials file with us-east-2 region (org login scenario)...") + creds_file = tmp_path / "kiro-auth-token.json" + creds_data = { + "accessToken": "test_access_token", + "refreshToken": "test_refresh_token", + "expiresAt": "2099-01-01T00:00:00.000Z", + "region": "us-east-2" + } + creds_file.write_text(json.dumps(creds_data)) + + manager = KiroAuthManager(creds_file=str(creds_file), region="us-east-1") + + print("Verification: API region stays at us-east-1...") + print(f"Comparing _region: Expected 'us-east-1', Got '{manager._region}'") + assert manager._region == "us-east-1" + + print("Verification: api_host points to us-east-1, not us-east-2...") + 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 points to us-east-1, not us-east-2...") + print(f"q_host: {manager._q_host}") + assert "us-east-1" in manager._q_host + assert "us-east-2" not in manager._q_host + + print("Verification: sso_region stores the credentials region...") + print(f"Comparing _sso_region: Expected 'us-east-2', Got '{manager._sso_region}'") + assert manager._sso_region == "us-east-2" + + print("Verification: refresh_url uses sso_region (us-east-2) with correct format...") + print(f"refresh_url: {manager._refresh_url}") + expected_refresh_url = "https://prod.us-east-2.auth.desktop.kiro.dev/refreshToken" + print(f"Comparing _refresh_url: Expected '{expected_refresh_url}', Got '{manager._refresh_url}'") + assert manager._refresh_url == expected_refresh_url + + def test_sso_region_set_from_credentials_file(self, tmp_path): + """ + What it does: Verifies _sso_region is populated from credentials file region field. + Purpose: Ensure auth/refresh requests use the correct regional endpoint. + """ + print("Setup: Creating credentials file with eu-central-1 region...") + creds_file = tmp_path / "kiro-auth-token.json" + creds_data = { + "refreshToken": "test_refresh_token", + "region": "eu-central-1" + } + creds_file.write_text(json.dumps(creds_data)) + + manager = KiroAuthManager(creds_file=str(creds_file), region="us-east-1") + + print("Verification: _sso_region set to eu-central-1...") + print(f"Comparing _sso_region: Expected 'eu-central-1', Got '{manager._sso_region}'") + assert manager._sso_region == "eu-central-1" + + print("Verification: API region unchanged...") + assert manager._region == "us-east-1" + + def test_sso_region_none_when_credentials_file_has_no_region(self, tmp_path): + """ + What it does: Verifies _sso_region is None when credentials file has no region field. + Purpose: Ensure backward compatibility with credentials files that omit region. + """ + print("Setup: Creating credentials file without region field...") + creds_file = tmp_path / "kiro-auth-token.json" + creds_data = { + "refreshToken": "test_refresh_token", + "accessToken": "test_access_token" + } + creds_file.write_text(json.dumps(creds_data)) + + manager = KiroAuthManager(creds_file=str(creds_file), region="us-east-1") + + print("Verification: _sso_region is None...") + print(f"Comparing _sso_region: Expected None, Got '{manager._sso_region}'") + assert manager._sso_region is None + class TestKiroAuthManagerTokenExpiration: """Tests for token expiration checking.""" @@ -1686,6 +1774,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 # =============================================================================