Skip to content

fix(models): fetch ListAvailableModels from q.{region}.amazonaws.com #293

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
harlan-huang wants to merge 1 commit into
base: main
Choose a base branch
from
Open

fix(models): fetch ListAvailableModels from q.{region}.amazonaws.com #293

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
119 changes: 37 additions & 82 deletions kiro/account_manager.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -64,35 +64,6 @@
from kiro.http_client import KiroHttpClient


def _is_runtime_endpoint(auth_manager: KiroAuthManager) -> bool:
"""
Check if auth manager uses runtime endpoint that doesn't provide /ListAvailableModels.

Runtime endpoint pattern: https://runtime.{region}.kiro.dev
Old endpoint pattern: https://q.{region}.amazonaws.com

Runtime endpoint does not provide /ListAvailableModels API (AWS limitation).

Args:
auth_manager: KiroAuthManager instance

Returns:
True if using runtime endpoint, False otherwise

Examples:
>>> auth_manager.api_host = "https://runtime.us-east-1.kiro.dev"
>>> _is_runtime_endpoint(auth_manager)
True
>>> auth_manager.api_host = "https://runtime.eu-central-1.kiro.dev"
>>> _is_runtime_endpoint(auth_manager)
True
>>> auth_manager.api_host = "https://q.us-east-1.amazonaws.com"
>>> _is_runtime_endpoint(auth_manager)
False
"""
return "://runtime." in auth_manager.api_host


