i/Aether
1
0
Fork 0
mirror of https://github.com/fawney19/Aether.git synced 2026-09-02 01:10:23 +08:00
Code Issues Packages Projects Releases Wiki Activity

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

Browse Source 
- 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
This commit is contained in:
AAEE86
2026-05-08 13:24:27 +08:00
parent fa22384f24
commit 2543676437
8 changed files with 475 additions and 21 deletions
Download Patch File Download Diff File Expand all files Collapse all files

5
crates/aether-model-fetch/src/lib.rs
View File

@@ -23,6 +23,7 @@ pub use strategy::{
};
pub use transport::{
build_antigravity_fetch_available_models_plan, build_gemini_cli_load_code_assist_plan,
build_models_fetch_execution_plan, build_standard_models_fetch_execution_plan,
build_vertex_models_fetch_execution_plan, ModelFetchTransportRuntime,
build_kiro_list_available_models_plan, build_models_fetch_execution_plan,
build_standard_models_fetch_execution_plan, build_vertex_models_fetch_execution_plan,
ModelFetchTransportRuntime,
};

44
crates/aether-model-fetch/src/logic.rs
View File

@@ -229,7 +229,7 @@ pub fn endpoint_supports_rust_models_fetch(api_format: &str) -> bool {
pub fn provider_type_uses_preset_models(provider_type: &str) -> bool {
matches!(
provider_type.trim().to_ascii_lowercase().as_str(),
"kiro" | "claude_code" | "gemini_cli"
"claude_code" | "gemini_cli"
)
}
@@ -244,11 +244,19 @@ pub fn preset_models_for_provider(provider_type: &str) -> Option<Vec<Value>> {
preset_model("gemini-3.1-pro-preview", "google", "Gemini 3.1 Pro Preview", "gemini:generate_content"),
],
"kiro" => vec![
preset_model("claude-sonnet-4.5", "anthropic", "Claude Sonnet 4.5", "claude:messages"),
preset_model("auto", "kiro", "Auto", "claude:messages"),
preset_model("claude-opus-4.7", "anthropic", "Claude Opus 4.7", "claude:messages"),
preset_model("claude-opus-4.6", "anthropic", "Claude Opus 4.6", "claude:messages"),
preset_model("claude-sonnet-4.6", "anthropic", "Claude Sonnet 4.6", "claude:messages"),
preset_model("claude-opus-4.5", "anthropic", "Claude Opus 4.5", "claude:messages"),
preset_model("claude-opus-4.6", "anthropic", "Claude Opus 4.6", "claude:messages"),
preset_model("claude-sonnet-4.5", "anthropic", "Claude Sonnet 4.5", "claude:messages"),
preset_model("claude-sonnet-4", "anthropic", "Claude Sonnet 4", "claude:messages"),
preset_model("claude-haiku-4.5", "anthropic", "Claude Haiku 4.5", "claude:messages"),
preset_model("deepseek-3.2", "deepseek", "Deepseek v3.2", "claude:messages"),
preset_model("minimax-m2.5", "minimax", "MiniMax M2.5", "claude:messages"),
preset_model("minimax-m2.1", "minimax", "MiniMax M2.1", "claude:messages"),
preset_model("glm-5", "zhipu", "GLM 5", "claude:messages"),
preset_model("qwen3-coder-next", "alibaba", "Qwen3 Coder Next", "claude:messages"),
],
"claude_code" => vec![
preset_model("claude-opus-4-5-20251101", "anthropic", "Claude Opus 4.5", "claude:messages"),
@@ -896,4 +904,34 @@ mod tests {
]
);
}
#[test]
fn preset_models_cover_kiro_catalog() {
let models = preset_models_for_provider("kiro").expect("preset models should exist");
let model_ids = models
.iter()
.map(|model| model["id"].as_str().expect("model id"))
.collect::<Vec<_>>();
assert_eq!(
model_ids,
vec![
"auto",
"claude-opus-4.7",
"claude-opus-4.6",
"claude-sonnet-4.6",
"claude-opus-4.5",
"claude-sonnet-4.5",
"claude-sonnet-4",
"claude-haiku-4.5",
"deepseek-3.2",
"minimax-m2.5",
"minimax-m2.1",
"glm-5",
"qwen3-coder-next",
]
);
assert!(models
.iter()
.all(|model| model["api_formats"] == json!(["claude:messages"])));
}
}

