client.rs 28 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444445446447448449450451452453454455456457458459460461462463464465466467468469470471472473474475476477478479480481482483484485486487488489490491492493494495496497498499500501502503504505506507508509510511512513514515516517518519520521522523524525526527528529530531532533534535536537538539540541542543544545546547548549550551552553554555556557558559560561562563564565566567568569570571572573574575576577578579580581582583584585586587588589590591592593594595596597598599600601602603604605606607608609610611612613614615616617618619620621622623624625626627628629630631632633634635636637638639640641642643644645646647648649650651652653654655656657658659660661662663664665666667668669670671672673674675676677678679680681682683684685686687688689690691692693694695696697698699700701702703704705706707708709710711712713714715716717718719720721722723724725726727728729730731732733734735736737738739740741742743744745746747748749750751752753754755756757758759760761762763764765766767768769770771772773774775776777778779780781782783784785786787788789790791792793794795796797798799800801802803804805806807808809810811812813814815816817818819820821822823824825826827828829830831832833834835836837838839840841842843844845846847848849850851852853854855856857858859860861862863
  1. use std::collections::VecDeque;
  2. use std::time::{Duration, SystemTime, UNIX_EPOCH};
  3. use runtime::{
  4. load_oauth_credentials, save_oauth_credentials, OAuthConfig, OAuthRefreshRequest,
  5. OAuthTokenExchangeRequest,
  6. };
  7. use serde::Deserialize;
  8. use crate::error::ApiError;
  9. use crate::sse::SseParser;
  10. use crate::types::{MessageRequest, MessageResponse, StreamEvent};
  11. const DEFAULT_BASE_URL: &str = "https://api.anthropic.com";
  12. const ANTHROPIC_VERSION: &str = "2023-06-01";
  13. const REQUEST_ID_HEADER: &str = "request-id";
  14. const ALT_REQUEST_ID_HEADER: &str = "x-request-id";
  15. const DEFAULT_INITIAL_BACKOFF: Duration = Duration::from_millis(200);
  16. const DEFAULT_MAX_BACKOFF: Duration = Duration::from_secs(2);
  17. const DEFAULT_MAX_RETRIES: u32 = 2;
  18. #[derive(Debug, Clone, PartialEq, Eq)]
  19. pub enum AuthSource {
  20. None,
  21. ApiKey(String),
  22. BearerToken(String),
  23. ApiKeyAndBearer {
  24. api_key: String,
  25. bearer_token: String,
  26. },
  27. }
  28. impl AuthSource {
  29. pub fn from_env() -> Result<Self, ApiError> {
  30. let api_key = read_env_non_empty("ANTHROPIC_API_KEY")?;
  31. let auth_token = read_env_non_empty("ANTHROPIC_AUTH_TOKEN")?;
  32. match (api_key, auth_token) {
  33. (Some(api_key), Some(bearer_token)) => Ok(Self::ApiKeyAndBearer {
  34. api_key,
  35. bearer_token,
  36. }),
  37. (Some(api_key), None) => Ok(Self::ApiKey(api_key)),
  38. (None, Some(bearer_token)) => Ok(Self::BearerToken(bearer_token)),
  39. (None, None) => Err(ApiError::MissingApiKey),
  40. }
  41. }
  42. #[must_use]
  43. pub fn api_key(&self) -> Option<&str> {
  44. match self {
  45. Self::ApiKey(api_key) | Self::ApiKeyAndBearer { api_key, .. } => Some(api_key),
  46. Self::None | Self::BearerToken(_) => None,
  47. }
  48. }
  49. #[must_use]
  50. pub fn bearer_token(&self) -> Option<&str> {
  51. match self {
  52. Self::BearerToken(token)
  53. | Self::ApiKeyAndBearer {
  54. bearer_token: token,
  55. ..
  56. } => Some(token),
  57. Self::None | Self::ApiKey(_) => None,
  58. }
  59. }
  60. #[must_use]
  61. pub fn masked_authorization_header(&self) -> &'static str {
  62. if self.bearer_token().is_some() {
  63. "Bearer [REDACTED]"
  64. } else {
  65. "<absent>"
  66. }
  67. }
  68. pub fn apply(&self, mut request_builder: reqwest::RequestBuilder) -> reqwest::RequestBuilder {
  69. if let Some(api_key) = self.api_key() {
  70. request_builder = request_builder.header("x-api-key", api_key);
  71. }
  72. if let Some(token) = self.bearer_token() {
  73. request_builder = request_builder.bearer_auth(token);
  74. }
  75. request_builder
  76. }
  77. }
  78. #[derive(Debug, Clone, PartialEq, Eq, Deserialize)]
  79. pub struct OAuthTokenSet {
  80. pub access_token: String,
  81. pub refresh_token: Option<String>,
  82. pub expires_at: Option<u64>,
  83. #[serde(default)]
  84. pub scopes: Vec<String>,
  85. }
  86. impl From<OAuthTokenSet> for AuthSource {
  87. fn from(value: OAuthTokenSet) -> Self {
  88. Self::BearerToken(value.access_token)
  89. }
  90. }
  91. #[derive(Debug, Clone)]
  92. pub struct AnthropicClient {
  93. http: reqwest::Client,
  94. auth: AuthSource,
  95. base_url: String,
  96. max_retries: u32,
  97. initial_backoff: Duration,
  98. max_backoff: Duration,
  99. }
  100. impl AnthropicClient {
  101. #[must_use]
  102. pub fn new(api_key: impl Into<String>) -> Self {
  103. Self {
  104. http: reqwest::Client::new(),
  105. auth: AuthSource::ApiKey(api_key.into()),
  106. base_url: DEFAULT_BASE_URL.to_string(),
  107. max_retries: DEFAULT_MAX_RETRIES,
  108. initial_backoff: DEFAULT_INITIAL_BACKOFF,
  109. max_backoff: DEFAULT_MAX_BACKOFF,
  110. }
  111. }
  112. #[must_use]
  113. pub fn from_auth(auth: AuthSource) -> Self {
  114. Self {
  115. http: reqwest::Client::new(),
  116. auth,
  117. base_url: DEFAULT_BASE_URL.to_string(),
  118. max_retries: DEFAULT_MAX_RETRIES,
  119. initial_backoff: DEFAULT_INITIAL_BACKOFF,
  120. max_backoff: DEFAULT_MAX_BACKOFF,
  121. }
  122. }
  123. pub fn from_env() -> Result<Self, ApiError> {
  124. Ok(Self::from_auth(AuthSource::from_env_or_saved()?).with_base_url(read_base_url()))
  125. }
  126. #[must_use]
  127. pub fn with_auth_source(mut self, auth: AuthSource) -> Self {
  128. self.auth = auth;
  129. self
  130. }
  131. #[must_use]
  132. pub fn with_auth_token(mut self, auth_token: Option<String>) -> Self {
  133. match (
  134. self.auth.api_key().map(ToOwned::to_owned),
  135. auth_token.filter(|token| !token.is_empty()),
  136. ) {
  137. (Some(api_key), Some(bearer_token)) => {
  138. self.auth = AuthSource::ApiKeyAndBearer {
  139. api_key,
  140. bearer_token,
  141. };
  142. }
  143. (Some(api_key), None) => {
  144. self.auth = AuthSource::ApiKey(api_key);
  145. }
  146. (None, Some(bearer_token)) => {
  147. self.auth = AuthSource::BearerToken(bearer_token);
  148. }
  149. (None, None) => {
  150. self.auth = AuthSource::None;
  151. }
  152. }
  153. self
  154. }
  155. #[must_use]
  156. pub fn with_base_url(mut self, base_url: impl Into<String>) -> Self {
  157. self.base_url = base_url.into();
  158. self
  159. }
  160. #[must_use]
  161. pub fn with_retry_policy(
  162. mut self,
  163. max_retries: u32,
  164. initial_backoff: Duration,
  165. max_backoff: Duration,
  166. ) -> Self {
  167. self.max_retries = max_retries;
  168. self.initial_backoff = initial_backoff;
  169. self.max_backoff = max_backoff;
  170. self
  171. }
  172. #[must_use]
  173. pub fn auth_source(&self) -> &AuthSource {
  174. &self.auth
  175. }
  176. pub async fn send_message(
  177. &self,
  178. request: &MessageRequest,
  179. ) -> Result<MessageResponse, ApiError> {
  180. let request = MessageRequest {
  181. stream: false,
  182. ..request.clone()
  183. };
  184. let response = self.send_with_retry(&request).await?;
  185. let request_id = request_id_from_headers(response.headers());
  186. let mut response = response
  187. .json::<MessageResponse>()
  188. .await
  189. .map_err(ApiError::from)?;
  190. if response.request_id.is_none() {
  191. response.request_id = request_id;
  192. }
  193. Ok(response)
  194. }
  195. pub async fn stream_message(
  196. &self,
  197. request: &MessageRequest,
  198. ) -> Result<MessageStream, ApiError> {
  199. let response = self
  200. .send_with_retry(&request.clone().with_streaming())
  201. .await?;
  202. Ok(MessageStream {
  203. request_id: request_id_from_headers(response.headers()),
  204. response,
  205. parser: SseParser::new(),
  206. pending: VecDeque::new(),
  207. done: false,
  208. })
  209. }
  210. pub async fn exchange_oauth_code(
  211. &self,
  212. config: &OAuthConfig,
  213. request: &OAuthTokenExchangeRequest,
  214. ) -> Result<OAuthTokenSet, ApiError> {
  215. let response = self
  216. .http
  217. .post(&config.token_url)
  218. .header("content-type", "application/x-www-form-urlencoded")
  219. .form(&request.form_params())
  220. .send()
  221. .await
  222. .map_err(ApiError::from)?;
  223. let response = expect_success(response).await?;
  224. response
  225. .json::<OAuthTokenSet>()
  226. .await
  227. .map_err(ApiError::from)
  228. }
  229. pub async fn refresh_oauth_token(
  230. &self,
  231. config: &OAuthConfig,
  232. request: &OAuthRefreshRequest,
  233. ) -> Result<OAuthTokenSet, ApiError> {
  234. let response = self
  235. .http
  236. .post(&config.token_url)
  237. .header("content-type", "application/x-www-form-urlencoded")
  238. .form(&request.form_params())
  239. .send()
  240. .await
  241. .map_err(ApiError::from)?;
  242. let response = expect_success(response).await?;
  243. response
  244. .json::<OAuthTokenSet>()
  245. .await
  246. .map_err(ApiError::from)
  247. }
  248. async fn send_with_retry(
  249. &self,
  250. request: &MessageRequest,
  251. ) -> Result<reqwest::Response, ApiError> {
  252. let mut attempts = 0;
  253. let mut last_error: Option<ApiError>;
  254. loop {
  255. attempts += 1;
  256. match self.send_raw_request(request).await {
  257. Ok(response) => match expect_success(response).await {
  258. Ok(response) => return Ok(response),
  259. Err(error) if error.is_retryable() && attempts <= self.max_retries + 1 => {
  260. last_error = Some(error);
  261. }
  262. Err(error) => return Err(error),
  263. },
  264. Err(error) if error.is_retryable() && attempts <= self.max_retries + 1 => {
  265. last_error = Some(error);
  266. }
  267. Err(error) => return Err(error),
  268. }
  269. if attempts > self.max_retries {
  270. break;
  271. }
  272. tokio::time::sleep(self.backoff_for_attempt(attempts)?).await;
  273. }
  274. Err(ApiError::RetriesExhausted {
  275. attempts,
  276. last_error: Box::new(last_error.expect("retry loop must capture an error")),
  277. })
  278. }
  279. async fn send_raw_request(
  280. &self,
  281. request: &MessageRequest,
  282. ) -> Result<reqwest::Response, ApiError> {
  283. let request_url = format!("{}/v1/messages", self.base_url.trim_end_matches('/'));
  284. let resolved_base_url = self.base_url.trim_end_matches('/');
  285. eprintln!("[anthropic-client] resolved_base_url={resolved_base_url}");
  286. eprintln!("[anthropic-client] request_url={request_url}");
  287. let request_builder = self
  288. .http
  289. .post(&request_url)
  290. .header("anthropic-version", ANTHROPIC_VERSION)
  291. .header("content-type", "application/json");
  292. let mut request_builder = self.auth.apply(request_builder);
  293. eprintln!(
  294. "[anthropic-client] headers x-api-key={} authorization={} anthropic-version={ANTHROPIC_VERSION} content-type=application/json",
  295. if self.auth.api_key().is_some() {
  296. "[REDACTED]"
  297. } else {
  298. "<absent>"
  299. },
  300. self.auth.masked_authorization_header()
  301. );
  302. request_builder = request_builder.json(request);
  303. request_builder.send().await.map_err(ApiError::from)
  304. }
  305. fn backoff_for_attempt(&self, attempt: u32) -> Result<Duration, ApiError> {
  306. let Some(multiplier) = 1_u32.checked_shl(attempt.saturating_sub(1)) else {
  307. return Err(ApiError::BackoffOverflow {
  308. attempt,
  309. base_delay: self.initial_backoff,
  310. });
  311. };
  312. Ok(self
  313. .initial_backoff
  314. .checked_mul(multiplier)
  315. .map_or(self.max_backoff, |delay| delay.min(self.max_backoff)))
  316. }
  317. }
  318. impl AuthSource {
  319. pub fn from_env_or_saved() -> Result<Self, ApiError> {
  320. if let Some(api_key) = read_env_non_empty("ANTHROPIC_API_KEY")? {
  321. return match read_env_non_empty("ANTHROPIC_AUTH_TOKEN")? {
  322. Some(bearer_token) => Ok(Self::ApiKeyAndBearer {
  323. api_key,
  324. bearer_token,
  325. }),
  326. None => Ok(Self::ApiKey(api_key)),
  327. };
  328. }
  329. if let Some(bearer_token) = read_env_non_empty("ANTHROPIC_AUTH_TOKEN")? {
  330. return Ok(Self::BearerToken(bearer_token));
  331. }
  332. match load_saved_oauth_token() {
  333. Ok(Some(token_set)) if oauth_token_is_expired(&token_set) => {
  334. if token_set.refresh_token.is_some() {
  335. Err(ApiError::Auth(
  336. "saved OAuth token is expired; load runtime OAuth config to refresh it"
  337. .to_string(),
  338. ))
  339. } else {
  340. Err(ApiError::ExpiredOAuthToken)
  341. }
  342. }
  343. Ok(Some(token_set)) => Ok(Self::BearerToken(token_set.access_token)),
  344. Ok(None) => Err(ApiError::MissingApiKey),
  345. Err(error) => Err(error),
  346. }
  347. }
  348. }
  349. #[must_use]
  350. pub fn oauth_token_is_expired(token_set: &OAuthTokenSet) -> bool {
  351. token_set
  352. .expires_at
  353. .is_some_and(|expires_at| expires_at <= now_unix_timestamp())
  354. }
  355. pub fn resolve_saved_oauth_token(config: &OAuthConfig) -> Result<Option<OAuthTokenSet>, ApiError> {
  356. let Some(token_set) = load_saved_oauth_token()? else {
  357. return Ok(None);
  358. };
  359. if !oauth_token_is_expired(&token_set) {
  360. return Ok(Some(token_set));
  361. }
  362. let Some(refresh_token) = token_set.refresh_token.clone() else {
  363. return Err(ApiError::ExpiredOAuthToken);
  364. };
  365. let client = AnthropicClient::from_auth(AuthSource::None).with_base_url(read_base_url());
  366. let refreshed = client_runtime_block_on(async {
  367. client
  368. .refresh_oauth_token(
  369. config,
  370. &OAuthRefreshRequest::from_config(config, refresh_token, Some(token_set.scopes)),
  371. )
  372. .await
  373. })?;
  374. save_oauth_credentials(&runtime::OAuthTokenSet {
  375. access_token: refreshed.access_token.clone(),
  376. refresh_token: refreshed.refresh_token.clone(),
  377. expires_at: refreshed.expires_at,
  378. scopes: refreshed.scopes.clone(),
  379. })
  380. .map_err(ApiError::from)?;
  381. Ok(Some(refreshed))
  382. }
  383. fn client_runtime_block_on<F, T>(future: F) -> Result<T, ApiError>
  384. where
  385. F: std::future::Future<Output = Result<T, ApiError>>,
  386. {
  387. tokio::runtime::Runtime::new()
  388. .map_err(ApiError::from)?
  389. .block_on(future)
  390. }
  391. fn load_saved_oauth_token() -> Result<Option<OAuthTokenSet>, ApiError> {
  392. let token_set = load_oauth_credentials().map_err(ApiError::from)?;
  393. Ok(token_set.map(|token_set| OAuthTokenSet {
  394. access_token: token_set.access_token,
  395. refresh_token: token_set.refresh_token,
  396. expires_at: token_set.expires_at,
  397. scopes: token_set.scopes,
  398. }))
  399. }
  400. fn now_unix_timestamp() -> u64 {
  401. SystemTime::now()
  402. .duration_since(UNIX_EPOCH)
  403. .map_or(0, |duration| duration.as_secs())
  404. }
  405. fn read_env_non_empty(key: &str) -> Result<Option<String>, ApiError> {
  406. match std::env::var(key) {
  407. Ok(value) if !value.is_empty() => Ok(Some(value)),
  408. Ok(_) | Err(std::env::VarError::NotPresent) => Ok(None),
  409. Err(error) => Err(ApiError::from(error)),
  410. }
  411. }
  412. #[cfg(test)]
  413. fn read_api_key() -> Result<String, ApiError> {
  414. let auth = AuthSource::from_env_or_saved()?;
  415. auth.api_key()
  416. .or_else(|| auth.bearer_token())
  417. .map(ToOwned::to_owned)
  418. .ok_or(ApiError::MissingApiKey)
  419. }
  420. #[cfg(test)]
  421. fn read_auth_token() -> Option<String> {
  422. read_env_non_empty("ANTHROPIC_AUTH_TOKEN")
  423. .ok()
  424. .and_then(std::convert::identity)
  425. }
  426. fn read_base_url() -> String {
  427. std::env::var("ANTHROPIC_BASE_URL").unwrap_or_else(|_| DEFAULT_BASE_URL.to_string())
  428. }
  429. fn request_id_from_headers(headers: &reqwest::header::HeaderMap) -> Option<String> {
  430. headers
  431. .get(REQUEST_ID_HEADER)
  432. .or_else(|| headers.get(ALT_REQUEST_ID_HEADER))
  433. .and_then(|value| value.to_str().ok())
  434. .map(ToOwned::to_owned)
  435. }
  436. #[derive(Debug)]
  437. pub struct MessageStream {
  438. request_id: Option<String>,
  439. response: reqwest::Response,
  440. parser: SseParser,
  441. pending: VecDeque<StreamEvent>,
  442. done: bool,
  443. }
  444. impl MessageStream {
  445. #[must_use]
  446. pub fn request_id(&self) -> Option<&str> {
  447. self.request_id.as_deref()
  448. }
  449. pub async fn next_event(&mut self) -> Result<Option<StreamEvent>, ApiError> {
  450. loop {
  451. if let Some(event) = self.pending.pop_front() {
  452. return Ok(Some(event));
  453. }
  454. if self.done {
  455. let remaining = self.parser.finish()?;
  456. self.pending.extend(remaining);
  457. if let Some(event) = self.pending.pop_front() {
  458. return Ok(Some(event));
  459. }
  460. return Ok(None);
  461. }
  462. match self.response.chunk().await? {
  463. Some(chunk) => {
  464. self.pending.extend(self.parser.push(&chunk)?);
  465. }
  466. None => {
  467. self.done = true;
  468. }
  469. }
  470. }
  471. }
  472. }
  473. async fn expect_success(response: reqwest::Response) -> Result<reqwest::Response, ApiError> {
  474. let status = response.status();
  475. if status.is_success() {
  476. return Ok(response);
  477. }
  478. let body = response.text().await.unwrap_or_else(|_| String::new());
  479. let parsed_error = serde_json::from_str::<AnthropicErrorEnvelope>(&body).ok();
  480. let retryable = is_retryable_status(status);
  481. Err(ApiError::Api {
  482. status,
  483. error_type: parsed_error
  484. .as_ref()
  485. .map(|error| error.error.error_type.clone()),
  486. message: parsed_error
  487. .as_ref()
  488. .map(|error| error.error.message.clone()),
  489. body,
  490. retryable,
  491. })
  492. }
  493. const fn is_retryable_status(status: reqwest::StatusCode) -> bool {
  494. matches!(status.as_u16(), 408 | 409 | 429 | 500 | 502 | 503 | 504)
  495. }
  496. #[derive(Debug, Deserialize)]
  497. struct AnthropicErrorEnvelope {
  498. error: AnthropicErrorBody,
  499. }
  500. #[derive(Debug, Deserialize)]
  501. struct AnthropicErrorBody {
  502. #[serde(rename = "type")]
  503. error_type: String,
  504. message: String,
  505. }
  506. #[cfg(test)]
  507. mod tests {
  508. use super::{ALT_REQUEST_ID_HEADER, REQUEST_ID_HEADER};
  509. use std::io::{Read, Write};
  510. use std::net::TcpListener;
  511. use std::sync::{Mutex, OnceLock};
  512. use std::thread;
  513. use std::time::{Duration, SystemTime, UNIX_EPOCH};
  514. use runtime::{clear_oauth_credentials, save_oauth_credentials, OAuthConfig};
  515. use crate::client::{
  516. now_unix_timestamp, oauth_token_is_expired, resolve_saved_oauth_token, AnthropicClient,
  517. AuthSource, OAuthTokenSet,
  518. };
  519. use crate::types::{ContentBlockDelta, MessageRequest};
  520. fn env_lock() -> std::sync::MutexGuard<'static, ()> {
  521. static LOCK: OnceLock<Mutex<()>> = OnceLock::new();
  522. LOCK.get_or_init(|| Mutex::new(()))
  523. .lock()
  524. .expect("env lock")
  525. }
  526. fn temp_config_home() -> std::path::PathBuf {
  527. std::env::temp_dir().join(format!(
  528. "api-oauth-test-{}-{}",
  529. std::process::id(),
  530. SystemTime::now()
  531. .duration_since(UNIX_EPOCH)
  532. .expect("time")
  533. .as_nanos()
  534. ))
  535. }
  536. fn sample_oauth_config(token_url: String) -> OAuthConfig {
  537. OAuthConfig {
  538. client_id: "runtime-client".to_string(),
  539. authorize_url: "https://console.test/oauth/authorize".to_string(),
  540. token_url,
  541. callback_port: Some(4545),
  542. manual_redirect_url: Some("https://console.test/oauth/callback".to_string()),
  543. scopes: vec!["org:read".to_string(), "user:write".to_string()],
  544. }
  545. }
  546. fn spawn_token_server(response_body: &'static str) -> String {
  547. let listener = TcpListener::bind("127.0.0.1:0").expect("bind listener");
  548. let address = listener.local_addr().expect("local addr");
  549. thread::spawn(move || {
  550. let (mut stream, _) = listener.accept().expect("accept connection");
  551. let mut buffer = [0_u8; 4096];
  552. let _ = stream.read(&mut buffer).expect("read request");
  553. let response = format!(
  554. "HTTP/1.1 200 OK\r\ncontent-type: application/json\r\ncontent-length: {}\r\n\r\n{}",
  555. response_body.len(),
  556. response_body
  557. );
  558. stream
  559. .write_all(response.as_bytes())
  560. .expect("write response");
  561. });
  562. format!("http://{address}/oauth/token")
  563. }
  564. #[test]
  565. fn read_api_key_requires_presence() {
  566. let _guard = env_lock();
  567. std::env::remove_var("ANTHROPIC_AUTH_TOKEN");
  568. std::env::remove_var("ANTHROPIC_API_KEY");
  569. std::env::remove_var("CLAUDE_CONFIG_HOME");
  570. let error = super::read_api_key().expect_err("missing key should error");
  571. assert!(matches!(error, crate::error::ApiError::MissingApiKey));
  572. }
  573. #[test]
  574. fn read_api_key_requires_non_empty_value() {
  575. let _guard = env_lock();
  576. std::env::set_var("ANTHROPIC_AUTH_TOKEN", "");
  577. std::env::remove_var("ANTHROPIC_API_KEY");
  578. let error = super::read_api_key().expect_err("empty key should error");
  579. assert!(matches!(error, crate::error::ApiError::MissingApiKey));
  580. std::env::remove_var("ANTHROPIC_AUTH_TOKEN");
  581. }
  582. #[test]
  583. fn read_api_key_prefers_api_key_env() {
  584. let _guard = env_lock();
  585. std::env::set_var("ANTHROPIC_AUTH_TOKEN", "auth-token");
  586. std::env::set_var("ANTHROPIC_API_KEY", "legacy-key");
  587. assert_eq!(
  588. super::read_api_key().expect("api key should load"),
  589. "legacy-key"
  590. );
  591. std::env::remove_var("ANTHROPIC_AUTH_TOKEN");
  592. std::env::remove_var("ANTHROPIC_API_KEY");
  593. }
  594. #[test]
  595. fn read_auth_token_reads_auth_token_env() {
  596. let _guard = env_lock();
  597. std::env::set_var("ANTHROPIC_AUTH_TOKEN", "auth-token");
  598. assert_eq!(super::read_auth_token().as_deref(), Some("auth-token"));
  599. std::env::remove_var("ANTHROPIC_AUTH_TOKEN");
  600. }
  601. #[test]
  602. fn oauth_token_maps_to_bearer_auth_source() {
  603. let auth = AuthSource::from(OAuthTokenSet {
  604. access_token: "access-token".to_string(),
  605. refresh_token: Some("refresh".to_string()),
  606. expires_at: Some(123),
  607. scopes: vec!["scope:a".to_string()],
  608. });
  609. assert_eq!(auth.bearer_token(), Some("access-token"));
  610. assert_eq!(auth.api_key(), None);
  611. }
  612. #[test]
  613. fn auth_source_from_env_combines_api_key_and_bearer_token() {
  614. let _guard = env_lock();
  615. std::env::set_var("ANTHROPIC_AUTH_TOKEN", "auth-token");
  616. std::env::set_var("ANTHROPIC_API_KEY", "legacy-key");
  617. let auth = AuthSource::from_env().expect("env auth");
  618. assert_eq!(auth.api_key(), Some("legacy-key"));
  619. assert_eq!(auth.bearer_token(), Some("auth-token"));
  620. std::env::remove_var("ANTHROPIC_AUTH_TOKEN");
  621. std::env::remove_var("ANTHROPIC_API_KEY");
  622. }
  623. #[test]
  624. fn auth_source_from_saved_oauth_when_env_absent() {
  625. let _guard = env_lock();
  626. let config_home = temp_config_home();
  627. std::env::set_var("CLAUDE_CONFIG_HOME", &config_home);
  628. std::env::remove_var("ANTHROPIC_AUTH_TOKEN");
  629. std::env::remove_var("ANTHROPIC_API_KEY");
  630. save_oauth_credentials(&runtime::OAuthTokenSet {
  631. access_token: "saved-access-token".to_string(),
  632. refresh_token: Some("refresh".to_string()),
  633. expires_at: Some(now_unix_timestamp() + 300),
  634. scopes: vec!["scope:a".to_string()],
  635. })
  636. .expect("save oauth credentials");
  637. let auth = AuthSource::from_env_or_saved().expect("saved auth");
  638. assert_eq!(auth.bearer_token(), Some("saved-access-token"));
  639. clear_oauth_credentials().expect("clear credentials");
  640. std::env::remove_var("CLAUDE_CONFIG_HOME");
  641. std::fs::remove_dir_all(config_home).expect("cleanup temp dir");
  642. }
  643. #[test]
  644. fn oauth_token_expiry_uses_expires_at_timestamp() {
  645. assert!(oauth_token_is_expired(&OAuthTokenSet {
  646. access_token: "access-token".to_string(),
  647. refresh_token: None,
  648. expires_at: Some(1),
  649. scopes: Vec::new(),
  650. }));
  651. assert!(!oauth_token_is_expired(&OAuthTokenSet {
  652. access_token: "access-token".to_string(),
  653. refresh_token: None,
  654. expires_at: Some(now_unix_timestamp() + 60),
  655. scopes: Vec::new(),
  656. }));
  657. }
  658. #[test]
  659. fn resolve_saved_oauth_token_refreshes_expired_credentials() {
  660. let _guard = env_lock();
  661. let config_home = temp_config_home();
  662. std::env::set_var("CLAUDE_CONFIG_HOME", &config_home);
  663. std::env::remove_var("ANTHROPIC_AUTH_TOKEN");
  664. std::env::remove_var("ANTHROPIC_API_KEY");
  665. save_oauth_credentials(&runtime::OAuthTokenSet {
  666. access_token: "expired-access-token".to_string(),
  667. refresh_token: Some("refresh-token".to_string()),
  668. expires_at: Some(1),
  669. scopes: vec!["scope:a".to_string()],
  670. })
  671. .expect("save expired oauth credentials");
  672. let token_url = spawn_token_server(
  673. "{\"access_token\":\"refreshed-token\",\"refresh_token\":\"fresh-refresh\",\"expires_at\":9999999999,\"scopes\":[\"scope:a\"]}",
  674. );
  675. let resolved = resolve_saved_oauth_token(&sample_oauth_config(token_url))
  676. .expect("resolve refreshed token")
  677. .expect("token set present");
  678. assert_eq!(resolved.access_token, "refreshed-token");
  679. let stored = runtime::load_oauth_credentials()
  680. .expect("load stored credentials")
  681. .expect("stored token set");
  682. assert_eq!(stored.access_token, "refreshed-token");
  683. clear_oauth_credentials().expect("clear credentials");
  684. std::env::remove_var("CLAUDE_CONFIG_HOME");
  685. std::fs::remove_dir_all(config_home).expect("cleanup temp dir");
  686. }
  687. #[test]
  688. fn message_request_stream_helper_sets_stream_true() {
  689. let request = MessageRequest {
  690. model: "claude-3-7-sonnet-latest".to_string(),
  691. max_tokens: 64,
  692. messages: vec![],
  693. system: None,
  694. tools: None,
  695. tool_choice: None,
  696. stream: false,
  697. };
  698. assert!(request.with_streaming().stream);
  699. }
  700. #[test]
  701. fn backoff_doubles_until_maximum() {
  702. let client = AnthropicClient::new("test-key").with_retry_policy(
  703. 3,
  704. Duration::from_millis(10),
  705. Duration::from_millis(25),
  706. );
  707. assert_eq!(
  708. client.backoff_for_attempt(1).expect("attempt 1"),
  709. Duration::from_millis(10)
  710. );
  711. assert_eq!(
  712. client.backoff_for_attempt(2).expect("attempt 2"),
  713. Duration::from_millis(20)
  714. );
  715. assert_eq!(
  716. client.backoff_for_attempt(3).expect("attempt 3"),
  717. Duration::from_millis(25)
  718. );
  719. }
  720. #[test]
  721. fn retryable_statuses_are_detected() {
  722. assert!(super::is_retryable_status(
  723. reqwest::StatusCode::TOO_MANY_REQUESTS
  724. ));
  725. assert!(super::is_retryable_status(
  726. reqwest::StatusCode::INTERNAL_SERVER_ERROR
  727. ));
  728. assert!(!super::is_retryable_status(
  729. reqwest::StatusCode::UNAUTHORIZED
  730. ));
  731. }
  732. #[test]
  733. fn tool_delta_variant_round_trips() {
  734. let delta = ContentBlockDelta::InputJsonDelta {
  735. partial_json: "{\"city\":\"Paris\"}".to_string(),
  736. };
  737. let encoded = serde_json::to_string(&delta).expect("delta should serialize");
  738. let decoded: ContentBlockDelta =
  739. serde_json::from_str(&encoded).expect("delta should deserialize");
  740. assert_eq!(decoded, delta);
  741. }
  742. #[test]
  743. fn request_id_uses_primary_or_fallback_header() {
  744. let mut headers = reqwest::header::HeaderMap::new();
  745. headers.insert(REQUEST_ID_HEADER, "req_primary".parse().expect("header"));
  746. assert_eq!(
  747. super::request_id_from_headers(&headers).as_deref(),
  748. Some("req_primary")
  749. );
  750. headers.clear();
  751. headers.insert(
  752. ALT_REQUEST_ID_HEADER,
  753. "req_fallback".parse().expect("header"),
  754. );
  755. assert_eq!(
  756. super::request_id_from_headers(&headers).as_deref(),
  757. Some("req_fallback")
  758. );
  759. }
  760. #[test]
  761. fn auth_source_applies_headers() {
  762. let auth = AuthSource::ApiKeyAndBearer {
  763. api_key: "test-key".to_string(),
  764. bearer_token: "proxy-token".to_string(),
  765. };
  766. let request = auth
  767. .apply(reqwest::Client::new().post("https://example.test"))
  768. .build()
  769. .expect("request build");
  770. let headers = request.headers();
  771. assert_eq!(
  772. headers.get("x-api-key").and_then(|v| v.to_str().ok()),
  773. Some("test-key")
  774. );
  775. assert_eq!(
  776. headers.get("authorization").and_then(|v| v.to_str().ok()),
  777. Some("Bearer proxy-token")
  778. );
  779. }
  780. }