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

Closes #33524

Release Notes:

- agent: Fix an issue where performing a web search request would
sometimes fail

crates/web_search_providers/src/cloud.rs

@@ -8,7 +8,8 @@ use http_client::{HttpClient, Method};
 use language_model::{LlmApiToken, RefreshLlmTokenListener};
 use web_search::{WebSearchProvider, WebSearchProviderId};
 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 {
@@ -73,9 +74,18 @@ async fn perform_web_search(
     llm_api_token: LlmApiToken,
     body: WebSearchBody,
 ) -> Result<WebSearchResponse> {
-    let http_client = &client.http_client();
+    const MAX_RETRIES: usize = 3;
 
-    let token = llm_api_token.acquire(&client).await?;
+    let http_client = &client.http_client();
+    let mut retries_remaining = MAX_RETRIES;
+    let mut 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()
             .method(Method::POST)
@@ -93,7 +103,16 @@ async fn perform_web_search(
             let mut body = String::new();
             response.body_mut().read_to_string(&mut body).await?;
             return Ok(serde_json::from_str(&body)?);
+        } else if response
+            .headers()
+            .get(EXPIRED_LLM_TOKEN_HEADER_NAME)
+            .is_some()
+        {
+            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!(
@@ -102,3 +121,4 @@ async fn perform_web_search(
             );
         }
     }
+}