214
crates/aether-model-fetch/src/strategy.rs
View File

@@ -18,8 +18,8 @@ use sha2::Sha256;
use crate::logic::{extract_error_message, parse_models_response_page, preset_models_for_provider};
use crate::transport::{
build_antigravity_fetch_available_models_plan, build_gemini_cli_load_code_assist_plan,
build_standard_models_fetch_execution_plan, build_vertex_models_fetch_execution_plan,
ModelFetchTransportRuntime,
build_kiro_list_available_models_plan, build_standard_models_fetch_execution_plan,
build_vertex_models_fetch_execution_plan, ModelFetchTransportRuntime,
};
const ANTIGRAVITY_SANDBOX_BASE_URL: &str = "https://daily-cloudcode-pa.sandbox.googleapis.com";
@@ -48,6 +48,7 @@ pub enum ModelFetchStrategyKind {
Vertex,
Antigravity,
GeminiCliPreset,
Kiro,
}
pub trait ModelFetchStrategy {
@@ -94,6 +95,13 @@ fn select_model_fetch_strategy(
.trim()
.to_ascii_lowercase();
if let Some(models) = preset_models_for_provider(&provider_type) {
if provider_type == "kiro" {
return Ok(SelectedModelFetchStrategy {
provider_type,
kind: ModelFetchStrategyKind::Kiro,
preset_models: None,
});
}
if provider_type == "codex" {
return Ok(SelectedModelFetchStrategy {
provider_type,
@@ -165,6 +173,7 @@ async fn execute_model_fetch_strategy(
)
.await
}
ModelFetchStrategyKind::Kiro => fetch_kiro_models(runtime, first_transport).await,
}
}
@@ -340,6 +349,21 @@ async fn fetch_gemini_cli_models(
Ok(build_success_outcome(models, upstream_metadata, true))
}
async fn fetch_kiro_models(
runtime: &(impl ModelFetchTransportRuntime + ?Sized),
transport: &GatewayProviderTransportSnapshot,
) -> Result<ModelsFetchOutcome, String> {
let plan = build_kiro_list_available_models_plan(runtime, transport).await?;
let result = runtime.execute_model_fetch_execution_plan(&plan).await?;
if !(200..300).contains(&result.status_code) {
return Err(execution_result_error_message(&result));
}
let body_json = execution_result_json_body_allow_empty(&result)?;
let (models, metadata) = parse_kiro_available_models_response(&body_json)?;
Ok(build_success_outcome(models, metadata, true))
}
async fn fetch_vertex_models(
runtime: &(impl ModelFetchTransportRuntime + ?Sized),
transports: &[GatewayProviderTransportSnapshot],
@@ -735,6 +759,94 @@ fn parse_antigravity_models_response(body: &Value) -> Result<(Vec<Value>, Option
Ok((models, upstream_metadata))
}
fn parse_kiro_available_models_response(
body: &Value,
) -> Result<(Vec<Value>, Option<Value>), String> {
let items = body
.get("models")
.and_then(Value::as_array)
.ok_or_else(|| "kiro: invalid response (missing models)".to_string())?;
let mut seen = BTreeSet::new();
let mut models = Vec::new();
for item in items {
let Some(model_id) = item
.get("modelId")
.or_else(|| item.get("id"))
.and_then(Value::as_str)
.map(str::trim)
.filter(|value| !value.is_empty())
else {
continue;
};
if !seen.insert(model_id.to_string()) {
continue;
}
let display_name = item
.get("modelName")
.or_else(|| item.get("display_name"))
.or_else(|| item.get("name"))
.and_then(Value::as_str)
.map(str::trim)
.filter(|value| !value.is_empty())
.unwrap_or(model_id);
let mut model = item.as_object().cloned().unwrap_or_default();
model.insert("id".to_string(), Value::String(model_id.to_string()));
model.insert("object".to_string(), Value::String("model".to_string()));
model.insert(
"owned_by".to_string(),
Value::String(infer_kiro_model_owner(model_id).to_string()),
);
model.insert(
"display_name".to_string(),
Value::String(display_name.to_string()),
);
model.insert(
"api_formats".to_string(),
Value::Array(vec![Value::String("claude:messages".to_string())]),
);
model.remove("api_format");
models.push(Value::Object(model));
}
let default_model = body.get("defaultModel").and_then(|value| {
json_string(value.get("modelId")).map(|model_id| {
json!({
"model_id": model_id,
"model_name": json_string(value.get("modelName")),
})
})
});
let upstream_metadata = default_model.map(|default_model| {
json!({
"kiro": {
"updated_at": now_unix_secs(),
"default_model": default_model,
}
})
});
Ok((models, upstream_metadata))
}
fn infer_kiro_model_owner(model_id: &str) -> &'static str {
let normalized = model_id.trim().to_ascii_lowercase();
if normalized.starts_with("claude-") {
"anthropic"
} else if normalized.starts_with("deepseek-") {
"deepseek"
} else if normalized.starts_with("minimax-") {
"minimax"
} else if normalized.starts_with("glm-") {
"zhipu"
} else if normalized.starts_with("qwen") {
"alibaba"
} else {
"kiro"
}
}
fn build_antigravity_quota_payload(quota_info: Option<&Value>) -> serde_json::Map<String, Value> {
let quota_info = quota_info.and_then(Value::as_object);
let reset_time = quota_info
@@ -1298,6 +1410,31 @@ mod tests {
transport
}
fn sample_kiro_transport() -> GatewayProviderTransportSnapshot {
let mut transport = sample_custom_aiplatform_transport();
transport.provider.provider_type = "kiro".to_string();
transport.provider.name = "Kiro".to_string();
transport.endpoint.api_format = "claude:messages".to_string();
transport.endpoint.api_family = Some("claude".to_string());
transport.endpoint.endpoint_kind = Some("messages".to_string());
transport.endpoint.base_url = "https://q.{region}.amazonaws.com".to_string();
transport.endpoint.custom_path = None;
transport.key.auth_type = "oauth".to_string();
transport.key.api_formats = Some(vec!["claude:messages".to_string()]);
transport.key.decrypted_api_key = "__placeholder__".to_string();
transport.key.decrypted_auth_config = Some(
r#"{
"access_token":"cached-token",
"expires_at":4102444800,
"profile_arn":"arn:aws:codewhisperer:us-east-1:123456789012:profile/demo",
"api_region":"us-east-1",
"machine_id":"123e4567-e89b-12d3-a456-426614174000"
}"#
.to_string(),
);
transport
}
#[test]
fn strategy_selection_keeps_codex_on_standard_transport_fetch() {
let strategy = select_model_fetch_strategy(&[sample_codex_transport()])
@@ -1319,6 +1456,15 @@ mod tests {
assert_eq!(strategy.kind(), ModelFetchStrategyKind::PresetCatalog);
}
#[test]
fn strategy_selection_uses_kiro_upstream_fetch() {
let strategy = select_model_fetch_strategy(&[sample_kiro_transport()])
.expect("strategy should select");
assert_eq!(strategy.provider_id(), "kiro");
assert_eq!(strategy.kind(), ModelFetchStrategyKind::Kiro);
}
#[tokio::test]
async fn custom_aiplatform_transport_uses_vertex_models_fetch_path_and_normalizes_chat_format()
{
@@ -1372,4 +1518,68 @@ mod tests {
assert_eq!(outcome.fetched_model_ids, vec!["gpt-5.4-upstream"]);
assert_eq!(outcome.cached_models.len(), 1);
}
#[tokio::test]
async fn kiro_transport_fetches_list_available_models() {
let executed_urls = Arc::new(Mutex::new(Vec::new()));
let runtime = TestRuntime {
executed_urls: Arc::clone(&executed_urls),
response_body: json!({
"defaultModel": {
"modelId": "auto",
"modelName": "Auto"
},
"models": [
{
"modelId": "auto",
"modelName": "Auto",
"tokenLimits": {
"maxInputTokens": 1000000,
"maxOutputTokens": 64000
}
},
{
"modelId": "claude-opus-4.7",
"modelName": "Claude Opus 4.7",
"description": "Experimental preview"
}
]
}),
};
let outcome = fetch_models_from_transports(&runtime, &[sample_kiro_transport()])
.await
.expect("models fetch should succeed");
let urls = executed_urls.lock().expect("executed_urls lock");
assert_eq!(
urls.as_slice(),
&["https://q.us-east-1.amazonaws.com/ListAvailableModels?origin=AI_EDITOR"]
);
assert_eq!(
outcome.fetched_model_ids,
vec!["auto".to_string(), "claude-opus-4.7".to_string()]
);
assert_eq!(outcome.cached_models.len(), 2);
assert_eq!(
outcome.cached_models[1]["display_name"].as_str(),
Some("Claude Opus 4.7")
);
assert_eq!(
outcome.cached_models[1]["owned_by"].as_str(),
Some("anthropic")
);
assert_eq!(
outcome.cached_models[1]["api_formats"],
json!(["claude:messages"])
);
assert_eq!(
outcome.upstream_metadata.as_ref().and_then(|value| {
value
.get("kiro")
.and_then(|value| value.get("default_model"))
.and_then(|value| value.get("model_id"))
}),
Some(&json!("auto"))
);
}
}

103
crates/aether-model-fetch/src/transport.rs
View File

@@ -9,6 +9,10 @@ use aether_provider_transport::auth::{
ensure_upstream_auth_header, resolve_local_gemini_auth, resolve_local_openai_bearer_auth,
resolve_local_standard_auth,
};
use aether_provider_transport::kiro::{
build_kiro_list_available_models_url, build_list_available_models_headers,
resolve_local_kiro_request_auth,
};
use aether_provider_transport::vertex::resolve_local_vertex_api_key_query_auth;
use aether_provider_transport::{
apply_local_header_rules, resolve_transport_execution_timeouts, resolve_transport_profile,
@@ -25,6 +29,7 @@ const GEMINI_CLI_USER_AGENT: &str = "GeminiCLI/0.1.5 (Windows; AMD64)";
const CLAUDE_VERSION_HEADER: &str = "2023-06-01";
const ANTIGRAVITY_FETCH_PROVIDER_API_FORMAT: &str = "antigravity:fetch_available_models";
const GEMINI_CLI_LOAD_CODE_ASSIST_PROVIDER_API_FORMAT: &str = "gemini_cli:load_code_assist";
const KIRO_LIST_AVAILABLE_MODELS_PROVIDER_API_FORMAT: &str = "kiro:list_available_models";
const BROWSER_FINGERPRINT_HEADERS: &[(&str, &str)] = &[
(
@@ -241,6 +246,55 @@ pub async fn build_gemini_cli_load_code_assist_plan(
.await
}
pub async fn build_kiro_list_available_models_plan(
runtime: &(impl ModelFetchTransportRuntime + ?Sized),
transport: &GatewayProviderTransportSnapshot,
) -> Result<ExecutionPlan, String> {
let kiro_auth = match runtime.resolve_local_oauth_request_auth(transport).await? {
Some(LocalResolvedOAuthRequestAuth::Kiro(auth)) => Some(auth),
_ => resolve_local_kiro_request_auth(transport),
}
.ok_or_else(|| "Kiro models fetch requires Kiro request auth".to_string())?;
let url = build_kiro_list_available_models_url(
&transport.endpoint.base_url,
Some(kiro_auth.auth_config.effective_api_region()),
)
.ok_or_else(|| "Kiro models fetch URL is unavailable".to_string())?;
let mut headers =
build_list_available_models_headers(&kiro_auth.auth_config, &kiro_auth.machine_id);
let mut protected_headers = Vec::new();
insert_non_empty_auth_header(
&mut headers,
&mut protected_headers,
kiro_auth.name,
&kiro_auth.value,
);
protected_headers.extend(["host".to_string(), "x-amz-user-agent".to_string()]);
headers = apply_fetch_header_rules(transport, headers, &protected_headers)?;
ensure_upstream_auth_header(&mut headers, kiro_auth.name, &kiro_auth.value);
build_execution_plan(
runtime,
transport,
ModelFetchExecutionPlanRequest {
method: "GET".to_string(),
url,
headers,
content_type: None,
body: RequestBody {
json_body: None,
body_bytes_b64: None,
body_ref: None,
},
client_api_format: "claude:messages".to_string(),
provider_api_format: KIRO_LIST_AVAILABLE_MODELS_PROVIDER_API_FORMAT.to_string(),
model_name: Some("ListAvailableModels".to_string()),
},
)
.await
}
pub async fn build_vertex_models_fetch_execution_plan(
runtime: &(impl ModelFetchTransportRuntime + ?Sized),
transport: &GatewayProviderTransportSnapshot,
@@ -544,8 +598,9 @@ mod tests {
use super::{
build_antigravity_fetch_available_models_plan, build_gemini_cli_load_code_assist_plan,
build_models_fetch_execution_plan, build_standard_models_fetch_execution_plan,
build_vertex_models_fetch_execution_plan, ModelFetchTransportRuntime,
build_kiro_list_available_models_plan, build_models_fetch_execution_plan,
build_standard_models_fetch_execution_plan, build_vertex_models_fetch_execution_plan,
ModelFetchTransportRuntime,
};
struct TestRuntime {
@@ -833,6 +888,50 @@ mod tests {
);
}
#[tokio::test]
async fn builds_kiro_list_available_models_plan() {
let runtime = TestRuntime {
oauth_auth: None,
proxy: None,
};
let mut transport = sample_transport("kiro", "claude:messages", "oauth");
transport.endpoint.base_url = "https://q.{region}.amazonaws.com".to_string();
transport.key.decrypted_api_key = "__placeholder__".to_string();
transport.key.decrypted_auth_config = Some(
r#"{
"access_token":"cached-token",
"expires_at":4102444800,
"api_region":"us-west-2",
"machine_id":"123e4567-e89b-12d3-a456-426614174000",
"kiro_version":"0.12.155"
}"#
.to_string(),
);
let plan = build_kiro_list_available_models_plan(&runtime, &transport)
.await
.expect("plan");
assert_eq!(plan.method, "GET");
assert_eq!(
plan.url,
"https://q.us-west-2.amazonaws.com/ListAvailableModels?origin=AI_EDITOR"
);
assert_eq!(plan.provider_api_format, "kiro:list_available_models");
assert_eq!(
plan.headers.get("authorization").map(String::as_str),
Some("Bearer cached-token")
);
assert_eq!(
plan.headers.get("host").map(String::as_str),
Some("q.us-west-2.amazonaws.com")
);
assert!(plan
.headers
.get("x-amz-user-agent")
.is_some_and(|value| value.starts_with("aws-sdk-js/1.0.0 KiroIDE-0.12.155-")));
}
#[tokio::test]
async fn builds_vertex_models_fetch_plan_with_auth_override() {
let runtime = TestRuntime {