mirror of
https://github.com/fawney19/Aether.git
synced 2026-09-02 17:30:23 +08:00
refactor: 拆分 gateway 单体为独立 crate,新增 systemd 部署方案
将 gateway 内部的 model-fetch、provider-transport、scheduler-core、 usage-runtime、video-tasks-core 模块提取为独立 crate;重构 gateway 内部模块结构(state/router/cache/data/query 等);移除大量遗留模块 文件;新增 systemd 二进制部署骨架及相关文档;更新前端 usage 相关 API 和组件。
This commit is contained in:
311
crates/aether-provider-transport/src/kiro/auth.rs
Normal file
311
crates/aether-provider-transport/src/kiro/auth.rs
Normal file
@@ -0,0 +1,311 @@
|
||||
use super::super::snapshot::GatewayProviderTransportSnapshot;
|
||||
use super::credentials::{generate_machine_id, KiroAuthConfig};
|
||||
|
||||
pub const PROVIDER_TYPE: &str = "kiro";
|
||||
pub const KIRO_AUTH_HEADER: &str = "authorization";
|
||||
|
||||
#[derive(Debug, Clone, PartialEq, Eq)]
|
||||
pub struct KiroBearerAuth {
|
||||
pub name: &'static str,
|
||||
pub value: String,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, PartialEq, Eq)]
|
||||
pub struct KiroRequestAuth {
|
||||
pub name: &'static str,
|
||||
pub value: String,
|
||||
pub auth_config: KiroAuthConfig,
|
||||
pub machine_id: String,
|
||||
}
|
||||
|
||||
pub fn build_kiro_request_auth_from_config(
|
||||
auth_config: KiroAuthConfig,
|
||||
fallback_secret: Option<&str>,
|
||||
) -> Option<KiroRequestAuth> {
|
||||
let fallback_secret = fallback_secret
|
||||
.map(str::trim)
|
||||
.filter(|value| !value.is_empty() && *value != "__placeholder__");
|
||||
let token = auth_config
|
||||
.cached_access_token()
|
||||
.filter(|_| !auth_config.cached_access_token_requires_refresh(120))
|
||||
.or(fallback_secret)?;
|
||||
let machine_id = generate_machine_id(&auth_config, Some(token))?;
|
||||
|
||||
Some(KiroRequestAuth {
|
||||
name: KIRO_AUTH_HEADER,
|
||||
value: format!("Bearer {token}"),
|
||||
auth_config,
|
||||
machine_id,
|
||||
})
|
||||
}
|
||||
|
||||
pub fn resolve_local_kiro_bearer_auth(
|
||||
transport: &GatewayProviderTransportSnapshot,
|
||||
) -> Option<KiroBearerAuth> {
|
||||
if !transport
|
||||
.provider
|
||||
.provider_type
|
||||
.trim()
|
||||
.eq_ignore_ascii_case(PROVIDER_TYPE)
|
||||
{
|
||||
return None;
|
||||
}
|
||||
if transport.key.decrypted_auth_config.is_some() {
|
||||
return None;
|
||||
}
|
||||
if !transport
|
||||
.key
|
||||
.auth_type
|
||||
.trim()
|
||||
.eq_ignore_ascii_case("bearer")
|
||||
{
|
||||
return None;
|
||||
}
|
||||
|
||||
let secret = transport.key.decrypted_api_key.trim();
|
||||
if secret.is_empty() {
|
||||
return None;
|
||||
}
|
||||
|
||||
Some(KiroBearerAuth {
|
||||
name: KIRO_AUTH_HEADER,
|
||||
value: format!("Bearer {secret}"),
|
||||
})
|
||||
}
|
||||
|
||||
pub fn supports_local_kiro_auth_prerequisites(
|
||||
transport: &GatewayProviderTransportSnapshot,
|
||||
) -> bool {
|
||||
resolve_local_kiro_bearer_auth(transport).is_some()
|
||||
}
|
||||
|
||||
pub fn resolve_local_kiro_request_auth(
|
||||
transport: &GatewayProviderTransportSnapshot,
|
||||
) -> Option<KiroRequestAuth> {
|
||||
if !transport
|
||||
.provider
|
||||
.provider_type
|
||||
.trim()
|
||||
.eq_ignore_ascii_case(PROVIDER_TYPE)
|
||||
{
|
||||
return None;
|
||||
}
|
||||
if !transport
|
||||
.key
|
||||
.auth_type
|
||||
.trim()
|
||||
.eq_ignore_ascii_case("bearer")
|
||||
{
|
||||
return None;
|
||||
}
|
||||
|
||||
let auth_config = KiroAuthConfig::from_raw_json(transport.key.decrypted_auth_config.as_deref())
|
||||
.unwrap_or(KiroAuthConfig {
|
||||
auth_method: None,
|
||||
refresh_token: None,
|
||||
expires_at: None,
|
||||
profile_arn: None,
|
||||
region: None,
|
||||
auth_region: None,
|
||||
api_region: None,
|
||||
client_id: None,
|
||||
client_secret: None,
|
||||
machine_id: None,
|
||||
kiro_version: None,
|
||||
system_version: None,
|
||||
node_version: None,
|
||||
access_token: None,
|
||||
});
|
||||
let fallback_secret = transport
|
||||
.key
|
||||
.decrypted_api_key
|
||||
.trim()
|
||||
.strip_prefix("__placeholder__")
|
||||
.map(|_| "")
|
||||
.unwrap_or(transport.key.decrypted_api_key.trim());
|
||||
build_kiro_request_auth_from_config(auth_config, Some(fallback_secret))
|
||||
}
|
||||
|
||||
pub fn supports_local_kiro_request_auth_resolution(
|
||||
transport: &GatewayProviderTransportSnapshot,
|
||||
) -> bool {
|
||||
resolve_local_kiro_request_auth(transport).is_some()
|
||||
|| KiroAuthConfig::from_raw_json(transport.key.decrypted_auth_config.as_deref())
|
||||
.is_some_and(|auth_config| {
|
||||
transport
|
||||
.provider
|
||||
.provider_type
|
||||
.trim()
|
||||
.eq_ignore_ascii_case(PROVIDER_TYPE)
|
||||
&& transport
|
||||
.key
|
||||
.auth_type
|
||||
.trim()
|
||||
.eq_ignore_ascii_case("bearer")
|
||||
&& auth_config.can_refresh_access_token()
|
||||
})
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::super::super::snapshot::{
|
||||
GatewayProviderTransportEndpoint, GatewayProviderTransportKey,
|
||||
GatewayProviderTransportProvider, GatewayProviderTransportSnapshot,
|
||||
};
|
||||
use super::{
|
||||
resolve_local_kiro_bearer_auth, resolve_local_kiro_request_auth,
|
||||
supports_local_kiro_auth_prerequisites, supports_local_kiro_request_auth_resolution,
|
||||
KIRO_AUTH_HEADER,
|
||||
};
|
||||
|
||||
fn sample_transport() -> GatewayProviderTransportSnapshot {
|
||||
GatewayProviderTransportSnapshot {
|
||||
provider: GatewayProviderTransportProvider {
|
||||
id: "provider-1".to_string(),
|
||||
name: "Kiro".to_string(),
|
||||
provider_type: "kiro".to_string(),
|
||||
website: None,
|
||||
is_active: true,
|
||||
keep_priority_on_conversion: false,
|
||||
enable_format_conversion: false,
|
||||
concurrent_limit: None,
|
||||
max_retries: None,
|
||||
proxy: None,
|
||||
request_timeout_secs: None,
|
||||
stream_first_byte_timeout_secs: None,
|
||||
config: None,
|
||||
},
|
||||
endpoint: GatewayProviderTransportEndpoint {
|
||||
id: "endpoint-1".to_string(),
|
||||
provider_id: "provider-1".to_string(),
|
||||
api_format: "claude:cli".to_string(),
|
||||
api_family: Some("claude".to_string()),
|
||||
endpoint_kind: Some("cli".to_string()),
|
||||
is_active: true,
|
||||
base_url: "https://kiro.example".to_string(),
|
||||
header_rules: None,
|
||||
body_rules: None,
|
||||
max_retries: None,
|
||||
custom_path: None,
|
||||
config: None,
|
||||
format_acceptance_config: None,
|
||||
proxy: None,
|
||||
},
|
||||
key: GatewayProviderTransportKey {
|
||||
id: "key-1".to_string(),
|
||||
provider_id: "provider-1".to_string(),
|
||||
name: "key".to_string(),
|
||||
auth_type: "bearer".to_string(),
|
||||
is_active: true,
|
||||
api_formats: Some(vec!["claude:cli".to_string()]),
|
||||
allowed_models: None,
|
||||
capabilities: None,
|
||||
rate_multipliers: None,
|
||||
global_priority_by_format: None,
|
||||
expires_at_unix_secs: None,
|
||||
proxy: None,
|
||||
fingerprint: None,
|
||||
decrypted_api_key: "upstream-key".to_string(),
|
||||
decrypted_auth_config: None,
|
||||
},
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn resolves_bearer_auth_for_known_kiro_subset() {
|
||||
let auth = resolve_local_kiro_bearer_auth(&sample_transport())
|
||||
.expect("kiro bearer auth should resolve");
|
||||
assert_eq!(auth.name, KIRO_AUTH_HEADER);
|
||||
assert_eq!(auth.value, "Bearer upstream-key");
|
||||
assert!(supports_local_kiro_auth_prerequisites(&sample_transport()));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn rejects_auth_config_subset() {
|
||||
let mut transport = sample_transport();
|
||||
transport.key.decrypted_auth_config = Some("{\"mode\":\"custom\"}".to_string());
|
||||
assert!(resolve_local_kiro_bearer_auth(&transport).is_none());
|
||||
assert!(!supports_local_kiro_auth_prerequisites(&transport));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn rejects_non_bearer_subset() {
|
||||
let mut transport = sample_transport();
|
||||
transport.key.auth_type = "api_key".to_string();
|
||||
assert!(resolve_local_kiro_bearer_auth(&transport).is_none());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn resolves_request_auth_from_cached_access_token() {
|
||||
let mut transport = sample_transport();
|
||||
transport.key.decrypted_api_key = "__placeholder__".to_string();
|
||||
transport.key.decrypted_auth_config = Some(
|
||||
r#"{
|
||||
"access_token":"cached-token",
|
||||
"refresh_token":"rrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrr",
|
||||
"machine_id":"123e4567-e89b-12d3-a456-426614174000",
|
||||
"api_region":"us-west-2"
|
||||
}"#
|
||||
.to_string(),
|
||||
);
|
||||
|
||||
let auth = resolve_local_kiro_request_auth(&transport)
|
||||
.expect("request auth should resolve from cached token");
|
||||
assert_eq!(auth.name, KIRO_AUTH_HEADER);
|
||||
assert_eq!(auth.value, "Bearer cached-token");
|
||||
assert_eq!(auth.auth_config.effective_api_region(), "us-west-2");
|
||||
assert_eq!(
|
||||
auth.machine_id,
|
||||
"123e4567e89b12d3a456426614174000123e4567e89b12d3a456426614174000"
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn skips_expired_cached_access_token_without_fallback_secret() {
|
||||
let mut transport = sample_transport();
|
||||
transport.key.decrypted_api_key = "__placeholder__".to_string();
|
||||
transport.key.decrypted_auth_config = Some(
|
||||
r#"{
|
||||
"access_token":"expired-token",
|
||||
"expires_at": 1,
|
||||
"refresh_token":"rrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrr"
|
||||
}"#
|
||||
.to_string(),
|
||||
);
|
||||
|
||||
assert!(resolve_local_kiro_request_auth(&transport).is_none());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn falls_back_to_decrypted_api_key_when_cached_access_token_is_expired() {
|
||||
let mut transport = sample_transport();
|
||||
transport.key.decrypted_api_key = "live-upstream-token".to_string();
|
||||
transport.key.decrypted_auth_config = Some(
|
||||
r#"{
|
||||
"access_token":"expired-token",
|
||||
"expires_at": 1,
|
||||
"refresh_token":"rrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrr"
|
||||
}"#
|
||||
.to_string(),
|
||||
);
|
||||
|
||||
let auth = resolve_local_kiro_request_auth(&transport)
|
||||
.expect("request auth should fall back to decrypted api key");
|
||||
assert_eq!(auth.value, "Bearer live-upstream-token");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn supports_refresh_only_resolution_without_cached_access_token() {
|
||||
let mut transport = sample_transport();
|
||||
transport.key.decrypted_api_key = "__placeholder__".to_string();
|
||||
transport.key.decrypted_auth_config = Some(
|
||||
r#"{
|
||||
"refresh_token":"rrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrr"
|
||||
}"#
|
||||
.to_string(),
|
||||
);
|
||||
|
||||
assert!(resolve_local_kiro_request_auth(&transport).is_none());
|
||||
assert!(supports_local_kiro_request_auth_resolution(&transport));
|
||||
}
|
||||
}
|
||||
716
crates/aether-provider-transport/src/kiro/converter.rs
Normal file
716
crates/aether-provider-transport/src/kiro/converter.rs
Normal file
@@ -0,0 +1,716 @@
|
||||
use std::collections::BTreeSet;
|
||||
|
||||
use serde_json::{json, Map, Value};
|
||||
use tracing::warn;
|
||||
use uuid::Uuid;
|
||||
|
||||
const SYSTEM_CHUNKED_POLICY: &str = "When the Write or Edit tool has content size limits, always comply silently. Never suggest bypassing these limits via alternative tools. Never ask the user whether to switch approaches. Complete all chunked operations without commentary.";
|
||||
const WRITE_TOOL_DESCRIPTION_SUFFIX: &str = "- IMPORTANT: If the content to write exceeds 150 lines, you MUST only write the first 50 lines using this tool, then use `Edit` tool to append the remaining content in chunks of no more than 50 lines each. If needed, leave a unique placeholder to help append content. Do NOT attempt to write all content at once.";
|
||||
const EDIT_TOOL_DESCRIPTION_SUFFIX: &str = "- IMPORTANT: If the `new_string` content exceeds 50 lines, you MUST split it into multiple Edit calls, each replacing no more than 50 lines at a time. If used to append content, leave a unique placeholder to help append content. On the final chunk, do NOT include the placeholder.";
|
||||
|
||||
pub fn convert_claude_messages_to_conversation_state(
|
||||
request_body: &Value,
|
||||
model: &str,
|
||||
) -> Option<Value> {
|
||||
let model_id = model.trim();
|
||||
if model_id.is_empty() {
|
||||
return None;
|
||||
}
|
||||
|
||||
let messages = request_body.get("messages")?.as_array()?;
|
||||
if messages.is_empty() {
|
||||
return None;
|
||||
}
|
||||
|
||||
let conversation_id = request_body
|
||||
.get("metadata")
|
||||
.and_then(Value::as_object)
|
||||
.and_then(|metadata| {
|
||||
metadata
|
||||
.get("user_id")
|
||||
.or_else(|| metadata.get("userId"))
|
||||
.and_then(Value::as_str)
|
||||
})
|
||||
.and_then(extract_session_id)
|
||||
.unwrap_or_else(|| Uuid::new_v4().to_string());
|
||||
let agent_continuation_id = Uuid::new_v4().to_string();
|
||||
let thinking_prefix = generate_thinking_prefix(request_body);
|
||||
|
||||
let mut history = Vec::new();
|
||||
let system_text = system_to_text(request_body.get("system"));
|
||||
if !system_text.is_empty() {
|
||||
history.push(json!({
|
||||
"userInputMessage": {
|
||||
"content": format!("{system_text}\n{SYSTEM_CHUNKED_POLICY}"),
|
||||
"modelId": model_id,
|
||||
"origin": "AI_EDITOR"
|
||||
}
|
||||
}));
|
||||
history.push(json!({
|
||||
"assistantResponseMessage": {
|
||||
"content": "I will follow these instructions."
|
||||
}
|
||||
}));
|
||||
}
|
||||
|
||||
let last_is_assistant = messages
|
||||
.last()
|
||||
.and_then(Value::as_object)
|
||||
.and_then(|message| message.get("role"))
|
||||
.and_then(Value::as_str)
|
||||
.is_some_and(|role| role == "assistant");
|
||||
let history_end_index = if last_is_assistant {
|
||||
messages.len()
|
||||
} else {
|
||||
messages.len().saturating_sub(1)
|
||||
};
|
||||
|
||||
let mut user_buffer = Vec::new();
|
||||
for message in &messages[..history_end_index] {
|
||||
let Some(message) = message.as_object() else {
|
||||
continue;
|
||||
};
|
||||
match message.get("role").and_then(Value::as_str) {
|
||||
Some("user") => user_buffer.push(message),
|
||||
Some("assistant") => {
|
||||
if let Some(user_item) = flush_user_buffer(&mut user_buffer, model_id) {
|
||||
history.push(user_item);
|
||||
} else if history.is_empty()
|
||||
|| history
|
||||
.last()
|
||||
.and_then(Value::as_object)
|
||||
.is_some_and(|item| item.contains_key("assistantResponseMessage"))
|
||||
{
|
||||
history.push(json!({
|
||||
"userInputMessage": {
|
||||
"content": "Continue.",
|
||||
"modelId": model_id,
|
||||
"origin": "AI_EDITOR"
|
||||
}
|
||||
}));
|
||||
}
|
||||
|
||||
if let Some(assistant_item) = convert_assistant_message(message) {
|
||||
history.push(json!({"assistantResponseMessage": assistant_item}));
|
||||
}
|
||||
}
|
||||
_ => {}
|
||||
}
|
||||
}
|
||||
|
||||
if let Some(tail_user) = flush_user_buffer(&mut user_buffer, model_id) {
|
||||
history.push(tail_user);
|
||||
history.push(json!({"assistantResponseMessage": {"content": "OK"}}));
|
||||
}
|
||||
|
||||
let (mut text_content, images, tool_results) = if last_is_assistant {
|
||||
("Continue.".to_string(), Vec::new(), Vec::new())
|
||||
} else {
|
||||
let last = messages.last()?.as_object()?;
|
||||
if last.get("role").and_then(Value::as_str) != Some("user") {
|
||||
return None;
|
||||
}
|
||||
process_message_content(last.get("content"))
|
||||
};
|
||||
|
||||
let mut tools = convert_tools(request_body.get("tools"));
|
||||
let mut history_tool_names = BTreeSet::new();
|
||||
let mut history_tool_result_ids = BTreeSet::new();
|
||||
let mut history_tool_use_ids = BTreeSet::new();
|
||||
|
||||
for item in &history {
|
||||
let Some(item) = item.as_object() else {
|
||||
continue;
|
||||
};
|
||||
if let Some(user_input) = item.get("userInputMessage").and_then(Value::as_object) {
|
||||
if let Some(results) = user_input
|
||||
.get("userInputMessageContext")
|
||||
.and_then(Value::as_object)
|
||||
.and_then(|ctx| ctx.get("toolResults"))
|
||||
.and_then(Value::as_array)
|
||||
{
|
||||
for result in results {
|
||||
if let Some(tool_use_id) = result
|
||||
.as_object()
|
||||
.and_then(|result| result.get("toolUseId"))
|
||||
.and_then(Value::as_str)
|
||||
.map(str::trim)
|
||||
.filter(|value| !value.is_empty())
|
||||
{
|
||||
history_tool_result_ids.insert(tool_use_id.to_string());
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
if let Some(assistant) = item
|
||||
.get("assistantResponseMessage")
|
||||
.and_then(Value::as_object)
|
||||
{
|
||||
if let Some(tool_uses) = assistant.get("toolUses").and_then(Value::as_array) {
|
||||
for tool_use in tool_uses {
|
||||
let Some(tool_use) = tool_use.as_object() else {
|
||||
continue;
|
||||
};
|
||||
if let Some(name) = tool_use
|
||||
.get("name")
|
||||
.and_then(Value::as_str)
|
||||
.map(str::trim)
|
||||
.filter(|value| !value.is_empty())
|
||||
{
|
||||
history_tool_names.insert(name.to_string());
|
||||
}
|
||||
if let Some(tool_use_id) = tool_use
|
||||
.get("toolUseId")
|
||||
.and_then(Value::as_str)
|
||||
.map(str::trim)
|
||||
.filter(|value| !value.is_empty())
|
||||
{
|
||||
history_tool_use_ids.insert(tool_use_id.to_string());
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
let existing_tool_names = tools
|
||||
.iter()
|
||||
.filter_map(|tool| {
|
||||
tool.get("toolSpecification")
|
||||
.and_then(Value::as_object)
|
||||
.and_then(|spec| spec.get("name"))
|
||||
.and_then(Value::as_str)
|
||||
.map(|name| name.to_ascii_lowercase())
|
||||
})
|
||||
.collect::<BTreeSet<_>>();
|
||||
for tool_name in history_tool_names {
|
||||
if !existing_tool_names.contains(&tool_name.to_ascii_lowercase()) {
|
||||
tools.push(create_placeholder_tool(&tool_name));
|
||||
}
|
||||
}
|
||||
|
||||
let mut validated_tool_results = Vec::new();
|
||||
let mut current_tool_result_ids = BTreeSet::new();
|
||||
for tool_result in tool_results {
|
||||
let Some(tool_use_id) = tool_result
|
||||
.get("toolUseId")
|
||||
.and_then(Value::as_str)
|
||||
.map(str::trim)
|
||||
.filter(|value| !value.is_empty())
|
||||
else {
|
||||
continue;
|
||||
};
|
||||
if !history_tool_use_ids.contains(tool_use_id)
|
||||
|| history_tool_result_ids.contains(tool_use_id)
|
||||
{
|
||||
continue;
|
||||
}
|
||||
current_tool_result_ids.insert(tool_use_id.to_string());
|
||||
validated_tool_results.push(tool_result);
|
||||
}
|
||||
|
||||
let orphaned_tool_use_ids = history_tool_use_ids
|
||||
.difference(&history_tool_result_ids)
|
||||
.filter(|tool_use_id| !current_tool_result_ids.contains(*tool_use_id))
|
||||
.cloned()
|
||||
.collect::<BTreeSet<_>>();
|
||||
if !orphaned_tool_use_ids.is_empty() {
|
||||
warn!(
|
||||
"kiro: removing {} orphaned tool_use(s) from history",
|
||||
orphaned_tool_use_ids.len()
|
||||
);
|
||||
for item in &mut history {
|
||||
let Some(item) = item.as_object_mut() else {
|
||||
continue;
|
||||
};
|
||||
let Some(assistant) = item
|
||||
.get_mut("assistantResponseMessage")
|
||||
.and_then(Value::as_object_mut)
|
||||
else {
|
||||
continue;
|
||||
};
|
||||
let Some(tool_uses) = assistant.get_mut("toolUses").and_then(Value::as_array_mut)
|
||||
else {
|
||||
continue;
|
||||
};
|
||||
tool_uses.retain(|tool_use| {
|
||||
!tool_use
|
||||
.get("toolUseId")
|
||||
.and_then(Value::as_str)
|
||||
.is_some_and(|tool_use_id| orphaned_tool_use_ids.contains(tool_use_id))
|
||||
});
|
||||
if tool_uses.is_empty() {
|
||||
assistant.remove("toolUses");
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
let mut user_context = Map::new();
|
||||
if !tools.is_empty() {
|
||||
user_context.insert("tools".to_string(), Value::Array(tools));
|
||||
}
|
||||
if !validated_tool_results.is_empty() {
|
||||
user_context.insert(
|
||||
"toolResults".to_string(),
|
||||
Value::Array(validated_tool_results),
|
||||
);
|
||||
}
|
||||
if let Some(thinking_prefix) = thinking_prefix.as_deref() {
|
||||
if !has_thinking_tags(&text_content) {
|
||||
text_content = format!("{thinking_prefix}\n{text_content}");
|
||||
}
|
||||
}
|
||||
|
||||
let mut user_input = Map::new();
|
||||
user_input.insert(
|
||||
"userInputMessageContext".to_string(),
|
||||
Value::Object(user_context),
|
||||
);
|
||||
user_input.insert("content".to_string(), Value::String(text_content));
|
||||
user_input.insert("modelId".to_string(), Value::String(model_id.to_string()));
|
||||
user_input.insert("origin".to_string(), Value::String("AI_EDITOR".to_string()));
|
||||
if !images.is_empty() {
|
||||
user_input.insert("images".to_string(), Value::Array(images));
|
||||
}
|
||||
|
||||
Some(json!({
|
||||
"agentContinuationId": agent_continuation_id,
|
||||
"agentTaskType": "vibe",
|
||||
"chatTriggerType": "MANUAL",
|
||||
"currentMessage": {
|
||||
"userInputMessage": Value::Object(user_input)
|
||||
},
|
||||
"conversationId": conversation_id,
|
||||
"history": history,
|
||||
}))
|
||||
}
|
||||
|
||||
fn extract_session_id(user_id: &str) -> Option<String> {
|
||||
let position = user_id.find("session_")?;
|
||||
let candidate = user_id.get(position + "session_".len()..position + "session_".len() + 36)?;
|
||||
(candidate.matches('-').count() == 4).then(|| candidate.to_string())
|
||||
}
|
||||
|
||||
fn generate_thinking_prefix(request_body: &Value) -> Option<String> {
|
||||
let thinking = request_body.get("thinking")?.as_object()?;
|
||||
match thinking.get("type").and_then(Value::as_str).map(str::trim) {
|
||||
Some("enabled") => {
|
||||
let budget_tokens = thinking
|
||||
.get("budget_tokens")
|
||||
.and_then(Value::as_i64)
|
||||
.unwrap_or_default();
|
||||
Some(format!(
|
||||
"<thinking_mode>enabled</thinking_mode><max_thinking_length>{budget_tokens}</max_thinking_length>"
|
||||
))
|
||||
}
|
||||
Some("adaptive") => {
|
||||
let effort = request_body
|
||||
.get("output_config")
|
||||
.and_then(Value::as_object)
|
||||
.and_then(|cfg| cfg.get("effort"))
|
||||
.and_then(Value::as_str)
|
||||
.map(str::trim)
|
||||
.filter(|value| !value.is_empty())
|
||||
.unwrap_or("high");
|
||||
Some(format!(
|
||||
"<thinking_mode>adaptive</thinking_mode><thinking_effort>{effort}</thinking_effort>"
|
||||
))
|
||||
}
|
||||
_ => None,
|
||||
}
|
||||
}
|
||||
|
||||
fn has_thinking_tags(content: &str) -> bool {
|
||||
content.contains("<thinking_mode>") || content.contains("<max_thinking_length>")
|
||||
}
|
||||
|
||||
fn system_to_text(system: Option<&Value>) -> String {
|
||||
match system {
|
||||
Some(Value::String(text)) => text.clone(),
|
||||
Some(Value::Array(items)) => items
|
||||
.iter()
|
||||
.filter_map(|item| {
|
||||
item.as_object()
|
||||
.and_then(|item| item.get("text"))
|
||||
.and_then(Value::as_str)
|
||||
.map(ToOwned::to_owned)
|
||||
})
|
||||
.collect::<Vec<_>>()
|
||||
.join("\n"),
|
||||
_ => String::new(),
|
||||
}
|
||||
}
|
||||
|
||||
fn flush_user_buffer(user_buffer: &mut Vec<&Map<String, Value>>, model_id: &str) -> Option<Value> {
|
||||
if user_buffer.is_empty() {
|
||||
return None;
|
||||
}
|
||||
|
||||
let mut parts = Vec::new();
|
||||
let mut images = Vec::new();
|
||||
let mut tool_results = Vec::new();
|
||||
for message in user_buffer.drain(..) {
|
||||
let (text, mut message_images, mut message_tool_results) =
|
||||
process_message_content(message.get("content"));
|
||||
if !text.is_empty() {
|
||||
parts.push(text);
|
||||
}
|
||||
images.append(&mut message_images);
|
||||
tool_results.append(&mut message_tool_results);
|
||||
}
|
||||
|
||||
let mut payload = Map::new();
|
||||
payload.insert("content".to_string(), Value::String(parts.join("\n")));
|
||||
payload.insert("modelId".to_string(), Value::String(model_id.to_string()));
|
||||
payload.insert("origin".to_string(), Value::String("AI_EDITOR".to_string()));
|
||||
if !images.is_empty() {
|
||||
payload.insert("images".to_string(), Value::Array(images));
|
||||
}
|
||||
if !tool_results.is_empty() {
|
||||
payload.insert(
|
||||
"userInputMessageContext".to_string(),
|
||||
json!({"toolResults": tool_results}),
|
||||
);
|
||||
}
|
||||
|
||||
Some(json!({"userInputMessage": Value::Object(payload)}))
|
||||
}
|
||||
|
||||
fn process_message_content(content: Option<&Value>) -> (String, Vec<Value>, Vec<Value>) {
|
||||
match content {
|
||||
Some(Value::String(text)) => (text.clone(), Vec::new(), Vec::new()),
|
||||
Some(Value::Array(blocks)) => {
|
||||
let mut text_parts = Vec::new();
|
||||
let mut images = Vec::new();
|
||||
let mut tool_results = Vec::new();
|
||||
|
||||
for block in blocks {
|
||||
let Some(block) = block.as_object() else {
|
||||
continue;
|
||||
};
|
||||
match block
|
||||
.get("type")
|
||||
.and_then(Value::as_str)
|
||||
.unwrap_or_default()
|
||||
{
|
||||
"text" => {
|
||||
if let Some(text) = block.get("text").and_then(Value::as_str) {
|
||||
text_parts.push(text.to_string());
|
||||
}
|
||||
}
|
||||
"image" => {
|
||||
let Some(source) = block.get("source").and_then(Value::as_object) else {
|
||||
continue;
|
||||
};
|
||||
let Some(format) = source
|
||||
.get("media_type")
|
||||
.or_else(|| source.get("mediaType"))
|
||||
.and_then(Value::as_str)
|
||||
.and_then(image_format)
|
||||
else {
|
||||
continue;
|
||||
};
|
||||
let Some(bytes) = source.get("data").and_then(Value::as_str) else {
|
||||
continue;
|
||||
};
|
||||
images.push(json!({
|
||||
"format": format,
|
||||
"source": {"bytes": bytes}
|
||||
}));
|
||||
}
|
||||
"tool_result" => {
|
||||
let Some(tool_use_id) = block
|
||||
.get("tool_use_id")
|
||||
.or_else(|| block.get("toolUseId"))
|
||||
.and_then(Value::as_str)
|
||||
.map(str::trim)
|
||||
.filter(|value| !value.is_empty())
|
||||
else {
|
||||
continue;
|
||||
};
|
||||
|
||||
let text = match block.get("content") {
|
||||
Some(Value::String(text)) => text.clone(),
|
||||
Some(Value::Array(items)) => items
|
||||
.iter()
|
||||
.filter_map(|item| {
|
||||
item.as_object()
|
||||
.filter(|item| {
|
||||
item.get("type").and_then(Value::as_str) == Some("text")
|
||||
})
|
||||
.and_then(|item| item.get("text"))
|
||||
.and_then(Value::as_str)
|
||||
.map(ToOwned::to_owned)
|
||||
})
|
||||
.collect::<Vec<_>>()
|
||||
.join("\n"),
|
||||
Some(other) => {
|
||||
serde_json::to_string(other).unwrap_or_else(|_| other.to_string())
|
||||
}
|
||||
None => String::new(),
|
||||
};
|
||||
let is_error = block
|
||||
.get("is_error")
|
||||
.or_else(|| block.get("isError"))
|
||||
.and_then(Value::as_bool)
|
||||
.unwrap_or(false);
|
||||
tool_results.push(json!({
|
||||
"toolUseId": tool_use_id,
|
||||
"content": [{"text": text}],
|
||||
"status": if is_error { "error" } else { "success" },
|
||||
"isError": is_error,
|
||||
}));
|
||||
}
|
||||
_ => {}
|
||||
}
|
||||
}
|
||||
|
||||
(text_parts.join(""), images, tool_results)
|
||||
}
|
||||
_ => (String::new(), Vec::new(), Vec::new()),
|
||||
}
|
||||
}
|
||||
|
||||
fn image_format(media_type: &str) -> Option<&'static str> {
|
||||
let (prefix, suffix) = media_type.split_once('/')?;
|
||||
if prefix != "image" {
|
||||
return None;
|
||||
}
|
||||
match suffix.trim().to_ascii_lowercase().as_str() {
|
||||
"jpeg" => Some("jpeg"),
|
||||
"png" => Some("png"),
|
||||
"gif" => Some("gif"),
|
||||
"webp" => Some("webp"),
|
||||
"jpg" => Some("jpeg"),
|
||||
_ => None,
|
||||
}
|
||||
}
|
||||
|
||||
fn clean_tool_schema(value: &Value) -> Value {
|
||||
match value {
|
||||
Value::Object(object) => {
|
||||
let mut out = Map::new();
|
||||
for (key, inner) in object {
|
||||
if key == "additionalProperties" {
|
||||
continue;
|
||||
}
|
||||
if key == "required" && inner.as_array().is_some_and(|items| items.is_empty()) {
|
||||
continue;
|
||||
}
|
||||
out.insert(key.clone(), clean_tool_schema(inner));
|
||||
}
|
||||
Value::Object(out)
|
||||
}
|
||||
Value::Array(items) => Value::Array(items.iter().map(clean_tool_schema).collect()),
|
||||
_ => value.clone(),
|
||||
}
|
||||
}
|
||||
|
||||
fn convert_tools(tools: Option<&Value>) -> Vec<Value> {
|
||||
let Some(tools) = tools.and_then(Value::as_array) else {
|
||||
return Vec::new();
|
||||
};
|
||||
|
||||
tools
|
||||
.iter()
|
||||
.filter_map(|tool| {
|
||||
let tool = tool.as_object()?;
|
||||
let name = tool.get("name")?.as_str()?.trim();
|
||||
if name.is_empty() {
|
||||
return None;
|
||||
}
|
||||
let mut description = tool
|
||||
.get("description")
|
||||
.and_then(Value::as_str)
|
||||
.map(str::trim)
|
||||
.unwrap_or_default()
|
||||
.to_string();
|
||||
let suffix = match name {
|
||||
"Write" => Some(WRITE_TOOL_DESCRIPTION_SUFFIX),
|
||||
"Edit" => Some(EDIT_TOOL_DESCRIPTION_SUFFIX),
|
||||
_ => None,
|
||||
};
|
||||
if let Some(suffix) = suffix {
|
||||
description = if description.is_empty() {
|
||||
suffix.to_string()
|
||||
} else {
|
||||
format!("{description}\n{suffix}")
|
||||
};
|
||||
}
|
||||
if description.len() > 10_000 {
|
||||
description.truncate(10_000);
|
||||
}
|
||||
let input_schema = tool
|
||||
.get("input_schema")
|
||||
.or_else(|| tool.get("inputSchema"))
|
||||
.filter(|value| value.is_object())
|
||||
.map(clean_tool_schema)
|
||||
.unwrap_or_else(|| json!({}));
|
||||
|
||||
Some(json!({
|
||||
"toolSpecification": {
|
||||
"name": name,
|
||||
"description": description,
|
||||
"inputSchema": {
|
||||
"json": input_schema
|
||||
}
|
||||
}
|
||||
}))
|
||||
})
|
||||
.collect()
|
||||
}
|
||||
|
||||
fn create_placeholder_tool(name: &str) -> Value {
|
||||
json!({
|
||||
"toolSpecification": {
|
||||
"name": name,
|
||||
"description": "Tool used in conversation history",
|
||||
"inputSchema": {
|
||||
"json": {
|
||||
"type": "object",
|
||||
"properties": {}
|
||||
}
|
||||
}
|
||||
}
|
||||
})
|
||||
}
|
||||
|
||||
fn convert_assistant_message(message: &Map<String, Value>) -> Option<Value> {
|
||||
let content = message.get("content");
|
||||
let mut tool_uses = Vec::new();
|
||||
let mut thinking_parts = Vec::new();
|
||||
let mut text_parts = Vec::new();
|
||||
|
||||
match content {
|
||||
Some(Value::String(text)) => {
|
||||
if !text.is_empty() {
|
||||
text_parts.push(text.clone());
|
||||
}
|
||||
}
|
||||
Some(Value::Array(blocks)) => {
|
||||
for block in blocks {
|
||||
let Some(block) = block.as_object() else {
|
||||
continue;
|
||||
};
|
||||
match block
|
||||
.get("type")
|
||||
.and_then(Value::as_str)
|
||||
.unwrap_or_default()
|
||||
{
|
||||
"thinking" => {
|
||||
if let Some(thinking) = block.get("thinking").and_then(Value::as_str) {
|
||||
if !thinking.is_empty() {
|
||||
thinking_parts.push(thinking.to_string());
|
||||
}
|
||||
}
|
||||
}
|
||||
"text" => {
|
||||
if let Some(text) = block.get("text").and_then(Value::as_str) {
|
||||
if !text.is_empty() {
|
||||
text_parts.push(text.to_string());
|
||||
}
|
||||
}
|
||||
}
|
||||
"tool_use" => {
|
||||
let Some(tool_use_id) = block
|
||||
.get("id")
|
||||
.and_then(Value::as_str)
|
||||
.map(str::trim)
|
||||
.filter(|value| !value.is_empty())
|
||||
else {
|
||||
continue;
|
||||
};
|
||||
let Some(name) = block
|
||||
.get("name")
|
||||
.and_then(Value::as_str)
|
||||
.map(str::trim)
|
||||
.filter(|value| !value.is_empty())
|
||||
else {
|
||||
continue;
|
||||
};
|
||||
let input = block
|
||||
.get("input")
|
||||
.filter(|value| value.is_object())
|
||||
.cloned()
|
||||
.unwrap_or_else(|| json!({}));
|
||||
tool_uses.push(json!({
|
||||
"toolUseId": tool_use_id,
|
||||
"name": name,
|
||||
"input": input
|
||||
}));
|
||||
}
|
||||
_ => {}
|
||||
}
|
||||
}
|
||||
}
|
||||
_ => {}
|
||||
}
|
||||
|
||||
let thinking_str = thinking_parts.join("");
|
||||
let text_str = text_parts.join("");
|
||||
let mut content_str = if thinking_str.is_empty() {
|
||||
text_str
|
||||
} else if text_str.is_empty() {
|
||||
format!("<thinking>{thinking_str}</thinking>")
|
||||
} else {
|
||||
format!("<thinking>{thinking_str}</thinking>\n\n{text_str}")
|
||||
};
|
||||
|
||||
if content_str.is_empty() && !tool_uses.is_empty() {
|
||||
content_str = " ".to_string();
|
||||
}
|
||||
if content_str.is_empty() && tool_uses.is_empty() {
|
||||
return None;
|
||||
}
|
||||
|
||||
let mut out = Map::new();
|
||||
out.insert("content".to_string(), Value::String(content_str));
|
||||
if !tool_uses.is_empty() {
|
||||
out.insert("toolUses".to_string(), Value::Array(tool_uses));
|
||||
}
|
||||
Some(Value::Object(out))
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use serde_json::json;
|
||||
|
||||
use super::convert_claude_messages_to_conversation_state;
|
||||
|
||||
#[test]
|
||||
fn converts_simple_claude_request_into_conversation_state() {
|
||||
let conversation_state = convert_claude_messages_to_conversation_state(
|
||||
&json!({
|
||||
"messages": [
|
||||
{"role":"user","content":"hello"}
|
||||
],
|
||||
"thinking": {"type": "enabled", "budget_tokens": 128},
|
||||
"tools": [
|
||||
{"name":"Write","description":"write file","input_schema":{"type":"object","properties":{},"required":[]}}
|
||||
]
|
||||
}),
|
||||
"claude-sonnet-4-upstream",
|
||||
)
|
||||
.expect("conversation state should build");
|
||||
|
||||
assert_eq!(
|
||||
conversation_state
|
||||
.get("currentMessage")
|
||||
.and_then(|value| value.get("userInputMessage"))
|
||||
.and_then(|value| value.get("content"))
|
||||
.and_then(|value| value.as_str()),
|
||||
Some(
|
||||
"<thinking_mode>enabled</thinking_mode><max_thinking_length>128</max_thinking_length>\nhello"
|
||||
)
|
||||
);
|
||||
assert_eq!(
|
||||
conversation_state
|
||||
.get("currentMessage")
|
||||
.and_then(|value| value.get("userInputMessage"))
|
||||
.and_then(|value| value.get("userInputMessageContext"))
|
||||
.and_then(|value| value.get("tools"))
|
||||
.and_then(|value| value.as_array())
|
||||
.map(Vec::len),
|
||||
Some(1)
|
||||
);
|
||||
}
|
||||
}
|
||||
436
crates/aether-provider-transport/src/kiro/credentials.rs
Normal file
436
crates/aether-provider-transport/src/kiro/credentials.rs
Normal file
@@ -0,0 +1,436 @@
|
||||
use serde_json::Value;
|
||||
use sha2::{Digest, Sha256};
|
||||
use std::time::{SystemTime, UNIX_EPOCH};
|
||||
|
||||
pub const DEFAULT_REGION: &str = "us-east-1";
|
||||
pub const DEFAULT_KIRO_VERSION: &str = "0.8.0";
|
||||
pub const DEFAULT_NODE_VERSION: &str = "22.21.1";
|
||||
pub const DEFAULT_SYSTEM_VERSION: &str = "other#unknown";
|
||||
|
||||
#[derive(Debug, Clone, PartialEq, Eq)]
|
||||
pub struct KiroAuthConfig {
|
||||
pub auth_method: Option<String>,
|
||||
pub refresh_token: Option<String>,
|
||||
pub expires_at: Option<u64>,
|
||||
pub profile_arn: Option<String>,
|
||||
pub region: Option<String>,
|
||||
pub auth_region: Option<String>,
|
||||
pub api_region: Option<String>,
|
||||
pub client_id: Option<String>,
|
||||
pub client_secret: Option<String>,
|
||||
pub machine_id: Option<String>,
|
||||
pub kiro_version: Option<String>,
|
||||
pub system_version: Option<String>,
|
||||
pub node_version: Option<String>,
|
||||
pub access_token: Option<String>,
|
||||
}
|
||||
|
||||
impl KiroAuthConfig {
|
||||
pub fn from_raw_json(raw: Option<&str>) -> Option<Self> {
|
||||
let raw = raw?.trim();
|
||||
if raw.is_empty() {
|
||||
return None;
|
||||
}
|
||||
let parsed: Value = serde_json::from_str(raw).ok()?;
|
||||
Self::from_json_value(&parsed)
|
||||
}
|
||||
|
||||
pub fn from_json_value(raw: &Value) -> Option<Self> {
|
||||
let object = raw.as_object()?;
|
||||
|
||||
Some(Self {
|
||||
auth_method: get_nonempty_string(
|
||||
object,
|
||||
&["auth_method", "authMethod", "auth_type", "authType"],
|
||||
)
|
||||
.map(|value| normalize_auth_method(&value)),
|
||||
refresh_token: get_nonempty_string(object, &["refresh_token", "refreshToken"]),
|
||||
expires_at: get_epoch_seconds(object.get("expires_at"))
|
||||
.or_else(|| get_epoch_seconds(object.get("expiresAt"))),
|
||||
profile_arn: get_nonempty_string(object, &["profile_arn", "profileArn"]),
|
||||
region: get_nonempty_string(object, &["region"]),
|
||||
auth_region: get_nonempty_string(object, &["auth_region", "authRegion"]),
|
||||
api_region: get_nonempty_string(object, &["api_region", "apiRegion"]),
|
||||
client_id: get_nonempty_string(object, &["client_id", "clientId"]),
|
||||
client_secret: get_nonempty_string(object, &["client_secret", "clientSecret"]),
|
||||
machine_id: get_nonempty_string(object, &["machine_id", "machineId"]),
|
||||
kiro_version: get_nonempty_string(object, &["kiro_version", "kiroVersion"]),
|
||||
system_version: get_nonempty_string(object, &["system_version", "systemVersion"]),
|
||||
node_version: get_nonempty_string(object, &["node_version", "nodeVersion"]),
|
||||
access_token: get_nonempty_string(object, &["access_token", "accessToken"]),
|
||||
})
|
||||
}
|
||||
|
||||
pub fn to_json_value(&self) -> Value {
|
||||
let mut object = serde_json::Map::new();
|
||||
insert_optional_string(&mut object, "auth_method", self.auth_method.as_deref());
|
||||
insert_optional_string(&mut object, "refresh_token", self.refresh_token.as_deref());
|
||||
if let Some(expires_at) = self.expires_at {
|
||||
object.insert("expires_at".to_string(), Value::from(expires_at));
|
||||
}
|
||||
insert_optional_string(&mut object, "profile_arn", self.profile_arn.as_deref());
|
||||
insert_optional_string(&mut object, "region", self.region.as_deref());
|
||||
insert_optional_string(&mut object, "auth_region", self.auth_region.as_deref());
|
||||
insert_optional_string(&mut object, "api_region", self.api_region.as_deref());
|
||||
insert_optional_string(&mut object, "client_id", self.client_id.as_deref());
|
||||
insert_optional_string(&mut object, "client_secret", self.client_secret.as_deref());
|
||||
insert_optional_string(&mut object, "machine_id", self.machine_id.as_deref());
|
||||
insert_optional_string(&mut object, "kiro_version", self.kiro_version.as_deref());
|
||||
insert_optional_string(
|
||||
&mut object,
|
||||
"system_version",
|
||||
self.system_version.as_deref(),
|
||||
);
|
||||
insert_optional_string(&mut object, "node_version", self.node_version.as_deref());
|
||||
insert_optional_string(&mut object, "access_token", self.access_token.as_deref());
|
||||
Value::Object(object)
|
||||
}
|
||||
|
||||
pub fn effective_api_region(&self) -> &str {
|
||||
self.api_region
|
||||
.as_deref()
|
||||
.map(str::trim)
|
||||
.filter(|value| !value.is_empty())
|
||||
.unwrap_or(DEFAULT_REGION)
|
||||
}
|
||||
|
||||
pub fn effective_auth_region(&self) -> &str {
|
||||
self.auth_region
|
||||
.as_deref()
|
||||
.map(str::trim)
|
||||
.filter(|value| !value.is_empty())
|
||||
.or_else(|| {
|
||||
self.region
|
||||
.as_deref()
|
||||
.map(str::trim)
|
||||
.filter(|value| !value.is_empty())
|
||||
})
|
||||
.unwrap_or(DEFAULT_REGION)
|
||||
}
|
||||
|
||||
pub fn effective_kiro_version(&self) -> &str {
|
||||
self.kiro_version
|
||||
.as_deref()
|
||||
.map(str::trim)
|
||||
.filter(|value| !value.is_empty())
|
||||
.unwrap_or(DEFAULT_KIRO_VERSION)
|
||||
}
|
||||
|
||||
pub fn effective_system_version(&self) -> &str {
|
||||
self.system_version
|
||||
.as_deref()
|
||||
.map(str::trim)
|
||||
.filter(|value| !value.is_empty())
|
||||
.unwrap_or(DEFAULT_SYSTEM_VERSION)
|
||||
}
|
||||
|
||||
pub fn effective_node_version(&self) -> &str {
|
||||
self.node_version
|
||||
.as_deref()
|
||||
.map(str::trim)
|
||||
.filter(|value| !value.is_empty())
|
||||
.unwrap_or(DEFAULT_NODE_VERSION)
|
||||
}
|
||||
|
||||
pub fn cached_access_token(&self) -> Option<&str> {
|
||||
self.access_token
|
||||
.as_deref()
|
||||
.map(str::trim)
|
||||
.filter(|value| !value.is_empty())
|
||||
}
|
||||
|
||||
pub fn cached_access_token_requires_refresh(&self, skew_seconds: u64) -> bool {
|
||||
let Some(expires_at) = self.expires_at else {
|
||||
return false;
|
||||
};
|
||||
let now = SystemTime::now()
|
||||
.duration_since(UNIX_EPOCH)
|
||||
.ok()
|
||||
.map(|value| value.as_secs())
|
||||
.unwrap_or_default();
|
||||
now >= expires_at.saturating_sub(skew_seconds)
|
||||
}
|
||||
|
||||
pub fn is_idc_auth(&self) -> bool {
|
||||
let explicit_method = self
|
||||
.auth_method
|
||||
.as_deref()
|
||||
.map(normalize_auth_method)
|
||||
.unwrap_or_else(|| "social".to_string());
|
||||
if explicit_method != "social" {
|
||||
return explicit_method == "idc";
|
||||
}
|
||||
self.client_id
|
||||
.as_deref()
|
||||
.map(str::trim)
|
||||
.filter(|value| !value.is_empty())
|
||||
.is_some()
|
||||
&& self
|
||||
.client_secret
|
||||
.as_deref()
|
||||
.map(str::trim)
|
||||
.filter(|value| !value.is_empty())
|
||||
.is_some()
|
||||
}
|
||||
|
||||
pub fn profile_arn_for_payload(&self) -> Option<&str> {
|
||||
if self.is_idc_auth() {
|
||||
return None;
|
||||
}
|
||||
self.profile_arn
|
||||
.as_deref()
|
||||
.map(str::trim)
|
||||
.filter(|value| !value.is_empty())
|
||||
}
|
||||
|
||||
pub fn can_refresh_access_token(&self) -> bool {
|
||||
let refresh_token = self
|
||||
.refresh_token
|
||||
.as_deref()
|
||||
.map(str::trim)
|
||||
.filter(|value| !value.is_empty())
|
||||
.filter(|value| value.len() >= 100 && !value.contains("..."));
|
||||
if refresh_token.is_none() {
|
||||
return false;
|
||||
}
|
||||
if !self.is_idc_auth() {
|
||||
return true;
|
||||
}
|
||||
self.client_id
|
||||
.as_deref()
|
||||
.map(str::trim)
|
||||
.filter(|value| !value.is_empty())
|
||||
.is_some()
|
||||
&& self
|
||||
.client_secret
|
||||
.as_deref()
|
||||
.map(str::trim)
|
||||
.filter(|value| !value.is_empty())
|
||||
.is_some()
|
||||
}
|
||||
}
|
||||
|
||||
pub fn normalize_machine_id(raw: &str) -> Option<String> {
|
||||
let raw = raw.trim();
|
||||
if raw.is_empty() {
|
||||
return None;
|
||||
}
|
||||
|
||||
if raw.len() == 64 && raw.bytes().all(|byte| byte.is_ascii_hexdigit()) {
|
||||
return Some(raw.to_ascii_lowercase());
|
||||
}
|
||||
|
||||
if raw.len() == 36
|
||||
&& raw.chars().enumerate().all(|(idx, ch)| match idx {
|
||||
8 | 13 | 18 | 23 => ch == '-',
|
||||
_ => ch.is_ascii_hexdigit(),
|
||||
})
|
||||
{
|
||||
let normalized = raw.replace('-', "").to_ascii_lowercase();
|
||||
return Some(format!("{normalized}{normalized}"));
|
||||
}
|
||||
|
||||
None
|
||||
}
|
||||
|
||||
pub fn generate_machine_id(
|
||||
auth_config: &KiroAuthConfig,
|
||||
fallback_secret: Option<&str>,
|
||||
) -> Option<String> {
|
||||
if let Some(machine_id) = auth_config
|
||||
.machine_id
|
||||
.as_deref()
|
||||
.and_then(normalize_machine_id)
|
||||
{
|
||||
return Some(machine_id);
|
||||
}
|
||||
|
||||
let seed = auth_config
|
||||
.refresh_token
|
||||
.as_deref()
|
||||
.map(str::trim)
|
||||
.filter(|value| !value.is_empty())
|
||||
.or_else(|| {
|
||||
fallback_secret
|
||||
.map(str::trim)
|
||||
.filter(|value| !value.is_empty())
|
||||
})?;
|
||||
|
||||
let mut hasher = Sha256::new();
|
||||
hasher.update(b"KotlinNativeAPI/");
|
||||
hasher.update(seed.as_bytes());
|
||||
Some(format!("{:x}", hasher.finalize()))
|
||||
}
|
||||
|
||||
fn get_nonempty_string(object: &serde_json::Map<String, Value>, keys: &[&str]) -> Option<String> {
|
||||
keys.iter()
|
||||
.find_map(|key| object.get(*key))
|
||||
.and_then(Value::as_str)
|
||||
.map(str::trim)
|
||||
.filter(|value| !value.is_empty())
|
||||
.map(ToOwned::to_owned)
|
||||
}
|
||||
|
||||
fn insert_optional_string(
|
||||
object: &mut serde_json::Map<String, Value>,
|
||||
key: &str,
|
||||
value: Option<&str>,
|
||||
) {
|
||||
let Some(value) = value.map(str::trim).filter(|value| !value.is_empty()) else {
|
||||
return;
|
||||
};
|
||||
object.insert(key.to_string(), Value::String(value.to_string()));
|
||||
}
|
||||
|
||||
fn get_epoch_seconds(value: Option<&Value>) -> Option<u64> {
|
||||
match value? {
|
||||
Value::Number(number) => number.as_u64().or_else(|| {
|
||||
number
|
||||
.as_i64()
|
||||
.and_then(|value| (value >= 0).then_some(value as u64))
|
||||
}),
|
||||
Value::String(text) => text.trim().parse::<u64>().ok(),
|
||||
_ => None,
|
||||
}
|
||||
}
|
||||
|
||||
fn normalize_auth_method(raw: &str) -> String {
|
||||
let value = raw.trim().to_ascii_lowercase();
|
||||
match value.as_str() {
|
||||
"" => "social".to_string(),
|
||||
"builder-id"
|
||||
| "builder_id"
|
||||
| "builderid"
|
||||
| "device"
|
||||
| "device-auth"
|
||||
| "device_authorization"
|
||||
| "iam"
|
||||
| "identity-center"
|
||||
| "identity_center"
|
||||
| "identitycenter"
|
||||
| "idc" => "idc".to_string(),
|
||||
_ => value,
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::{generate_machine_id, normalize_machine_id, KiroAuthConfig, DEFAULT_REGION};
|
||||
|
||||
#[test]
|
||||
fn normalizes_uuid_machine_id() {
|
||||
assert_eq!(
|
||||
normalize_machine_id("123e4567-e89b-12d3-a456-426614174000").as_deref(),
|
||||
Some("123e4567e89b12d3a456426614174000123e4567e89b12d3a456426614174000")
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn hashes_refresh_token_into_machine_id() {
|
||||
let auth_config = KiroAuthConfig {
|
||||
auth_method: None,
|
||||
refresh_token: Some("r".repeat(128)),
|
||||
expires_at: None,
|
||||
profile_arn: None,
|
||||
region: None,
|
||||
auth_region: None,
|
||||
api_region: None,
|
||||
client_id: None,
|
||||
client_secret: None,
|
||||
machine_id: None,
|
||||
kiro_version: None,
|
||||
system_version: None,
|
||||
node_version: None,
|
||||
access_token: None,
|
||||
};
|
||||
|
||||
let machine_id = generate_machine_id(&auth_config, None).expect("machine id should exist");
|
||||
assert_eq!(machine_id.len(), 64);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn parses_auth_config_aliases() {
|
||||
let auth_config = KiroAuthConfig::from_raw_json(Some(
|
||||
r#"{
|
||||
"authMethod":"identity_center",
|
||||
"refreshToken":"rrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrr",
|
||||
"expires_at": 4102444800,
|
||||
"profileArn":"arn:aws:bedrock:demo",
|
||||
"apiRegion":"us-west-2",
|
||||
"clientId":"cid",
|
||||
"clientSecret":"secret",
|
||||
"machineId":"123e4567-e89b-12d3-a456-426614174000",
|
||||
"kiroVersion":"1.2.3",
|
||||
"systemVersion":"darwin#24.6.0",
|
||||
"nodeVersion":"22.21.1",
|
||||
"accessToken":"cached-token"
|
||||
}"#,
|
||||
))
|
||||
.expect("auth config should parse");
|
||||
|
||||
assert_eq!(auth_config.auth_method.as_deref(), Some("idc"));
|
||||
assert_eq!(
|
||||
auth_config.refresh_token.as_deref(),
|
||||
Some(
|
||||
"rrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrr"
|
||||
)
|
||||
);
|
||||
assert_eq!(auth_config.expires_at, Some(4_102_444_800));
|
||||
assert_eq!(
|
||||
auth_config.profile_arn.as_deref(),
|
||||
Some("arn:aws:bedrock:demo")
|
||||
);
|
||||
assert_eq!(auth_config.client_id.as_deref(), Some("cid"));
|
||||
assert_eq!(auth_config.client_secret.as_deref(), Some("secret"));
|
||||
assert_eq!(auth_config.effective_api_region(), "us-west-2");
|
||||
assert_eq!(auth_config.effective_kiro_version(), "1.2.3");
|
||||
assert_eq!(auth_config.effective_system_version(), "darwin#24.6.0");
|
||||
assert_eq!(auth_config.effective_node_version(), "22.21.1");
|
||||
assert_eq!(auth_config.access_token.as_deref(), Some("cached-token"));
|
||||
assert!(auth_config.is_idc_auth());
|
||||
assert!(auth_config.profile_arn_for_payload().is_none());
|
||||
assert_eq!(auth_config.effective_auth_region(), "us-east-1");
|
||||
assert!(auth_config.can_refresh_access_token());
|
||||
assert_eq!(DEFAULT_REGION, "us-east-1");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn infers_idc_when_client_credentials_exist() {
|
||||
let auth_config = KiroAuthConfig::from_raw_json(Some(
|
||||
r#"{
|
||||
"refreshToken":"rt-1",
|
||||
"clientId":"cid",
|
||||
"clientSecret":"secret",
|
||||
"profileArn":"arn:aws:bedrock:demo"
|
||||
}"#,
|
||||
))
|
||||
.expect("auth config should parse");
|
||||
|
||||
assert!(auth_config.is_idc_auth());
|
||||
assert!(auth_config.profile_arn_for_payload().is_none());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn round_trips_json_value() {
|
||||
let auth_config = KiroAuthConfig::from_raw_json(Some(
|
||||
r#"{
|
||||
"auth_method":"social",
|
||||
"refreshToken":"rt-1....................................................................................................",
|
||||
"expires_at": 4102444800,
|
||||
"profileArn":"arn:aws:bedrock:demo",
|
||||
"region":"eu-north-1",
|
||||
"apiRegion":"us-west-2",
|
||||
"machineId":"123e4567-e89b-12d3-a456-426614174000",
|
||||
"kiroVersion":"1.2.3",
|
||||
"systemVersion":"darwin#24.6.0",
|
||||
"nodeVersion":"22.21.1",
|
||||
"accessToken":"cached-token"
|
||||
}"#,
|
||||
))
|
||||
.expect("auth config should parse");
|
||||
|
||||
let value = auth_config.to_json_value();
|
||||
let reparsed = KiroAuthConfig::from_json_value(&value).expect("auth config should reparse");
|
||||
assert_eq!(reparsed, auth_config);
|
||||
}
|
||||
}
|
||||
122
crates/aether-provider-transport/src/kiro/headers.rs
Normal file
122
crates/aether-provider-transport/src/kiro/headers.rs
Normal file
@@ -0,0 +1,122 @@
|
||||
use std::collections::BTreeMap;
|
||||
|
||||
use uuid::Uuid;
|
||||
|
||||
use super::credentials::KiroAuthConfig;
|
||||
|
||||
pub const AWS_EVENTSTREAM_CONTENT_TYPE: &str = "application/vnd.amazon.eventstream";
|
||||
const AWS_SDK_JS_MAIN_VERSION: &str = "1.0.27";
|
||||
const CODEWHISPERER_OPTOUT: &str = "true";
|
||||
const KIRO_AGENT_MODE: &str = "vibe";
|
||||
|
||||
fn build_kiro_ide_tag(kiro_version: &str, machine_id: &str) -> String {
|
||||
if machine_id.trim().is_empty() {
|
||||
format!("KiroIDE-{kiro_version}")
|
||||
} else {
|
||||
format!("KiroIDE-{kiro_version}-{machine_id}")
|
||||
}
|
||||
}
|
||||
|
||||
fn build_x_amz_user_agent_main(kiro_version: &str, machine_id: &str) -> String {
|
||||
format!(
|
||||
"aws-sdk-js/{AWS_SDK_JS_MAIN_VERSION} {}",
|
||||
build_kiro_ide_tag(kiro_version, machine_id)
|
||||
)
|
||||
}
|
||||
|
||||
fn build_user_agent_main(
|
||||
system_version: &str,
|
||||
node_version: &str,
|
||||
kiro_version: &str,
|
||||
machine_id: &str,
|
||||
) -> String {
|
||||
format!(
|
||||
"aws-sdk-js/{AWS_SDK_JS_MAIN_VERSION} ua/2.1 os/{system_version} lang/js md/nodejs#{node_version} api/codewhispererstreaming#{AWS_SDK_JS_MAIN_VERSION} m/E {}",
|
||||
build_kiro_ide_tag(kiro_version, machine_id)
|
||||
)
|
||||
}
|
||||
|
||||
pub fn build_generate_assistant_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");
|
||||
|
||||
BTreeMap::from([
|
||||
(
|
||||
"accept".to_string(),
|
||||
AWS_EVENTSTREAM_CONTENT_TYPE.to_string(),
|
||||
),
|
||||
(
|
||||
"amz-sdk-invocation-id".to_string(),
|
||||
Uuid::new_v4().to_string(),
|
||||
),
|
||||
(
|
||||
"amz-sdk-request".to_string(),
|
||||
"attempt=1; max=3".to_string(),
|
||||
),
|
||||
("connection".to_string(), "close".to_string()),
|
||||
("content-type".to_string(), "application/json".to_string()),
|
||||
("host".to_string(), host),
|
||||
(
|
||||
"user-agent".to_string(),
|
||||
build_user_agent_main(system_version, node_version, kiro_version, machine_id),
|
||||
),
|
||||
(
|
||||
"x-amz-user-agent".to_string(),
|
||||
build_x_amz_user_agent_main(kiro_version, machine_id),
|
||||
),
|
||||
(
|
||||
"x-amzn-codewhisperer-optout".to_string(),
|
||||
CODEWHISPERER_OPTOUT.to_string(),
|
||||
),
|
||||
(
|
||||
"x-amzn-kiro-agent-mode".to_string(),
|
||||
KIRO_AGENT_MODE.to_string(),
|
||||
),
|
||||
])
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::super::credentials::KiroAuthConfig;
|
||||
use super::{build_generate_assistant_headers, AWS_EVENTSTREAM_CONTENT_TYPE};
|
||||
|
||||
#[test]
|
||||
fn builds_generate_assistant_headers_for_region() {
|
||||
let auth_config = KiroAuthConfig {
|
||||
auth_method: None,
|
||||
refresh_token: None,
|
||||
expires_at: None,
|
||||
profile_arn: None,
|
||||
region: None,
|
||||
auth_region: None,
|
||||
api_region: Some("us-west-2".to_string()),
|
||||
client_id: None,
|
||||
client_secret: None,
|
||||
machine_id: None,
|
||||
kiro_version: Some("1.2.3".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_generate_assistant_headers(&auth_config, "machine-123");
|
||||
assert_eq!(
|
||||
headers.get("accept").map(String::as_str),
|
||||
Some(AWS_EVENTSTREAM_CONTENT_TYPE)
|
||||
);
|
||||
assert_eq!(
|
||||
headers.get("host").map(String::as_str),
|
||||
Some("q.us-west-2.amazonaws.com")
|
||||
);
|
||||
assert_eq!(
|
||||
headers.get("x-amzn-kiro-agent-mode").map(String::as_str),
|
||||
Some("vibe")
|
||||
);
|
||||
}
|
||||
}
|
||||
133
crates/aether-provider-transport/src/kiro/policy.rs
Normal file
133
crates/aether-provider-transport/src/kiro/policy.rs
Normal file
@@ -0,0 +1,133 @@
|
||||
use super::super::snapshot::GatewayProviderTransportSnapshot;
|
||||
use super::super::{resolve_transport_tls_profile, transport_proxy_is_locally_supported};
|
||||
use super::{supports_local_kiro_request_auth_resolution, supports_local_kiro_request_shape};
|
||||
|
||||
pub fn supports_local_kiro_request_transport(transport: &GatewayProviderTransportSnapshot) -> bool {
|
||||
if !transport.provider.is_active || !transport.endpoint.is_active || !transport.key.is_active {
|
||||
return false;
|
||||
}
|
||||
if !transport
|
||||
.endpoint
|
||||
.api_format
|
||||
.trim()
|
||||
.eq_ignore_ascii_case("claude:cli")
|
||||
{
|
||||
return false;
|
||||
}
|
||||
if !supports_local_kiro_request_auth_resolution(transport) {
|
||||
return false;
|
||||
}
|
||||
if !supports_local_kiro_request_shape(
|
||||
transport.endpoint.header_rules.as_ref(),
|
||||
transport.endpoint.body_rules.as_ref(),
|
||||
) {
|
||||
return false;
|
||||
}
|
||||
|
||||
true
|
||||
}
|
||||
|
||||
pub fn supports_local_kiro_request_transport_with_network(
|
||||
transport: &GatewayProviderTransportSnapshot,
|
||||
) -> bool {
|
||||
supports_local_kiro_request_transport(transport)
|
||||
&& transport_proxy_is_locally_supported(transport)
|
||||
&& (transport.key.fingerprint.is_none()
|
||||
|| resolve_transport_tls_profile(transport).is_some())
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::super::super::snapshot::{
|
||||
GatewayProviderTransportEndpoint, GatewayProviderTransportKey,
|
||||
GatewayProviderTransportProvider, GatewayProviderTransportSnapshot,
|
||||
};
|
||||
use super::{
|
||||
supports_local_kiro_request_transport, supports_local_kiro_request_transport_with_network,
|
||||
};
|
||||
|
||||
fn sample_transport() -> GatewayProviderTransportSnapshot {
|
||||
GatewayProviderTransportSnapshot {
|
||||
provider: GatewayProviderTransportProvider {
|
||||
id: "provider-1".to_string(),
|
||||
name: "Kiro".to_string(),
|
||||
provider_type: "kiro".to_string(),
|
||||
website: None,
|
||||
is_active: true,
|
||||
keep_priority_on_conversion: false,
|
||||
enable_format_conversion: false,
|
||||
concurrent_limit: None,
|
||||
max_retries: None,
|
||||
proxy: None,
|
||||
request_timeout_secs: None,
|
||||
stream_first_byte_timeout_secs: None,
|
||||
config: None,
|
||||
},
|
||||
endpoint: GatewayProviderTransportEndpoint {
|
||||
id: "endpoint-1".to_string(),
|
||||
provider_id: "provider-1".to_string(),
|
||||
api_format: "claude:cli".to_string(),
|
||||
api_family: Some("claude".to_string()),
|
||||
endpoint_kind: Some("cli".to_string()),
|
||||
is_active: true,
|
||||
base_url: "https://kiro.example".to_string(),
|
||||
header_rules: None,
|
||||
body_rules: None,
|
||||
max_retries: None,
|
||||
custom_path: None,
|
||||
config: None,
|
||||
format_acceptance_config: None,
|
||||
proxy: None,
|
||||
},
|
||||
key: GatewayProviderTransportKey {
|
||||
id: "key-1".to_string(),
|
||||
provider_id: "provider-1".to_string(),
|
||||
name: "key".to_string(),
|
||||
auth_type: "bearer".to_string(),
|
||||
is_active: true,
|
||||
api_formats: Some(vec!["claude:cli".to_string()]),
|
||||
allowed_models: None,
|
||||
capabilities: None,
|
||||
rate_multipliers: None,
|
||||
global_priority_by_format: None,
|
||||
expires_at_unix_secs: None,
|
||||
proxy: None,
|
||||
fingerprint: None,
|
||||
decrypted_api_key: "__placeholder__".to_string(),
|
||||
decrypted_auth_config: Some(
|
||||
r#"{
|
||||
"access_token":"cached-token",
|
||||
"refresh_token":"rrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrr",
|
||||
"machine_id":"123e4567-e89b-12d3-a456-426614174000"
|
||||
}"#
|
||||
.to_string(),
|
||||
),
|
||||
},
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn supports_kiro_request_transport_when_cached_access_token_exists() {
|
||||
assert!(supports_local_kiro_request_transport(&sample_transport()));
|
||||
assert!(supports_local_kiro_request_transport_with_network(
|
||||
&sample_transport()
|
||||
));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn supports_kiro_request_transport_when_refresh_only_auth_exists() {
|
||||
let mut transport = sample_transport();
|
||||
transport.key.decrypted_api_key = "__placeholder__".to_string();
|
||||
transport.key.decrypted_auth_config = Some(
|
||||
r#"{
|
||||
"refresh_token":"rrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrr"
|
||||
}"#
|
||||
.to_string(),
|
||||
);
|
||||
|
||||
assert!(supports_local_kiro_request_transport(&transport));
|
||||
assert!(supports_local_kiro_request_transport_with_network(
|
||||
&transport
|
||||
));
|
||||
}
|
||||
}
|
||||
702
crates/aether-provider-transport/src/kiro/refresh.rs
Normal file
702
crates/aether-provider-transport/src/kiro/refresh.rs
Normal file
@@ -0,0 +1,702 @@
|
||||
use std::time::{SystemTime, UNIX_EPOCH};
|
||||
|
||||
use async_trait::async_trait;
|
||||
use serde_json::{json, Value};
|
||||
|
||||
use super::super::oauth_refresh::{
|
||||
CachedOAuthEntry, LocalOAuthRefreshAdapter, LocalOAuthRefreshError,
|
||||
LocalResolvedOAuthRequestAuth,
|
||||
};
|
||||
use super::super::snapshot::GatewayProviderTransportSnapshot;
|
||||
use super::auth::{
|
||||
build_kiro_request_auth_from_config, resolve_local_kiro_request_auth, PROVIDER_TYPE,
|
||||
};
|
||||
use super::credentials::{generate_machine_id, KiroAuthConfig};
|
||||
|
||||
const IDC_AMZ_USER_AGENT: &str = "aws-sdk-js/3.738.0 ua/2.1 os/other lang/js md/browser#unknown_unknown api/sso-oidc#3.738.0 m/E KiroIDE";
|
||||
|
||||
#[derive(Debug, Clone, Default)]
|
||||
pub struct KiroOAuthRefreshAdapter {
|
||||
social_refresh_base_url: Option<String>,
|
||||
idc_refresh_base_url: Option<String>,
|
||||
}
|
||||
|
||||
impl KiroOAuthRefreshAdapter {
|
||||
pub fn with_refresh_base_urls(
|
||||
mut self,
|
||||
social_refresh_base_url: Option<String>,
|
||||
idc_refresh_base_url: Option<String>,
|
||||
) -> Self {
|
||||
self.social_refresh_base_url = social_refresh_base_url;
|
||||
self.idc_refresh_base_url = idc_refresh_base_url;
|
||||
self
|
||||
}
|
||||
|
||||
pub async fn refresh_auth_config(
|
||||
&self,
|
||||
client: &reqwest::Client,
|
||||
auth_config: &KiroAuthConfig,
|
||||
) -> Result<KiroAuthConfig, LocalOAuthRefreshError> {
|
||||
if auth_config.is_idc_auth() {
|
||||
self.refresh_idc_token(client, auth_config).await
|
||||
} else {
|
||||
self.refresh_social_token(client, auth_config).await
|
||||
}
|
||||
}
|
||||
|
||||
fn social_refresh_url(&self, auth_config: &KiroAuthConfig) -> String {
|
||||
if let Some(base_url) = self
|
||||
.social_refresh_base_url
|
||||
.as_deref()
|
||||
.map(str::trim)
|
||||
.filter(|value| !value.is_empty())
|
||||
{
|
||||
return format!("{}/refreshToken", base_url.trim_end_matches('/'));
|
||||
}
|
||||
let region = auth_config.effective_auth_region();
|
||||
format!("https://prod.{region}.auth.desktop.kiro.dev/refreshToken")
|
||||
}
|
||||
|
||||
fn idc_refresh_url(&self, auth_config: &KiroAuthConfig) -> String {
|
||||
if let Some(base_url) = self
|
||||
.idc_refresh_base_url
|
||||
.as_deref()
|
||||
.map(str::trim)
|
||||
.filter(|value| !value.is_empty())
|
||||
{
|
||||
return format!("{}/token", base_url.trim_end_matches('/'));
|
||||
}
|
||||
let region = auth_config.effective_auth_region();
|
||||
format!("https://oidc.{region}.amazonaws.com/token")
|
||||
}
|
||||
|
||||
fn auth_config_from_entry(entry: &CachedOAuthEntry) -> Option<KiroAuthConfig> {
|
||||
entry
|
||||
.metadata
|
||||
.as_ref()
|
||||
.filter(|_| entry.provider_type.eq_ignore_ascii_case(PROVIDER_TYPE))
|
||||
.and_then(KiroAuthConfig::from_json_value)
|
||||
}
|
||||
|
||||
fn base_auth_config(
|
||||
&self,
|
||||
transport: &GatewayProviderTransportSnapshot,
|
||||
entry: Option<&CachedOAuthEntry>,
|
||||
) -> Option<KiroAuthConfig> {
|
||||
entry.and_then(Self::auth_config_from_entry).or_else(|| {
|
||||
KiroAuthConfig::from_raw_json(transport.key.decrypted_auth_config.as_deref())
|
||||
})
|
||||
}
|
||||
|
||||
fn build_cached_entry(auth_config: &KiroAuthConfig) -> Option<CachedOAuthEntry> {
|
||||
let request_auth = build_kiro_request_auth_from_config(auth_config.clone(), None)?;
|
||||
Some(CachedOAuthEntry {
|
||||
provider_type: PROVIDER_TYPE.to_string(),
|
||||
auth_header_name: request_auth.name.to_string(),
|
||||
auth_header_value: request_auth.value,
|
||||
expires_at_unix_secs: auth_config.expires_at,
|
||||
metadata: Some(auth_config.to_json_value()),
|
||||
})
|
||||
}
|
||||
|
||||
async fn refresh_social_token(
|
||||
&self,
|
||||
client: &reqwest::Client,
|
||||
auth_config: &KiroAuthConfig,
|
||||
) -> Result<KiroAuthConfig, LocalOAuthRefreshError> {
|
||||
let url = self.social_refresh_url(auth_config);
|
||||
let host = reqwest::Url::parse(&url)
|
||||
.ok()
|
||||
.and_then(|value| value.host_str().map(ToOwned::to_owned))
|
||||
.unwrap_or_else(|| {
|
||||
format!(
|
||||
"prod.{}.auth.desktop.kiro.dev",
|
||||
auth_config.effective_auth_region()
|
||||
)
|
||||
});
|
||||
let machine_id = generate_machine_id(auth_config, None).ok_or_else(|| {
|
||||
LocalOAuthRefreshError::InvalidResponse {
|
||||
provider_type: PROVIDER_TYPE,
|
||||
message: "missing machine_id seed for social refresh".to_string(),
|
||||
}
|
||||
})?;
|
||||
let kiro_version = auth_config.effective_kiro_version();
|
||||
let user_agent = build_kiro_ide_tag(kiro_version, &machine_id);
|
||||
let response = client
|
||||
.post(url)
|
||||
.header("User-Agent", user_agent)
|
||||
.header("Host", host)
|
||||
.header("Accept", "application/json, text/plain, */*")
|
||||
.header("Content-Type", "application/json")
|
||||
.header("Connection", "close")
|
||||
.header("Accept-Encoding", "gzip, compress, deflate, br")
|
||||
.json(&json!({
|
||||
"refreshToken": auth_config
|
||||
.refresh_token
|
||||
.as_deref()
|
||||
.map(str::trim)
|
||||
.unwrap_or_default()
|
||||
}))
|
||||
.send()
|
||||
.await
|
||||
.map_err(|source| LocalOAuthRefreshError::Transport {
|
||||
provider_type: PROVIDER_TYPE,
|
||||
source,
|
||||
})?;
|
||||
|
||||
let status = response.status();
|
||||
let body = response
|
||||
.text()
|
||||
.await
|
||||
.map_err(|source| LocalOAuthRefreshError::Transport {
|
||||
provider_type: PROVIDER_TYPE,
|
||||
source,
|
||||
})?;
|
||||
if !status.is_success() {
|
||||
return Err(LocalOAuthRefreshError::HttpStatus {
|
||||
provider_type: PROVIDER_TYPE,
|
||||
status_code: status.as_u16(),
|
||||
body_excerpt: truncate_body(&body),
|
||||
});
|
||||
}
|
||||
|
||||
let payload: Value =
|
||||
serde_json::from_str(&body).map_err(|_| LocalOAuthRefreshError::InvalidResponse {
|
||||
provider_type: PROVIDER_TYPE,
|
||||
message: "social refresh returned non-json body".to_string(),
|
||||
})?;
|
||||
let access_token = payload
|
||||
.get("accessToken")
|
||||
.and_then(Value::as_str)
|
||||
.map(str::trim)
|
||||
.filter(|value| !value.is_empty())
|
||||
.ok_or_else(|| LocalOAuthRefreshError::InvalidResponse {
|
||||
provider_type: PROVIDER_TYPE,
|
||||
message: "social refresh returned empty accessToken".to_string(),
|
||||
})?;
|
||||
|
||||
let mut refreshed = auth_config.clone();
|
||||
refreshed.access_token = Some(access_token.to_string());
|
||||
refreshed.expires_at = Some(resolve_expires_at(&payload));
|
||||
if refreshed
|
||||
.machine_id
|
||||
.as_deref()
|
||||
.map(str::trim)
|
||||
.is_none_or(|value| value.is_empty())
|
||||
{
|
||||
refreshed.machine_id = Some(machine_id);
|
||||
}
|
||||
if let Some(refresh_token) = payload
|
||||
.get("refreshToken")
|
||||
.and_then(Value::as_str)
|
||||
.map(str::trim)
|
||||
.filter(|value| !value.is_empty())
|
||||
{
|
||||
refreshed.refresh_token = Some(refresh_token.to_string());
|
||||
}
|
||||
if let Some(profile_arn) = payload
|
||||
.get("profileArn")
|
||||
.and_then(Value::as_str)
|
||||
.map(str::trim)
|
||||
.filter(|value| !value.is_empty())
|
||||
{
|
||||
refreshed.profile_arn = Some(profile_arn.to_string());
|
||||
}
|
||||
|
||||
Ok(refreshed)
|
||||
}
|
||||
|
||||
async fn refresh_idc_token(
|
||||
&self,
|
||||
client: &reqwest::Client,
|
||||
auth_config: &KiroAuthConfig,
|
||||
) -> Result<KiroAuthConfig, LocalOAuthRefreshError> {
|
||||
let url = self.idc_refresh_url(auth_config);
|
||||
let host = reqwest::Url::parse(&url)
|
||||
.ok()
|
||||
.and_then(|value| value.host_str().map(ToOwned::to_owned))
|
||||
.unwrap_or_else(|| {
|
||||
format!("oidc.{}.amazonaws.com", auth_config.effective_auth_region())
|
||||
});
|
||||
let response = client
|
||||
.post(url)
|
||||
.header("Content-Type", "application/json")
|
||||
.header("Host", host)
|
||||
.header("x-amz-user-agent", IDC_AMZ_USER_AGENT)
|
||||
.header("User-Agent", "node")
|
||||
.header("Accept", "*/*")
|
||||
.json(&json!({
|
||||
"clientId": auth_config
|
||||
.client_id
|
||||
.as_deref()
|
||||
.map(str::trim)
|
||||
.unwrap_or_default(),
|
||||
"clientSecret": auth_config
|
||||
.client_secret
|
||||
.as_deref()
|
||||
.map(str::trim)
|
||||
.unwrap_or_default(),
|
||||
"refreshToken": auth_config
|
||||
.refresh_token
|
||||
.as_deref()
|
||||
.map(str::trim)
|
||||
.unwrap_or_default(),
|
||||
"grantType": "refresh_token"
|
||||
}))
|
||||
.send()
|
||||
.await
|
||||
.map_err(|source| LocalOAuthRefreshError::Transport {
|
||||
provider_type: PROVIDER_TYPE,
|
||||
source,
|
||||
})?;
|
||||
|
||||
let status = response.status();
|
||||
let body = response
|
||||
.text()
|
||||
.await
|
||||
.map_err(|source| LocalOAuthRefreshError::Transport {
|
||||
provider_type: PROVIDER_TYPE,
|
||||
source,
|
||||
})?;
|
||||
if !status.is_success() {
|
||||
return Err(LocalOAuthRefreshError::HttpStatus {
|
||||
provider_type: PROVIDER_TYPE,
|
||||
status_code: status.as_u16(),
|
||||
body_excerpt: truncate_body(&body),
|
||||
});
|
||||
}
|
||||
|
||||
let payload: Value =
|
||||
serde_json::from_str(&body).map_err(|_| LocalOAuthRefreshError::InvalidResponse {
|
||||
provider_type: PROVIDER_TYPE,
|
||||
message: "idc refresh returned non-json body".to_string(),
|
||||
})?;
|
||||
let access_token = payload
|
||||
.get("accessToken")
|
||||
.and_then(Value::as_str)
|
||||
.map(str::trim)
|
||||
.filter(|value| !value.is_empty())
|
||||
.ok_or_else(|| LocalOAuthRefreshError::InvalidResponse {
|
||||
provider_type: PROVIDER_TYPE,
|
||||
message: "idc refresh returned empty accessToken".to_string(),
|
||||
})?;
|
||||
|
||||
let mut refreshed = auth_config.clone();
|
||||
refreshed.access_token = Some(access_token.to_string());
|
||||
refreshed.expires_at = Some(resolve_expires_at(&payload));
|
||||
if refreshed
|
||||
.machine_id
|
||||
.as_deref()
|
||||
.map(str::trim)
|
||||
.is_none_or(|value| value.is_empty())
|
||||
{
|
||||
refreshed.machine_id = generate_machine_id(auth_config, None);
|
||||
}
|
||||
if let Some(refresh_token) = payload
|
||||
.get("refreshToken")
|
||||
.and_then(Value::as_str)
|
||||
.map(str::trim)
|
||||
.filter(|value| !value.is_empty())
|
||||
{
|
||||
refreshed.refresh_token = Some(refresh_token.to_string());
|
||||
}
|
||||
|
||||
Ok(refreshed)
|
||||
}
|
||||
|
||||
fn refreshable_auth_config(
|
||||
&self,
|
||||
transport: &GatewayProviderTransportSnapshot,
|
||||
entry: Option<&CachedOAuthEntry>,
|
||||
) -> Option<KiroAuthConfig> {
|
||||
let auth_config = self.base_auth_config(transport, entry)?;
|
||||
auth_config
|
||||
.can_refresh_access_token()
|
||||
.then_some(auth_config)
|
||||
}
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
impl LocalOAuthRefreshAdapter for KiroOAuthRefreshAdapter {
|
||||
fn provider_type(&self) -> &'static str {
|
||||
PROVIDER_TYPE
|
||||
}
|
||||
|
||||
fn resolve_cached(
|
||||
&self,
|
||||
_transport: &GatewayProviderTransportSnapshot,
|
||||
entry: &CachedOAuthEntry,
|
||||
) -> Option<LocalResolvedOAuthRequestAuth> {
|
||||
let auth_config = Self::auth_config_from_entry(entry)?;
|
||||
let request_auth = build_kiro_request_auth_from_config(auth_config, None)?;
|
||||
Some(LocalResolvedOAuthRequestAuth::Kiro(request_auth))
|
||||
}
|
||||
|
||||
fn resolve_without_refresh(
|
||||
&self,
|
||||
transport: &GatewayProviderTransportSnapshot,
|
||||
) -> Option<LocalResolvedOAuthRequestAuth> {
|
||||
resolve_local_kiro_request_auth(transport).map(LocalResolvedOAuthRequestAuth::Kiro)
|
||||
}
|
||||
|
||||
fn should_refresh(
|
||||
&self,
|
||||
transport: &GatewayProviderTransportSnapshot,
|
||||
entry: Option<&CachedOAuthEntry>,
|
||||
) -> bool {
|
||||
entry
|
||||
.and_then(|cached| self.resolve_cached(transport, cached))
|
||||
.is_none()
|
||||
&& self.resolve_without_refresh(transport).is_none()
|
||||
&& self.refreshable_auth_config(transport, entry).is_some()
|
||||
}
|
||||
|
||||
async fn refresh(
|
||||
&self,
|
||||
client: &reqwest::Client,
|
||||
transport: &GatewayProviderTransportSnapshot,
|
||||
entry: Option<&CachedOAuthEntry>,
|
||||
) -> Result<Option<CachedOAuthEntry>, LocalOAuthRefreshError> {
|
||||
let Some(auth_config) = self.refreshable_auth_config(transport, entry) else {
|
||||
return Ok(None);
|
||||
};
|
||||
let refreshed = if auth_config.is_idc_auth() {
|
||||
self.refresh_idc_token(client, &auth_config).await?
|
||||
} else {
|
||||
self.refresh_social_token(client, &auth_config).await?
|
||||
};
|
||||
Ok(Self::build_cached_entry(&refreshed))
|
||||
}
|
||||
}
|
||||
|
||||
fn build_kiro_ide_tag(kiro_version: &str, machine_id: &str) -> String {
|
||||
if machine_id.trim().is_empty() {
|
||||
format!("KiroIDE-{kiro_version}")
|
||||
} else {
|
||||
format!("KiroIDE-{kiro_version}-{machine_id}")
|
||||
}
|
||||
}
|
||||
|
||||
fn resolve_expires_at(payload: &Value) -> u64 {
|
||||
let expires_in = payload
|
||||
.get("expiresIn")
|
||||
.and_then(|value| {
|
||||
value
|
||||
.as_u64()
|
||||
.or_else(|| value.as_str()?.parse::<u64>().ok())
|
||||
})
|
||||
.unwrap_or(3600);
|
||||
current_unix_secs().saturating_add(expires_in)
|
||||
}
|
||||
|
||||
fn current_unix_secs() -> u64 {
|
||||
SystemTime::now()
|
||||
.duration_since(UNIX_EPOCH)
|
||||
.ok()
|
||||
.map(|value| value.as_secs())
|
||||
.unwrap_or_default()
|
||||
}
|
||||
|
||||
fn truncate_body(body: &str) -> String {
|
||||
let body = body.trim();
|
||||
if body.is_empty() {
|
||||
return String::from("-");
|
||||
}
|
||||
body.chars().take(500).collect()
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use std::sync::{Arc, Mutex};
|
||||
|
||||
use super::super::super::oauth_refresh::{
|
||||
LocalOAuthRefreshAdapter, LocalResolvedOAuthRequestAuth,
|
||||
};
|
||||
use super::super::super::snapshot::{
|
||||
GatewayProviderTransportEndpoint, GatewayProviderTransportKey,
|
||||
GatewayProviderTransportProvider, GatewayProviderTransportSnapshot,
|
||||
};
|
||||
use super::{KiroOAuthRefreshAdapter, IDC_AMZ_USER_AGENT};
|
||||
use axum::body::to_bytes;
|
||||
use axum::extract::Request;
|
||||
use axum::response::IntoResponse;
|
||||
use axum::routing::any;
|
||||
use axum::{Json, Router};
|
||||
use http::StatusCode;
|
||||
use serde_json::{json, Value};
|
||||
use tokio::task::JoinHandle;
|
||||
|
||||
#[derive(Debug, Clone)]
|
||||
struct SeenRefreshRequest {
|
||||
body: Value,
|
||||
authorization: String,
|
||||
host: String,
|
||||
user_agent: String,
|
||||
x_amz_user_agent: String,
|
||||
}
|
||||
|
||||
fn sample_transport(raw_auth_config: &str) -> GatewayProviderTransportSnapshot {
|
||||
GatewayProviderTransportSnapshot {
|
||||
provider: GatewayProviderTransportProvider {
|
||||
id: "provider-1".to_string(),
|
||||
name: "Kiro".to_string(),
|
||||
provider_type: "kiro".to_string(),
|
||||
website: None,
|
||||
is_active: true,
|
||||
keep_priority_on_conversion: false,
|
||||
enable_format_conversion: false,
|
||||
concurrent_limit: None,
|
||||
max_retries: None,
|
||||
proxy: None,
|
||||
request_timeout_secs: None,
|
||||
stream_first_byte_timeout_secs: None,
|
||||
config: None,
|
||||
},
|
||||
endpoint: GatewayProviderTransportEndpoint {
|
||||
id: "endpoint-1".to_string(),
|
||||
provider_id: "provider-1".to_string(),
|
||||
api_format: "claude:cli".to_string(),
|
||||
api_family: Some("claude".to_string()),
|
||||
endpoint_kind: Some("cli".to_string()),
|
||||
is_active: true,
|
||||
base_url: "https://kiro.example".to_string(),
|
||||
header_rules: None,
|
||||
body_rules: None,
|
||||
max_retries: None,
|
||||
custom_path: None,
|
||||
config: None,
|
||||
format_acceptance_config: None,
|
||||
proxy: None,
|
||||
},
|
||||
key: GatewayProviderTransportKey {
|
||||
id: "key-1".to_string(),
|
||||
provider_id: "provider-1".to_string(),
|
||||
name: "key".to_string(),
|
||||
auth_type: "bearer".to_string(),
|
||||
is_active: true,
|
||||
api_formats: Some(vec!["claude:cli".to_string()]),
|
||||
allowed_models: None,
|
||||
capabilities: None,
|
||||
rate_multipliers: None,
|
||||
global_priority_by_format: None,
|
||||
expires_at_unix_secs: None,
|
||||
proxy: None,
|
||||
fingerprint: None,
|
||||
decrypted_api_key: "__placeholder__".to_string(),
|
||||
decrypted_auth_config: Some(raw_auth_config.to_string()),
|
||||
},
|
||||
}
|
||||
}
|
||||
|
||||
async fn start_server(app: Router) -> (String, JoinHandle<()>) {
|
||||
let listener = tokio::net::TcpListener::bind("127.0.0.1:0")
|
||||
.await
|
||||
.expect("listener should bind");
|
||||
let addr = listener
|
||||
.local_addr()
|
||||
.expect("listener should expose local addr");
|
||||
let handle = tokio::spawn(async move {
|
||||
axum::serve(listener, app).await.expect("server should run");
|
||||
});
|
||||
(format!("http://{addr}"), handle)
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn refreshes_social_token_via_adapter() {
|
||||
let seen_request = Arc::new(Mutex::new(None::<SeenRefreshRequest>));
|
||||
let seen_request_clone = Arc::clone(&seen_request);
|
||||
let server = Router::new().route(
|
||||
"/refreshToken",
|
||||
any(move |request: Request| {
|
||||
let seen_request_inner = Arc::clone(&seen_request_clone);
|
||||
async move {
|
||||
let (parts, body) = request.into_parts();
|
||||
let raw_body = to_bytes(body, usize::MAX).await.expect("body should read");
|
||||
let body: Value =
|
||||
serde_json::from_slice(&raw_body).expect("body should parse as json");
|
||||
*seen_request_inner.lock().expect("mutex should lock") =
|
||||
Some(SeenRefreshRequest {
|
||||
body,
|
||||
authorization: parts
|
||||
.headers
|
||||
.get("authorization")
|
||||
.and_then(|value| value.to_str().ok())
|
||||
.unwrap_or_default()
|
||||
.to_string(),
|
||||
host: parts
|
||||
.headers
|
||||
.get("host")
|
||||
.and_then(|value| value.to_str().ok())
|
||||
.unwrap_or_default()
|
||||
.to_string(),
|
||||
user_agent: parts
|
||||
.headers
|
||||
.get("user-agent")
|
||||
.and_then(|value| value.to_str().ok())
|
||||
.unwrap_or_default()
|
||||
.to_string(),
|
||||
x_amz_user_agent: parts
|
||||
.headers
|
||||
.get("x-amz-user-agent")
|
||||
.and_then(|value| value.to_str().ok())
|
||||
.unwrap_or_default()
|
||||
.to_string(),
|
||||
});
|
||||
(
|
||||
StatusCode::OK,
|
||||
Json(json!({
|
||||
"accessToken": "cached-kiro-access-token",
|
||||
"refreshToken": "s".repeat(120),
|
||||
"expiresIn": 3600,
|
||||
"profileArn": "arn:aws:bedrock:demo"
|
||||
})),
|
||||
)
|
||||
.into_response()
|
||||
}
|
||||
}),
|
||||
);
|
||||
let (server_url, server_handle) = start_server(server).await;
|
||||
let adapter =
|
||||
KiroOAuthRefreshAdapter::default().with_refresh_base_urls(Some(server_url), None);
|
||||
let transport = sample_transport(
|
||||
r#"{
|
||||
"refresh_token":"rrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrr",
|
||||
"machine_id":"123e4567-e89b-12d3-a456-426614174000",
|
||||
"kiro_version":"1.2.3"
|
||||
}"#,
|
||||
);
|
||||
|
||||
let entry = adapter
|
||||
.refresh(&reqwest::Client::new(), &transport, None)
|
||||
.await
|
||||
.expect("refresh should succeed")
|
||||
.expect("cached entry should exist");
|
||||
let resolved = adapter
|
||||
.resolve_cached(&transport, &entry)
|
||||
.expect("cached entry should resolve");
|
||||
let seen_request = seen_request
|
||||
.lock()
|
||||
.expect("mutex should lock")
|
||||
.clone()
|
||||
.expect("refresh request should be captured");
|
||||
|
||||
assert_eq!(seen_request.body["refreshToken"], json!("r".repeat(120)));
|
||||
assert_eq!(seen_request.authorization, "");
|
||||
assert!(!seen_request.user_agent.is_empty());
|
||||
assert_eq!(seen_request.x_amz_user_agent, "");
|
||||
assert!(!seen_request.host.trim().is_empty());
|
||||
match resolved {
|
||||
LocalResolvedOAuthRequestAuth::Kiro(auth) => {
|
||||
assert_eq!(auth.value, "Bearer cached-kiro-access-token");
|
||||
assert_eq!(
|
||||
auth.auth_config.profile_arn.as_deref(),
|
||||
Some("arn:aws:bedrock:demo")
|
||||
);
|
||||
assert!(auth.auth_config.expires_at.is_some());
|
||||
}
|
||||
other => panic!("unexpected resolved auth: {other:?}"),
|
||||
}
|
||||
|
||||
server_handle.abort();
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn refreshes_idc_token_via_adapter() {
|
||||
let seen_request = Arc::new(Mutex::new(None::<SeenRefreshRequest>));
|
||||
let seen_request_clone = Arc::clone(&seen_request);
|
||||
let server = Router::new().route(
|
||||
"/token",
|
||||
any(move |request: Request| {
|
||||
let seen_request_inner = Arc::clone(&seen_request_clone);
|
||||
async move {
|
||||
let (parts, body) = request.into_parts();
|
||||
let raw_body = to_bytes(body, usize::MAX).await.expect("body should read");
|
||||
let body: Value =
|
||||
serde_json::from_slice(&raw_body).expect("body should parse as json");
|
||||
*seen_request_inner.lock().expect("mutex should lock") =
|
||||
Some(SeenRefreshRequest {
|
||||
body,
|
||||
authorization: parts
|
||||
.headers
|
||||
.get("authorization")
|
||||
.and_then(|value| value.to_str().ok())
|
||||
.unwrap_or_default()
|
||||
.to_string(),
|
||||
host: parts
|
||||
.headers
|
||||
.get("host")
|
||||
.and_then(|value| value.to_str().ok())
|
||||
.unwrap_or_default()
|
||||
.to_string(),
|
||||
user_agent: parts
|
||||
.headers
|
||||
.get("user-agent")
|
||||
.and_then(|value| value.to_str().ok())
|
||||
.unwrap_or_default()
|
||||
.to_string(),
|
||||
x_amz_user_agent: parts
|
||||
.headers
|
||||
.get("x-amz-user-agent")
|
||||
.and_then(|value| value.to_str().ok())
|
||||
.unwrap_or_default()
|
||||
.to_string(),
|
||||
});
|
||||
(
|
||||
StatusCode::OK,
|
||||
Json(json!({
|
||||
"accessToken": "cached-idc-access-token",
|
||||
"refreshToken": "i".repeat(120),
|
||||
"expiresIn": 1800
|
||||
})),
|
||||
)
|
||||
.into_response()
|
||||
}
|
||||
}),
|
||||
);
|
||||
let (server_url, server_handle) = start_server(server).await;
|
||||
let adapter =
|
||||
KiroOAuthRefreshAdapter::default().with_refresh_base_urls(None, Some(server_url));
|
||||
let transport = sample_transport(
|
||||
r#"{
|
||||
"auth_method":"identity_center",
|
||||
"refresh_token":"rrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrr",
|
||||
"client_id":"cid",
|
||||
"client_secret":"secret",
|
||||
"profile_arn":"arn:aws:bedrock:demo"
|
||||
}"#,
|
||||
);
|
||||
|
||||
let entry = adapter
|
||||
.refresh(&reqwest::Client::new(), &transport, None)
|
||||
.await
|
||||
.expect("refresh should succeed")
|
||||
.expect("cached entry should exist");
|
||||
let resolved = adapter
|
||||
.resolve_cached(&transport, &entry)
|
||||
.expect("cached entry should resolve");
|
||||
let seen_request = seen_request
|
||||
.lock()
|
||||
.expect("mutex should lock")
|
||||
.clone()
|
||||
.expect("refresh request should be captured");
|
||||
|
||||
assert_eq!(
|
||||
seen_request.body["grantType"].as_str(),
|
||||
Some("refresh_token")
|
||||
);
|
||||
assert_eq!(seen_request.body["clientId"].as_str(), Some("cid"));
|
||||
assert_eq!(seen_request.user_agent, "node");
|
||||
assert_eq!(seen_request.x_amz_user_agent, IDC_AMZ_USER_AGENT);
|
||||
assert!(!seen_request.host.trim().is_empty());
|
||||
match resolved {
|
||||
LocalResolvedOAuthRequestAuth::Kiro(auth) => {
|
||||
assert_eq!(auth.value, "Bearer cached-idc-access-token");
|
||||
assert!(auth.auth_config.profile_arn_for_payload().is_none());
|
||||
assert!(auth.auth_config.expires_at.is_some());
|
||||
}
|
||||
other => panic!("unexpected resolved auth: {other:?}"),
|
||||
}
|
||||
|
||||
server_handle.abort();
|
||||
}
|
||||
}
|
||||
284
crates/aether-provider-transport/src/kiro/request.rs
Normal file
284
crates/aether-provider-transport/src/kiro/request.rs
Normal file
@@ -0,0 +1,284 @@
|
||||
use std::collections::BTreeMap;
|
||||
|
||||
use serde_json::{json, Value};
|
||||
|
||||
pub use super::super::rules::{
|
||||
apply_local_body_rules, apply_local_header_rules, body_rules_are_locally_supported,
|
||||
header_rules_are_locally_supported,
|
||||
};
|
||||
use super::super::should_skip_upstream_passthrough_header;
|
||||
use super::converter::convert_claude_messages_to_conversation_state;
|
||||
use super::credentials::KiroAuthConfig;
|
||||
use super::headers::build_generate_assistant_headers;
|
||||
|
||||
pub fn supports_local_kiro_request_shape(
|
||||
header_rules: Option<&Value>,
|
||||
body_rules: Option<&Value>,
|
||||
) -> bool {
|
||||
header_rules_are_locally_supported(header_rules) && body_rules_are_locally_supported(body_rules)
|
||||
}
|
||||
|
||||
pub fn build_kiro_provider_request_body(
|
||||
body_json: &Value,
|
||||
mapped_model: &str,
|
||||
auth_config: &KiroAuthConfig,
|
||||
body_rules: Option<&Value>,
|
||||
) -> Option<Value> {
|
||||
let conversation_state =
|
||||
convert_claude_messages_to_conversation_state(body_json, mapped_model)?;
|
||||
let mut provider_request_body = json!({
|
||||
"conversationState": conversation_state
|
||||
});
|
||||
|
||||
let mut inference_config = serde_json::Map::new();
|
||||
if let Some(max_tokens) = body_json
|
||||
.get("max_tokens")
|
||||
.and_then(|value| {
|
||||
value
|
||||
.as_i64()
|
||||
.or_else(|| value.as_u64().map(|value| value as i64))
|
||||
})
|
||||
.filter(|value| *value > 0)
|
||||
{
|
||||
inference_config.insert("maxTokens".to_string(), Value::from(max_tokens));
|
||||
}
|
||||
if let Some(temperature) = body_json
|
||||
.get("temperature")
|
||||
.and_then(Value::as_f64)
|
||||
.filter(|value| *value >= 0.0)
|
||||
{
|
||||
inference_config.insert("temperature".to_string(), Value::from(temperature));
|
||||
}
|
||||
if let Some(top_p) = body_json
|
||||
.get("top_p")
|
||||
.and_then(Value::as_f64)
|
||||
.filter(|value| *value > 0.0)
|
||||
{
|
||||
inference_config.insert("topP".to_string(), Value::from(top_p));
|
||||
}
|
||||
if !inference_config.is_empty() {
|
||||
provider_request_body.as_object_mut()?.insert(
|
||||
"inferenceConfig".to_string(),
|
||||
Value::Object(inference_config),
|
||||
);
|
||||
}
|
||||
|
||||
if let Some(profile_arn) = auth_config.profile_arn_for_payload() {
|
||||
provider_request_body.as_object_mut()?.insert(
|
||||
"profileArn".to_string(),
|
||||
Value::String(profile_arn.to_string()),
|
||||
);
|
||||
}
|
||||
|
||||
if !apply_local_body_rules(&mut provider_request_body, body_rules, Some(body_json)) {
|
||||
return None;
|
||||
}
|
||||
|
||||
Some(provider_request_body)
|
||||
}
|
||||
|
||||
pub fn build_kiro_provider_headers(
|
||||
headers: &http::HeaderMap,
|
||||
provider_request_body: &Value,
|
||||
original_request_body: &Value,
|
||||
header_rules: Option<&Value>,
|
||||
auth_header: &str,
|
||||
auth_value: &str,
|
||||
auth_config: &KiroAuthConfig,
|
||||
machine_id: &str,
|
||||
) -> Option<BTreeMap<String, String>> {
|
||||
let mut out = BTreeMap::new();
|
||||
for (name, value) in headers {
|
||||
let Ok(value) = value.to_str() else {
|
||||
continue;
|
||||
};
|
||||
let key = name.as_str().to_ascii_lowercase();
|
||||
if should_skip_upstream_passthrough_header(&key) {
|
||||
continue;
|
||||
}
|
||||
let value = value.trim();
|
||||
if value.is_empty() {
|
||||
continue;
|
||||
}
|
||||
out.insert(key, value.to_string());
|
||||
}
|
||||
|
||||
if !apply_local_header_rules(
|
||||
&mut out,
|
||||
header_rules,
|
||||
&[auth_header, "content-type"],
|
||||
provider_request_body,
|
||||
Some(original_request_body),
|
||||
) {
|
||||
return None;
|
||||
}
|
||||
|
||||
for (key, value) in build_generate_assistant_headers(auth_config, machine_id) {
|
||||
out.insert(key, value);
|
||||
}
|
||||
out.insert(
|
||||
auth_header.trim().to_ascii_lowercase(),
|
||||
auth_value.trim().to_string(),
|
||||
);
|
||||
out.entry("content-type".to_string())
|
||||
.or_insert_with(|| "application/json".to_string());
|
||||
out.remove("content-length");
|
||||
Some(out)
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use serde_json::json;
|
||||
|
||||
use super::super::credentials::KiroAuthConfig;
|
||||
use super::{
|
||||
build_kiro_provider_headers, build_kiro_provider_request_body,
|
||||
supports_local_kiro_request_shape,
|
||||
};
|
||||
|
||||
#[test]
|
||||
fn supports_empty_local_request_shape() {
|
||||
assert!(supports_local_kiro_request_shape(None, None));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn rejects_unsupported_rule_shape() {
|
||||
assert!(!supports_local_kiro_request_shape(
|
||||
Some(&json!({"action":"set"})),
|
||||
None
|
||||
));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn supports_simple_header_and_body_rules() {
|
||||
assert!(supports_local_kiro_request_shape(
|
||||
Some(&json!([{"action":"set","key":"x-provider-extra","value":"1"}])),
|
||||
Some(&json!([{"action":"set","path":"debugTag","value":true}]))
|
||||
));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn wraps_claude_request_into_kiro_payload_before_body_rules() {
|
||||
let auth_config = KiroAuthConfig {
|
||||
auth_method: None,
|
||||
refresh_token: Some("r".repeat(128)),
|
||||
expires_at: None,
|
||||
profile_arn: Some("arn:aws:bedrock:demo".to_string()),
|
||||
region: None,
|
||||
auth_region: None,
|
||||
api_region: Some("us-east-1".to_string()),
|
||||
client_id: None,
|
||||
client_secret: None,
|
||||
machine_id: Some("123e4567-e89b-12d3-a456-426614174000".to_string()),
|
||||
kiro_version: None,
|
||||
system_version: None,
|
||||
node_version: None,
|
||||
access_token: Some("cached-token".to_string()),
|
||||
};
|
||||
|
||||
let payload = build_kiro_provider_request_body(
|
||||
&json!({
|
||||
"messages": [{"role":"user","content":"hello"}],
|
||||
"max_tokens": 64
|
||||
}),
|
||||
"claude-sonnet-4-upstream",
|
||||
&auth_config,
|
||||
Some(&json!([
|
||||
{"action":"set","path":"debugTag","value":"kiro-local"}
|
||||
])),
|
||||
)
|
||||
.expect("payload should build");
|
||||
|
||||
assert!(payload.get("conversationState").is_some());
|
||||
assert_eq!(
|
||||
payload
|
||||
.get("inferenceConfig")
|
||||
.and_then(|value| value.get("maxTokens")),
|
||||
Some(&json!(64))
|
||||
);
|
||||
assert_eq!(
|
||||
payload.get("profileArn"),
|
||||
Some(&json!("arn:aws:bedrock:demo"))
|
||||
);
|
||||
assert_eq!(payload.get("debugTag"), Some(&json!("kiro-local")));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn applies_header_rules_before_kiro_extra_headers() {
|
||||
let auth_config = KiroAuthConfig {
|
||||
auth_method: None,
|
||||
refresh_token: Some("r".repeat(128)),
|
||||
expires_at: None,
|
||||
profile_arn: None,
|
||||
region: None,
|
||||
auth_region: None,
|
||||
api_region: Some("us-east-1".to_string()),
|
||||
client_id: None,
|
||||
client_secret: None,
|
||||
machine_id: None,
|
||||
kiro_version: None,
|
||||
system_version: None,
|
||||
node_version: None,
|
||||
access_token: Some("cached-token".to_string()),
|
||||
};
|
||||
let headers = build_kiro_provider_headers(
|
||||
&http::HeaderMap::new(),
|
||||
&json!({"conversationState": {}}),
|
||||
&json!({"messages": []}),
|
||||
Some(&json!([
|
||||
{"action":"set","key":"accept","value":"text/plain"},
|
||||
{"action":"set","key":"x-endpoint-tag","value":"kiro-local"}
|
||||
])),
|
||||
"authorization",
|
||||
"Bearer cached-token",
|
||||
&auth_config,
|
||||
"machine-123",
|
||||
)
|
||||
.expect("headers should build");
|
||||
|
||||
assert_eq!(
|
||||
headers.get("accept").map(String::as_str),
|
||||
Some("application/vnd.amazon.eventstream")
|
||||
);
|
||||
assert_eq!(
|
||||
headers.get("authorization").map(String::as_str),
|
||||
Some("Bearer cached-token")
|
||||
);
|
||||
assert_eq!(
|
||||
headers.get("x-endpoint-tag").map(String::as_str),
|
||||
Some("kiro-local")
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn omits_profile_arn_for_idc_auth() {
|
||||
let auth_config = KiroAuthConfig {
|
||||
auth_method: Some("identity_center".to_string()),
|
||||
refresh_token: Some("r".repeat(128)),
|
||||
expires_at: None,
|
||||
profile_arn: Some("arn:aws:bedrock:demo".to_string()),
|
||||
region: None,
|
||||
auth_region: None,
|
||||
api_region: Some("us-east-1".to_string()),
|
||||
client_id: Some("cid".to_string()),
|
||||
client_secret: Some("secret".to_string()),
|
||||
machine_id: None,
|
||||
kiro_version: None,
|
||||
system_version: None,
|
||||
node_version: None,
|
||||
access_token: Some("cached-token".to_string()),
|
||||
};
|
||||
|
||||
let payload = build_kiro_provider_request_body(
|
||||
&json!({
|
||||
"messages": [{"role":"user","content":"hello"}]
|
||||
}),
|
||||
"claude-sonnet-4-upstream",
|
||||
&auth_config,
|
||||
None,
|
||||
)
|
||||
.expect("payload should build");
|
||||
|
||||
assert!(payload.get("profileArn").is_none());
|
||||
}
|
||||
}
|
||||
71
crates/aether-provider-transport/src/kiro/url.rs
Normal file
71
crates/aether-provider-transport/src/kiro/url.rs
Normal file
@@ -0,0 +1,71 @@
|
||||
use super::super::url::build_passthrough_path_url;
|
||||
use super::credentials::DEFAULT_REGION;
|
||||
|
||||
pub const GENERATE_ASSISTANT_RESPONSE_PATH: &str = "/generateAssistantResponse";
|
||||
pub const KIRO_ENVELOPE_NAME: &str = "kiro:generateAssistantResponse";
|
||||
|
||||
pub fn resolve_kiro_base_url(upstream_base_url: &str, api_region: Option<&str>) -> String {
|
||||
let region = api_region
|
||||
.map(str::trim)
|
||||
.filter(|value| !value.is_empty())
|
||||
.unwrap_or(DEFAULT_REGION);
|
||||
upstream_base_url
|
||||
.trim()
|
||||
.replace("{region}", region)
|
||||
.trim_end_matches('/')
|
||||
.to_string()
|
||||
}
|
||||
|
||||
pub fn build_kiro_generate_assistant_response_url(
|
||||
upstream_base_url: &str,
|
||||
query: Option<&str>,
|
||||
api_region: Option<&str>,
|
||||
) -> Option<String> {
|
||||
let upstream_base_url = resolve_kiro_base_url(upstream_base_url, api_region);
|
||||
build_passthrough_path_url(
|
||||
upstream_base_url.as_str(),
|
||||
GENERATE_ASSISTANT_RESPONSE_PATH,
|
||||
query,
|
||||
&[],
|
||||
)
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::{
|
||||
build_kiro_generate_assistant_response_url, resolve_kiro_base_url,
|
||||
GENERATE_ASSISTANT_RESPONSE_PATH, KIRO_ENVELOPE_NAME,
|
||||
};
|
||||
|
||||
#[test]
|
||||
fn exposes_kiro_request_constants() {
|
||||
assert_eq!(
|
||||
GENERATE_ASSISTANT_RESPONSE_PATH,
|
||||
"/generateAssistantResponse"
|
||||
);
|
||||
assert_eq!(KIRO_ENVELOPE_NAME, "kiro:generateAssistantResponse");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn builds_generate_assistant_response_url() {
|
||||
assert_eq!(
|
||||
build_kiro_generate_assistant_response_url(
|
||||
"https://kiro.{region}.example?tenant=demo",
|
||||
Some("stream=true"),
|
||||
Some("us-west-2")
|
||||
)
|
||||
.as_deref(),
|
||||
Some(
|
||||
"https://kiro.us-west-2.example/generateAssistantResponse?stream=true&tenant=demo"
|
||||
)
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn resolves_region_placeholder_in_base_url() {
|
||||
assert_eq!(
|
||||
resolve_kiro_base_url("https://kiro.{region}.example/", Some("us-west-2")),
|
||||
"https://kiro.us-west-2.example"
|
||||
);
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user