Skip to content

Commit 91204a2

Browse files
jsgrrchgjsgrrchg
committed
Add Claude Bedrock gateway setup
1 parent 74a2fe3 commit 91204a2

6 files changed

Lines changed: 262 additions & 29 deletions

File tree

‎apps/desktop/native-backend/src/ai.rs‎

Lines changed: 160 additions & 15 deletions
Original file line numberDiff line numberDiff line change
@@ -150,6 +150,7 @@ struct AiRuntimeSetupPayload {
150150
#[serde(default)]
151151
gateway_headers: Option<AiSecretPatch>,
152152
anthropic_base_url: Option<String>,
153+
anthropic_bedrock_base_url: Option<String>,
153154
#[serde(default)]
154155
anthropic_custom_headers: Option<AiSecretPatch>,
155156
#[serde(default)]
@@ -4240,6 +4241,20 @@ fn acp_process_spec(
42404241
gemini_cli_auth_type(method).to_string(),
42414242
);
42424243
}
4244+
if runtime_id == CLAUDE_RUNTIME_ID && method == "gateway-bedrock" {
4245+
env.insert("CLAUDE_CODE_USE_BEDROCK".to_string(), "1".to_string());
4246+
env.entry("AWS_BEARER_TOKEN_BEDROCK".to_string())
4247+
.or_default();
4248+
}
4249+
}
4250+
if runtime_id == CLAUDE_RUNTIME_ID
4251+
&& env
4252+
.get("ANTHROPIC_BEDROCK_BASE_URL")
4253+
.is_some_and(|value| !value.is_empty())
4254+
{
4255+
env.insert("CLAUDE_CODE_USE_BEDROCK".to_string(), "1".to_string());
4256+
env.entry("AWS_BEARER_TOKEN_BEDROCK".to_string())
4257+
.or_default();
42434258
}
42444259
Ok(AcpProcessSpec {
42454260
program,
@@ -4635,6 +4650,10 @@ fn inherited_auth_method(runtime_id: &str, include_persisted: bool) -> Option<St
46354650
.or_else(|| {
46364651
env_secret_present("ANTHROPIC_API_KEY").then(|| "anthropic-api-key".to_string())
46374652
})
4653+
.or_else(|| {
4654+
env_secret_present("ANTHROPIC_BEDROCK_BASE_URL")
4655+
.then(|| "gateway-bedrock".to_string())
4656+
})
46384657
.or_else(|| env_secret_present("ANTHROPIC_BASE_URL").then(|| "gateway".to_string()))
46394658
.or_else(|| inherited_persisted_auth_method(runtime_id, include_persisted)),
46404659
GEMINI_RUNTIME_ID => env_secret_present("GEMINI_API_KEY")
@@ -4754,7 +4773,17 @@ fn auth_method_has_local_config(setup: &RuntimeSetupState, method_id: &str) -> b
47544773
.env
47554774
.get("KILO_API_KEY")
47564775
.is_some_and(|value| !value.is_empty()),
4757-
"gateway" => setup.has_gateway_config,
4776+
"gateway" => {
4777+
setup.has_gateway_config
4778+
&& !setup
4779+
.env
4780+
.get("ANTHROPIC_BEDROCK_BASE_URL")
4781+
.is_some_and(|value| !value.is_empty())
4782+
}
4783+
"gateway-bedrock" => setup
4784+
.env
4785+
.get("ANTHROPIC_BEDROCK_BASE_URL")
4786+
.is_some_and(|value| !value.is_empty()),
47584787
_ => false,
47594788
}
47604789
}
@@ -4768,6 +4797,7 @@ fn is_local_auth_method(method_id: &str) -> bool {
47684797
| "use_gemini"
47694798
| "kilo-api-key"
47704799
| "gateway"
4800+
| "gateway-bedrock"
47714801
)
47724802
}
47734803

