feat(model-fetch): fetch Kiro models from upstream

- add Kiro ListAvailableModels request planning and headers - route Kiro model refresh through upstream fetch - normalize Kiro model payloads and default model metadata - remove profileArn requirement from ListAvailableModels
2026-09-01 17:00:21 +08:00 · 2026-05-08 13:24:27 +08:00
parent fa22384f24
commit 2543676437
8 changed files with 475 additions and 21 deletions
--- a/crates/aether-provider-transport/src/kiro/headers.rs
+++ b/crates/aether-provider-transport/src/kiro/headers.rs
@@ -9,6 +9,7 @@ pub const KIRO_PROFILE_ARN_HEADER: &str = "x-amzn-kiro-profile-arn";
 pub const KIRO_TOKEN_TYPE_HEADER: &str = "TokenType";
 pub const KIRO_EXTERNAL_IDP_TOKEN_TYPE: &str = "EXTERNAL_IDP";
 const AWS_SDK_JS_MAIN_VERSION: &str = "1.0.27";
+const AWS_SDK_JS_LIST_MODELS_VERSION: &str = "1.0.0";
 const CODEWHISPERER_OPTOUT: &str = "true";
 const KIRO_AGENT_MODE: &str = "vibe";

@@ -126,12 +127,48 @@ pub fn build_mcp_headers(
    headers
 }