def _format_duration(seconds: float) -> str:
"""
Format duration in human-readable format.
Expand Down Expand Up @@ -497,48 +468,42 @@ async def _initialize_account(self, account_id: str) -> bool:
# Get token to verify credentials
token = await auth_manager.get_access_token()

# Determine if we should fetch models or use static list
if _is_runtime_endpoint(auth_manager):
# New runtime endpoint does not provide /ListAvailableModels (AWS limitation)
# Use static list without attempting request
logger.debug(f"Account {account_id}: Using static model list for runtime.kiro.dev endpoint")
# Fetch the account's real model list. ListAvailableModels lives on
# models_host, not on the runtime host used for generation, so this
# works for every endpoint. Fall back to the static list on failure.
params = {"origin": "AI_EDITOR"}
if auth_manager.auth_type == AuthType.KIRO_DESKTOP and auth_manager.profile_arn:
params["profileArn"] = auth_manager.profile_arn

list_models_url = f"{auth_manager.models_host}/ListAvailableModels"

# Use KiroHttpClient for retry logic (3 attempts with exponential backoff)
http_client = KiroHttpClient(auth_manager, shared_client=None)

try:
response = await http_client.request_with_retry(
method="GET",
url=list_models_url,
json_data=None,
params=params,
stream=False
)

if response.status_code == 200:
data = response.json()
models_list = data.get("models", [])
else:
# Shouldn't happen (retry handles non-200), but keep for safety
raise Exception(f"HTTP {response.status_code}")

except Exception as e:
# All retries exhausted - use fallback
logger.error(f"Failed to fetch models for {account_id} after retries: {e}")
logger.warning("Using pre-configured fallback models. Models will be refreshed on next TTL cycle when network recovers.")
models_list = FALLBACK_MODELS
else:
# Old endpoint - attempt to fetch dynamic model list
# Fetch models list with retry + fallback
params = {"origin": "AI_EDITOR"}
if auth_manager.auth_type == AuthType.KIRO_DESKTOP and auth_manager.profile_arn:
params["profileArn"] = auth_manager.profile_arn

list_models_url = f"{auth_manager.q_host}/ListAvailableModels"

# Use KiroHttpClient for retry logic (3 attempts with exponential backoff)
http_client = KiroHttpClient(auth_manager, shared_client=None)

try:
response = await http_client.request_with_retry(
method="GET",
url=list_models_url,
json_data=None,
params=params,
stream=False
)

if response.status_code == 200:
data = response.json()
models_list = data.get("models", [])
else:
# Shouldn't happen (retry handles non-200), but keep for safety
raise Exception(f"HTTP {response.status_code}")

except Exception as e:
# All retries exhausted - use fallback
logger.error(f"Failed to fetch models for {account_id} after retries: {e}")
logger.warning("Using pre-configured fallback models. Models will be refreshed on next TTL cycle when network recovers.")
models_list = FALLBACK_MODELS

finally:
await http_client.close()

finally:
await http_client.close()

# Create model cache and update
model_cache = ModelInfoCache()
Expand Down Expand Up @@ -589,17 +554,7 @@ async def _refresh_account_models(self, account_id: str) -> None:
if not account or not account.auth_manager:
return

# Check if using runtime endpoint (no dynamic model list available)
if _is_runtime_endpoint(account.auth_manager):
# Runtime endpoint does not provide /ListAvailableModels
# Use static list and update cache timestamp
logger.debug(f"Account {account_id}: Skipping model refresh for runtime.kiro.dev endpoint (using static list)")
await account.model_cache.update(FALLBACK_MODELS)
account.models_cached_at = time.time()
self._dirty = True
return

# Old endpoint - attempt to fetch dynamic model list
# Fetch the account's real model list from models_host
# Use KiroHttpClient for retry logic
http_client = KiroHttpClient(account.auth_manager, shared_client=None)

Expand All @@ -608,7 +563,7 @@ async def _refresh_account_models(self, account_id: str) -> None:
if account.auth_manager.auth_type == AuthType.KIRO_DESKTOP and account.auth_manager.profile_arn:
params["profileArn"] = account.auth_manager.profile_arn

list_models_url = f"{account.auth_manager.q_host}/ListAvailableModels"
list_models_url = f"{account.auth_manager.models_host}/ListAvailableModels"

response = await http_client.request_with_retry(
method="GET",
Expand Down
15 changes: 12 additions & 3 deletions kiro/auth.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -46,6 +46,7 @@
get_kiro_refresh_url,
get_kiro_api_host,
get_kiro_q_host,
get_kiro_models_host,
get_aws_sso_oidc_url,
)
from kiro.utils import get_machine_fingerprint
Expand Down Expand Up @@ -98,6 +99,7 @@ class KiroAuthManager:
region: AWS region
api_host: API host for the current region
q_host: Q API host for the current region
models_host: ListAvailableModels host for the current region
fingerprint: Unique machine fingerprint
auth_type: Type of authentication (KIRO_DESKTOP or AWS_SSO_OIDC)

Expand Down Expand Up @@ -221,14 +223,16 @@ def __init__(
self._refresh_url = get_kiro_refresh_url(sso_region_for_oidc)
self._api_host = get_kiro_api_host(final_api_region)
self._q_host = get_kiro_q_host(final_api_region)

self._models_host = get_kiro_models_host(final_api_region)

# Log initialized endpoints for diagnostics (helps with DNS issues like #58, #132, #133)
logger.info(
f"Auth manager initialized: "
f"sso_region={sso_region_for_oidc}, "
f"api_region={final_api_region}, "
f"api_host={self._api_host}, "
f"q_host={self._q_host}"
f"q_host={self._q_host}, "
f"models_host={self._models_host}"
)

def _detect_auth_type(self) -> None:
Expand Down Expand Up @@ -965,7 +969,12 @@ def api_host(self) -> str:
def q_host(self) -> str:
"""Q API host for the current region."""
return self._q_host


@property
def models_host(self) -> str:
"""ListAvailableModels host for the current region."""
return self._models_host

@property
def fingerprint(self) -> str:
"""Unique machine fingerprint."""
Expand Down
16 changes: 15 additions & 1 deletion kiro/config.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -181,9 +181,18 @@ def _get_raw_env_value(var_name: str, env_file: str = ".env") -> Optional[str]:
# Fixed in issue #58 - codewhisperer.{region}.amazonaws.com doesn't exist for non-us-east-1 regions
KIRO_API_HOST_TEMPLATE: str = "https://runtime.{region}.kiro.dev"

# Host for Q API (ListAvailableModels)
# Host for Q API (MCP tools)
KIRO_Q_HOST_TEMPLATE: str = "https://runtime.{region}.kiro.dev"

# Host for ListAvailableModels
# runtime.{region}.kiro.dev does not implement ListAvailableModels (it answers
# UnknownOperationException), while q.{region}.amazonaws.com still serves it and
# returns the models the account is actually entitled to, including models
# enabled or disabled by organization admins. Fetching from here keeps
# /v1/models in sync with what generateAssistantResponse accepts.
# If the request fails, FALLBACK_MODELS is used instead.
KIRO_MODELS_HOST_TEMPLATE: str = "https://q.{region}.amazonaws.com"

# ==================================================================================================
# Token Settings
# ==================================================================================================
Expand Down Expand Up @@ -579,3 +588,8 @@ def get_kiro_q_host(region: str) -> str:
"""Return Q API host for the specified region."""
return KIRO_Q_HOST_TEMPLATE.format(region=region)


def get_kiro_models_host(region: str) -> str:
"""Return ListAvailableModels host for the specified region."""
return KIRO_MODELS_HOST_TEMPLATE.format(region=region)

Loading