@@ -4786,8 +4816,11 @@ fn local_auth_method_for_runtime(runtime_id: &str, setup: &RuntimeSetupState) ->
47864816
.then(|| "openai-api-key".to_string())
47874817
}),
47884818
CLAUDE_RUNTIME_ID => setup
4789-
.has_gateway_config
4790-
.then(|| "gateway".to_string())
4819+
.env
4820+
.get("ANTHROPIC_BEDROCK_BASE_URL")
4821+
.is_some_and(|value| !value.is_empty())
4822+
.then(|| "gateway-bedrock".to_string())
4823+
.or_else(|| setup.has_gateway_config.then(|| "gateway".to_string()))
47914824
.or_else(|| {
47924825
setup
47934826
.env
@@ -4823,10 +4856,9 @@ fn has_local_auth_config(runtime_id: &str, setup: &RuntimeSetupState) -> bool {
48234856
}
48244857

48254858
fn refresh_runtime_setup_flags(runtime_id: &str, setup: &mut RuntimeSetupState) {
4826-
setup.has_gateway_url = setup
4827-
.env
4828-
.get("ANTHROPIC_BASE_URL")
4829-
.is_some_and(|value| !value.is_empty());
4859+
setup.has_gateway_url = ["ANTHROPIC_BASE_URL", "ANTHROPIC_BEDROCK_BASE_URL"]
4860+
.into_iter()
4861+
.any(|key| setup.env.get(key).is_some_and(|value| !value.is_empty()));
48304862
setup.has_gateway_config = runtime_id == CLAUDE_RUNTIME_ID
48314863
&& (setup.has_gateway_url
48324864
|| setup
@@ -4849,6 +4881,9 @@ fn clear_runtime_auth_state(setup: &mut RuntimeSetupState) {
48494881
"ANTHROPIC_API_KEY",
48504882
"ANTHROPIC_CUSTOM_HEADERS",
48514883
"ANTHROPIC_BASE_URL",
4884+
"ANTHROPIC_BEDROCK_BASE_URL",
4885+
"CLAUDE_CODE_USE_BEDROCK",
4886+
"AWS_BEARER_TOKEN_BEDROCK",
48524887
"GEMINI_API_KEY",
48534888
"GOOGLE_API_KEY",
48544889
"KILO_API_KEY",
@@ -4973,13 +5008,19 @@ fn default_claude_terminal_auth_method() -> &'static str {
49735008

49745009
fn claude_auth_method_ids_for_environment(is_remote: bool) -> Vec<&'static str> {
49755010
if is_remote {
4976-
vec!["claude-login", "anthropic-api-key", "gateway"]
5011+
vec![
5012+
"claude-login",
5013+
"anthropic-api-key",
5014+
"gateway",
5015+
"gateway-bedrock",
5016+
]
49775017
} else {
49785018
vec![
49795019
"claude-ai-login",
49805020
"console-login",
49815021
"anthropic-api-key",
49825022
"gateway",
5023+
"gateway-bedrock",
49835024
]
49845025
}
49855026
}
@@ -4990,6 +5031,11 @@ fn claude_auth_methods_for_environment(is_remote: bool) -> Vec<AiAuthMethod> {
49905031
name: "Custom gateway".to_string(),
49915032
description: "Use a custom Anthropic-compatible gateway.".to_string(),
49925033
};
5034+
let bedrock_gateway = AiAuthMethod {
5035+
id: "gateway-bedrock".to_string(),
5036+
name: "Bedrock gateway".to_string(),
5037+
description: "Use a custom Bedrock-compatible Claude gateway.".to_string(),
5038+
};
49935039

49945040
if is_remote {
49955041
return vec![
@@ -5006,6 +5052,7 @@ fn claude_auth_methods_for_environment(is_remote: bool) -> Vec<AiAuthMethod> {
50065052
description: "Use an Anthropic API key stored locally.".to_string(),
50075053
},
50085054
gateway,
5055+
bedrock_gateway,
50095056
];
50105057
}
50115058

@@ -5026,6 +5073,7 @@ fn claude_auth_methods_for_environment(is_remote: bool) -> Vec<AiAuthMethod> {
50265073
description: "Use an Anthropic API key stored locally.".to_string(),
50275074
},
50285075
gateway,
5076+
bedrock_gateway,
50295077
]
50305078
}
50315079

@@ -5069,8 +5117,21 @@ fn update_auth_state(
50695117
runtime_id: &str,
50705118
input: AiRuntimeSetupPayload,
50715119
) -> Result<(), String> {
5072-
let gateway_url_touched =
5120+
let anthropic_gateway_url_touched =
50735121
input.gateway_base_url.is_some() || input.anthropic_base_url.is_some();
5122+
let bedrock_gateway_url_touched = input.anthropic_bedrock_base_url.is_some();
5123+
let gateway_url_touched = anthropic_gateway_url_touched || bedrock_gateway_url_touched;
5124+
let gateway_auth_method = if bedrock_gateway_url_touched
5125+
|| (!anthropic_gateway_url_touched
5126+
&& setup
5127+
.env
5128+
.get("ANTHROPIC_BEDROCK_BASE_URL")
5129+
.is_some_and(|value| !value.is_empty()))
5130+
{
5131+
"gateway-bedrock"
5132+
} else {
5133+
"gateway"
5134+
};
50745135
let gateway_headers_patch = input
50755136
.anthropic_custom_headers
50765137
.clone()
@@ -5084,9 +5145,16 @@ fn update_auth_state(
50845145
.as_ref()
50855146
.or(input.anthropic_base_url.as_ref())
50865147
.and_then(|value| normalize_optional_string(value.clone()));
5148+
let bedrock_gateway_base_url = input
5149+
.anthropic_bedrock_base_url
5150+
.as_ref()
5151+
.and_then(|value| normalize_optional_string(value.clone()));
50875152
if let Some(value) = gateway_base_url.as_deref() {
50885153
validate_claude_gateway_url(value)?;
50895154
}
5155+
if let Some(value) = bedrock_gateway_base_url.as_deref() {
5156+
validate_claude_gateway_url(value)?;
5157+
}
50905158

50915159
let mut touched_auth = false;
50925160
if runtime_id == CODEX_RUNTIME_ID {
@@ -5104,23 +5172,48 @@ fn update_auth_state(
51045172
}
51055173
if let Some(patch) = input.anthropic_auth_token.clone() {
51065174
let auth_method = if gateway_url_touched || gateway_headers_patch.is_some() {
5107-
"gateway"
5175+
gateway_auth_method
51085176
} else {
51095177
"console-login"
51105178
};
51115179
touched_auth |= apply_secret_patch(setup, "ANTHROPIC_AUTH_TOKEN", patch, auth_method);
51125180
}
51135181
if let Some(patch) = gateway_headers_patch {
5114-
touched_auth |= apply_secret_patch(setup, "ANTHROPIC_CUSTOM_HEADERS", patch, "gateway");
5182+
touched_auth |= apply_secret_patch(
5183+
setup,
5184+
"ANTHROPIC_CUSTOM_HEADERS",
5185+
patch,
5186+
gateway_auth_method,
5187+
);
51155188
}
5116-
if gateway_url_touched {
5189+
if anthropic_gateway_url_touched {
51175190
if let Some(value) = gateway_base_url {
51185191
setup.env.insert("ANTHROPIC_BASE_URL".to_string(), value);
5192+
setup.env.remove("ANTHROPIC_BEDROCK_BASE_URL");
5193+
setup.env.remove("CLAUDE_CODE_USE_BEDROCK");
5194+
setup.env.remove("AWS_BEARER_TOKEN_BEDROCK");
51195195
} else {
51205196
setup.env.remove("ANTHROPIC_BASE_URL");
51215197
}
51225198
touched_auth = true;
51235199
}
5200+
if bedrock_gateway_url_touched {
5201+
if let Some(value) = bedrock_gateway_base_url {
5202+
setup
5203+
.env
5204+
.insert("ANTHROPIC_BEDROCK_BASE_URL".to_string(), value);
5205+
setup
5206+
.env
5207+
.insert("CLAUDE_CODE_USE_BEDROCK".to_string(), "1".to_string());
5208+
setup.env.remove("ANTHROPIC_BASE_URL");
5209+
setup.env.remove("ANTHROPIC_AUTH_TOKEN");
5210+
} else {
5211+
setup.env.remove("ANTHROPIC_BEDROCK_BASE_URL");
5212+
setup.env.remove("CLAUDE_CODE_USE_BEDROCK");
5213+
setup.env.remove("AWS_BEARER_TOKEN_BEDROCK");
5214+
}
5215+
touched_auth = true;
5216+
}
51245217
}
51255218
if runtime_id == GEMINI_RUNTIME_ID {
51265219
if let Some(patch) = input.gemini_api_key.clone() {
@@ -5159,7 +5252,7 @@ fn update_auth_state(
51595252

51605253
refresh_runtime_setup_flags(runtime_id, setup);
51615254
if setup.has_gateway_config && gateway_config_touched {
5162-
setup.auth_method = Some("gateway".to_string());
5255+
setup.auth_method = Some(gateway_auth_method.to_string());
51635256
touched_auth = true;
51645257
}
51655258
if touched_auth {
@@ -7138,6 +7231,52 @@ mod tests {
71387231
);
71397232
}
71407233

7234+
#[test]
7235+
fn setup_uses_bedrock_gateway_auth_method_for_bedrock_gateway_url() {
7236+
let (event_tx, _event_rx) = mpsc::channel();
7237+
let ai = NativeAi::new(event_tx);
7238+
7239+
let status = ai
7240+
.update_setup(&json!({
7241+
"runtimeId": CLAUDE_RUNTIME_ID,
7242+
"input": {
7243+
"anthropic_bedrock_base_url": "https://bedrock-gateway.example",
7244+
"anthropic_custom_headers": {
7245+
"action": "set",
7246+
"value": "x-api-key: test-token"
7247+
}
7248+
}
7249+
}))
7250+
.expect("Bedrock gateway setup should update");
7251+
7252+
assert_eq!(
7253+
status.get("auth_method").and_then(Value::as_str),
7254+
Some("gateway-bedrock")
7255+
);
7256+
assert_eq!(
7257+
status.get("has_gateway_config").and_then(Value::as_bool),
7258+
Some(true)
7259+
);
7260+
7261+
let state = ai.inner.lock().unwrap();
7262+
let setup = state
7263+
.setup
7264+
.get(CLAUDE_RUNTIME_ID)
7265+
.expect("Claude setup should be stored");
7266+
assert_eq!(
7267+
setup
7268+
.env
7269+
.get("ANTHROPIC_BEDROCK_BASE_URL")
7270+
.map(String::as_str),
7271+
Some("https://bedrock-gateway.example")
7272+
);
7273+
assert_eq!(
7274+
setup.env.get("CLAUDE_CODE_USE_BEDROCK").map(String::as_str),
7275+
Some("1")
7276+
);
7277+
assert!(!setup.env.contains_key("ANTHROPIC_BASE_URL"));
7278+
}
7279+
71417280
#[test]
71427281
fn setup_accepts_anthropic_api_key_auth() {
71437282
let (event_tx, _event_rx) = mpsc::channel();
@@ -8295,7 +8434,8 @@ mod tests {
82958434
"claude-ai-login",
82968435
"console-login",
82978436
"anthropic-api-key",
8298-
"gateway"
8437+
"gateway",
8438+
"gateway-bedrock"
82998439
]
83008440
);
83018441

@@ -8311,7 +8451,12 @@ mod tests {
83118451
let method_ids = claude_auth_method_ids_for_environment(true);
83128452
assert_eq!(
83138453
method_ids,
8314-
vec!["claude-login", "anthropic-api-key", "gateway"]
8454+
vec![
8455+
"claude-login",
8456+
"anthropic-api-key",
8457+
"gateway",
8458+
"gateway-bedrock"
8459+
]
83158460
);
83168461

83178462
let methods = claude_auth_methods_for_environment(true)

‎apps/desktop/src/features/ai/api.ts‎

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -284,6 +284,7 @@ export async function aiUpdateSetup(input: {
284284
gatewayBaseUrl?: string;
285285
gatewayHeaders: AISecretPatch;
286286
anthropicBaseUrl?: string;
287+
anthropicBedrockBaseUrl?: string;
287288
anthropicCustomHeaders: AISecretPatch;
288289
anthropicAuthToken: AISecretPatch;
289290
anthropicApiKey?: AISecretPatch;
@@ -303,6 +304,8 @@ export async function aiUpdateSetup(input: {
303304
gateway_base_url: input.gatewayBaseUrl ?? null,
304305
gateway_headers: input.gatewayHeaders,
305306
anthropic_base_url: input.anthropicBaseUrl ?? null,
307+
anthropic_bedrock_base_url:
308+
input.anthropicBedrockBaseUrl ?? null,
306309
anthropic_custom_headers: input.anthropicCustomHeaders,
307310
anthropic_auth_token: input.anthropicAuthToken,
308311
anthropic_api_key: input.anthropicApiKey ?? { action: "unchanged" },

‎apps/desktop/src/features/ai/store/chatStore.ts‎

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1193,6 +1193,7 @@ interface ChatStore {
11931193
gatewayBaseUrl?: string;
11941194
gatewayHeaders: AISecretPatch;
11951195
anthropicBaseUrl?: string;
1196+
anthropicBedrockBaseUrl?: string;
11961197
anthropicCustomHeaders: AISecretPatch;
11971198
anthropicAuthToken: AISecretPatch;
11981199
anthropicApiKey?: AISecretPatch;
@@ -1211,6 +1212,7 @@ interface ChatStore {
12111212
gatewayBaseUrl?: string;
12121213
gatewayHeaders: AISecretPatch;
12131214
anthropicBaseUrl?: string;
1215+
anthropicBedrockBaseUrl?: string;
12141216
anthropicCustomHeaders: AISecretPatch;
12151217
anthropicAuthToken: AISecretPatch;
12161218
anthropicApiKey?: AISecretPatch;
@@ -6860,6 +6862,7 @@ export const useChatStore = create<ChatStore>((set, get) => {
68606862
input.gatewayBaseUrl ||
68616863
secretPatchChanged(input.gatewayHeaders) ||
68626864
input.anthropicBaseUrl ||
6865+
input.anthropicBedrockBaseUrl ||
68636866
secretPatchChanged(input.anthropicCustomHeaders) ||
68646867
secretPatchChanged(input.anthropicAuthToken) ||
68656868
secretPatchChanged(
@@ -6879,6 +6882,7 @@ export const useChatStore = create<ChatStore>((set, get) => {
68796882
gatewayBaseUrl: input.gatewayBaseUrl,
68806883
gatewayHeaders: input.gatewayHeaders,
68816884
anthropicBaseUrl: input.anthropicBaseUrl,
6885+
anthropicBedrockBaseUrl: input.anthropicBedrockBaseUrl,
68826886
anthropicCustomHeaders: input.anthropicCustomHeaders,
68836887
anthropicAuthToken: input.anthropicAuthToken,
68846888
anthropicApiKey: input.anthropicApiKey,

0 commit comments

Comments
 (0)