+pub fn build_list_available_models_headers(
+    auth_config: &KiroAuthConfig,
+    machine_id: &str,
+) -> BTreeMap<String, String> {
+    let kiro_version = auth_config.effective_kiro_version();
+    let system_version = auth_config.effective_system_version();
+    let node_version = auth_config.effective_node_version();
+    let region = auth_config.effective_api_region();
+    let host = format!("q.{region}.amazonaws.com");
+    let ide_tag = build_kiro_ide_tag(kiro_version, machine_id);
+
+    BTreeMap::from([
+        ("accept".to_string(), "application/json".to_string()),
+        (
+            "amz-sdk-invocation-id".to_string(),
+            Uuid::new_v4().to_string(),
+        ),
+        (
+            "amz-sdk-request".to_string(),
+            "attempt=1; max=1".to_string(),
+        ),
+        ("connection".to_string(), "close".to_string()),
+        ("host".to_string(), host),
+        (
+            "user-agent".to_string(),
+            format!(
+                "aws-sdk-js/{AWS_SDK_JS_LIST_MODELS_VERSION} ua/2.1 os/{system_version} lang/js md/nodejs#{node_version} api/codewhispererruntime#1.0.0 m/N,E {ide_tag}"
+            ),
+        ),
+        (
+            "x-amz-user-agent".to_string(),
+            format!("aws-sdk-js/{AWS_SDK_JS_LIST_MODELS_VERSION} {ide_tag}"),
+        ),
+    ])
+}
+
 #[cfg(test)]
 mod tests {
    use super::super::credentials::KiroAuthConfig;
    use super::{
-        build_generate_assistant_headers, build_mcp_headers, AWS_EVENTSTREAM_CONTENT_TYPE,
-        KIRO_PROFILE_ARN_HEADER, KIRO_TOKEN_TYPE_HEADER,
+        build_generate_assistant_headers, build_list_available_models_headers, build_mcp_headers,
+        AWS_EVENTSTREAM_CONTENT_TYPE, KIRO_PROFILE_ARN_HEADER, KIRO_TOKEN_TYPE_HEADER,
    };

    #[test]
@@ -169,6 +206,42 @@ mod tests {
        assert!(!headers.contains_key(KIRO_TOKEN_TYPE_HEADER));
    }

+    #[test]
+    fn builds_list_available_models_headers() {
+        let auth_config = KiroAuthConfig {
+            auth_method: Some("social".to_string()),
+            refresh_token: None,
+            expires_at: None,
+            profile_arn: Some("arn:aws:codewhisperer:us-east-1:123456789012:profile/demo".into()),
+            region: None,
+            auth_region: None,
+            api_region: Some("us-east-1".to_string()),
+            client_id: None,
+            client_secret: None,
+            machine_id: None,
+            kiro_version: Some("0.12.155".to_string()),
+            system_version: Some("darwin#24.6.0".to_string()),
+            node_version: Some("22.21.1".to_string()),
+            access_token: None,
+        };
+
+        let headers = build_list_available_models_headers(&auth_config, "machine-123");
+
+        assert_eq!(
+            headers.get("host").map(String::as_str),
+            Some("q.us-east-1.amazonaws.com")
+        );
+        assert_eq!(
+            headers.get("amz-sdk-request").map(String::as_str),
+            Some("attempt=1; max=1")
+        );
+        assert_eq!(
+            headers.get("x-amz-user-agent").map(String::as_str),
+            Some("aws-sdk-js/1.0.0 KiroIDE-0.12.155-machine-123")
+        );
+        assert!(!headers.contains_key(KIRO_TOKEN_TYPE_HEADER));
+    }
+
    #[test]
    fn keeps_generate_assistant_headers_without_external_idp_token_type() {
        let auth_config = KiroAuthConfig {
--- a/crates/aether-provider-transport/src/kiro/mod.rs
+++ b/crates/aether-provider-transport/src/kiro/mod.rs
@@ -17,7 +17,7 @@ pub const RUNTIME_POLICY: ProviderRuntimePolicy = ProviderRuntimePolicy {
    enable_format_conversion_by_default: true,
    allow_auth_channel_mismatch_by_default: true,
    oauth_is_bearer_like: true,
-    supports_model_fetch: false,
+    supports_model_fetch: true,
    supports_local_openai_chat_transport: false,
    supports_local_same_format_transport: false,
    local_embedding_support: ProviderLocalEmbeddingSupport::None,
@@ -32,8 +32,9 @@ pub use auth::{
 pub use converter::convert_claude_messages_to_conversation_state;
 pub use credentials::{generate_machine_id, normalize_machine_id, KiroAuthConfig};
 pub use headers::{
-    build_generate_assistant_headers, build_mcp_headers, AWS_EVENTSTREAM_CONTENT_TYPE,
-    KIRO_EXTERNAL_IDP_TOKEN_TYPE, KIRO_PROFILE_ARN_HEADER, KIRO_TOKEN_TYPE_HEADER,
+    build_generate_assistant_headers, build_list_available_models_headers, build_mcp_headers,
+    AWS_EVENTSTREAM_CONTENT_TYPE, KIRO_EXTERNAL_IDP_TOKEN_TYPE, KIRO_PROFILE_ARN_HEADER,
+    KIRO_TOKEN_TYPE_HEADER,
 };
 pub use policy::{
    local_kiro_request_transport_unsupported_reason_with_network,
@@ -47,7 +48,8 @@ pub use request::{
    supports_local_kiro_request_shape, KiroProviderHeadersInput,
 };
 pub use url::{
-    build_kiro_generate_assistant_response_url, build_kiro_mcp_url,
-    build_kiro_mcp_url_from_resolved_url, resolve_kiro_base_url, GENERATE_ASSISTANT_RESPONSE_PATH,
-    KIRO_ENVELOPE_NAME, MCP_PATH, MCP_STREAM_PATH,
+    build_kiro_generate_assistant_response_url, build_kiro_list_available_models_url,
+    build_kiro_mcp_url, build_kiro_mcp_url_from_resolved_url, resolve_kiro_base_url,
+    GENERATE_ASSISTANT_RESPONSE_PATH, KIRO_ENVELOPE_NAME, LIST_AVAILABLE_MODELS_PATH, MCP_PATH,
+    MCP_STREAM_PATH,
 };
--- a/crates/aether-provider-transport/src/kiro/url.rs
+++ b/crates/aether-provider-transport/src/kiro/url.rs
@@ -2,6 +2,7 @@ use super::super::url::build_passthrough_path_url;
 use super::credentials::DEFAULT_REGION;

 pub const GENERATE_ASSISTANT_RESPONSE_PATH: &str = "/generateAssistantResponse";
+pub const LIST_AVAILABLE_MODELS_PATH: &str = "/ListAvailableModels";
 pub const MCP_PATH: &str = "/mcp";
 pub const MCP_STREAM_PATH: &str = "/mcp/stream";
 pub const KIRO_ENVELOPE_NAME: &str = "kiro:generateAssistantResponse";
@@ -37,6 +38,22 @@ pub fn build_kiro_mcp_url(upstream_base_url: &str, api_region: Option<&str>) ->
    build_passthrough_path_url(upstream_base_url.as_str(), MCP_PATH, None, &[])
 }

+pub fn build_kiro_list_available_models_url(
+    upstream_base_url: &str,
+    api_region: Option<&str>,
+) -> Option<String> {
+    let mut serializer = url::form_urlencoded::Serializer::new(String::new());
+    serializer.append_pair("origin", "AI_EDITOR");
+    let query = serializer.finish();
+    let upstream_base_url = resolve_kiro_base_url(upstream_base_url, api_region);
+    build_passthrough_path_url(
+        upstream_base_url.as_str(),
+        LIST_AVAILABLE_MODELS_PATH,
+        Some(&query),
+        &[],
+    )
+}
+
 pub fn build_kiro_mcp_url_from_resolved_url(resolved_url: &str) -> Option<String> {
    let mut parsed = url::Url::parse(resolved_url).ok()?;
    parsed.set_path(MCP_PATH);
@@ -47,9 +64,10 @@ pub fn build_kiro_mcp_url_from_resolved_url(resolved_url: &str) -> Option<String
 #[cfg(test)]
 mod tests {
    use super::{
-        build_kiro_generate_assistant_response_url, build_kiro_mcp_url,
-        build_kiro_mcp_url_from_resolved_url, resolve_kiro_base_url,
-        GENERATE_ASSISTANT_RESPONSE_PATH, KIRO_ENVELOPE_NAME, MCP_PATH, MCP_STREAM_PATH,
+        build_kiro_generate_assistant_response_url, build_kiro_list_available_models_url,
+        build_kiro_mcp_url, build_kiro_mcp_url_from_resolved_url, resolve_kiro_base_url,
+        GENERATE_ASSISTANT_RESPONSE_PATH, KIRO_ENVELOPE_NAME, LIST_AVAILABLE_MODELS_PATH, MCP_PATH,
+        MCP_STREAM_PATH,
    };

    #[test]
@@ -58,6 +76,7 @@ mod tests {
            GENERATE_ASSISTANT_RESPONSE_PATH,
            "/generateAssistantResponse"
        );
+        assert_eq!(LIST_AVAILABLE_MODELS_PATH, "/ListAvailableModels");
        assert_eq!(MCP_PATH, "/mcp");
        assert_eq!(MCP_STREAM_PATH, "/mcp/stream");
        assert_eq!(KIRO_ENVELOPE_NAME, "kiro:generateAssistantResponse");
@@ -94,6 +113,18 @@ mod tests {
        );
    }

+    #[test]
+    fn builds_list_available_models_url_for_profile() {
+        assert_eq!(
+            build_kiro_list_available_models_url(
+                "https://q.{region}.amazonaws.com",
+                Some("us-west-2")
+            )
+            .as_deref(),
+            Some("https://q.us-west-2.amazonaws.com/ListAvailableModels?origin=AI_EDITOR")
+        );
+    }
+
    #[test]
    fn rewrites_generate_assistant_url_to_mcp_url() {
        assert_eq!(
--- a/crates/aether-provider-transport/src/provider_types.rs
+++ b/crates/aether-provider-transport/src/provider_types.rs
@@ -627,7 +627,7 @@ mod tests {
        assert!(policy.fixed_provider);
        assert!(policy.enable_format_conversion_by_default);
        assert!(policy.oauth_is_bearer_like);
-        assert!(!policy.supports_model_fetch);
+        assert!(policy.supports_model_fetch);
        assert!(!policy.supports_local_openai_chat_transport);
        assert!(!policy.supports_local_same_format_transport);
        assert!(policy.key_inherits_api_formats("oauth", None));