Yehowshua/ZIm
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

agent: Fix issue where web search could return 401 (#33639)

Browse source 
Closes #33524

Release Notes:

- agent: Fix an issue where performing a web search request would
sometimes fail
This commit is contained in:
Bennet Bo Fenner 2025-06-30 11:12:12 +02:00 • committed by GitHub
parent d63909c598
commit ac3328adb6
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
1 changed files with 44 additions and 24 deletions
Download patch file Download diff file Expand all files Collapse all files

68
crates/web_search_providers/src/cloud.rs
View file

@ -8,7 +8,8 @@ use http_client::{HttpClient, Method};
use language_model::{LlmApiToken, RefreshLlmTokenListener}; use language_model::{LlmApiToken, RefreshLlmTokenListener};
use web_search::{WebSearchProvider, WebSearchProviderId}; use web_search::{WebSearchProvider, WebSearchProviderId};
use zed_llm_client::{ use zed_llm_client::{
CLIENT_SUPPORTS_EXA_WEB_SEARCH_PROVIDER_HEADER_NAME, WebSearchBody, WebSearchResponse, CLIENT_SUPPORTS_EXA_WEB_SEARCH_PROVIDER_HEADER_NAME, EXPIRED_LLM_TOKEN_HEADER_NAME,
WebSearchBody, WebSearchResponse,
}; };
pub struct CloudWebSearchProvider { pub struct CloudWebSearchProvider {
@ -73,32 +74,51 @@ async fn perform_web_search(
llm_api_token: LlmApiToken, llm_api_token: LlmApiToken,
body: WebSearchBody, body: WebSearchBody,
) -> Result<WebSearchResponse> { ) -> Result<WebSearchResponse> {
const MAX_RETRIES: usize = 3;
let http_client = &client.http_client(); let http_client = &client.http_client();
let mut retries_remaining = MAX_RETRIES;
let mut token = llm_api_token.acquire(&client).await?;
let token = llm_api_token.acquire(&client).await?; loop {
if retries_remaining == 0 {
return Err(anyhow::anyhow!(
"error performing web search, max retries exceeded"
));
}
let request = http_client::Request::builder() let request = http_client::Request::builder()
.method(Method::POST) .method(Method::POST)
.uri(http_client.build_zed_llm_url("/web_search", &[])?.as_ref()) .uri(http_client.build_zed_llm_url("/web_search", &[])?.as_ref())
.header("Content-Type", "application/json") .header("Content-Type", "application/json")
.header("Authorization", format!("Bearer {token}")) .header("Authorization", format!("Bearer {token}"))
.header(CLIENT_SUPPORTS_EXA_WEB_SEARCH_PROVIDER_HEADER_NAME, "true") .header(CLIENT_SUPPORTS_EXA_WEB_SEARCH_PROVIDER_HEADER_NAME, "true")
.body(serde_json::to_string(&body)?.into())?; .body(serde_json::to_string(&body)?.into())?;
let mut response = http_client let mut response = http_client
.send(request) .send(request)
.await .await
.context("failed to send web search request")?; .context("failed to send web search request")?;
if response.status().is_success() { if response.status().is_success() {
let mut body = String::new(); let mut body = String::new();
response.body_mut().read_to_string(&mut body).await?; response.body_mut().read_to_string(&mut body).await?;
return Ok(serde_json::from_str(&body)?); return Ok(serde_json::from_str(&body)?);
} else { } else if response
let mut body = String::new(); .headers()
response.body_mut().read_to_string(&mut body).await?; .get(EXPIRED_LLM_TOKEN_HEADER_NAME)
anyhow::bail!( .is_some()
"error performing web search.\nStatus: {:?}\nBody: {body}", {
response.status(), token = llm_api_token.refresh(&client).await?;
); retries_remaining -= 1;
} else {
// For now we will only retry if the LLM token is expired,
// not if the request failed for any other reason.
let mut body = String::new();
response.body_mut().read_to_string(&mut body).await?;
anyhow::bail!(
"error performing web search.\nStatus: {:?}\nBody: {body}",
response.status(),
);
}
} }
} }