From bfa70f0ef2c8fdf4ba34a7e7b481999d72c0cd55 Mon Sep 17 00:00:00 2001 From: mulfyx Date: Wed, 15 Jul 2026 05:40:06 +0500 Subject: [PATCH 1/2] feat(codex): use standalone endpoint for forced web search --- docs/src/content/docs/providers/codex.md | 8 +- src/providers/codex/client.rs | 295 ++++++++++ src/providers/codex/mod.rs | 78 +++ src/providers/codex/search.rs | 539 ++++++++++++++++++ .../codex/translate/web_search_compat.rs | 15 +- 5 files changed, 932 insertions(+), 3 deletions(-) create mode 100644 src/providers/codex/search.rs diff --git a/docs/src/content/docs/providers/codex.md b/docs/src/content/docs/providers/codex.md index a2734b52..82547516 100644 --- a/docs/src/content/docs/providers/codex.md +++ b/docs/src/content/docs/providers/codex.md @@ -35,7 +35,13 @@ Claude Code summary compaction requests are capped at low effort by default beca ## Tools and multimodal input - Claude function tools and tool results map to Responses API function calls and outputs. -- Claude Code hosted web search maps to Codex `web_search`, including supported domain filters and forced tool choice. +- Claude Code's forced `web_search_20250305` subrequest uses Codex's standalone + `/alpha/search` endpoint. It keeps the resolved model, omits search reasoning, + and preserves non-empty domain filters, so Luna searches do not require a Sol + Responses turn. Automatic hosted-search requests remain on the full Responses + API because the standalone endpoint cannot decide whether to invoke a tool. + Results map back to Anthropic `server_tool_use` and + `web_search_tool_result` blocks with search usage. - Top-level base64 user images map to `input_image`. - Supported base64 images nested in tool results also map to `input_image`. - Remote image URLs, malformed images, and unsupported tool-result image forms remain textual placeholders. diff --git a/src/providers/codex/client.rs b/src/providers/codex/client.rs index d3e7fb84..56902f12 100644 --- a/src/providers/codex/client.rs +++ b/src/providers/codex/client.rs @@ -11,6 +11,7 @@ use crate::traffic::TrafficCapture; use super::auth::constants::{CODEX_API_ENDPOINT, ORIGINATOR, RESPONSES_LITE_ORIGINATOR}; use super::auth::manager::CodexAuthManager; use super::auth::token_store::{DefaultCodexAuthStore, StoredAuth, file_store}; +use super::search::{SearchRequest, SearchResponse}; use super::translate::request::ResponsesRequest; // --------------------------------------------------------------------------- @@ -180,6 +181,27 @@ pub fn build_native_codex_headers( Ok(headers) } +pub fn build_codex_search_headers( + auth: &StoredAuth, + ctx: &RequestContext, +) -> Result { + let mut headers = build_codex_headers(auth, ctx, false)?; + headers.insert( + http::header::ACCEPT, + header_value("accept", "application/json")?, + ); + let originator = config::codex_originator(RESPONSES_LITE_ORIGINATOR); + headers.insert("originator", header_value("originator", &originator)?); + let user_agent = config::codex_user_agent(RESPONSES_LITE_ORIGINATOR); + if !user_agent.is_empty() { + headers.insert( + http::header::USER_AGENT, + header_value("user-agent", &user_agent)?, + ); + } + Ok(headers) +} + fn header_value(name: &str, value: &str) -> Result { http::HeaderValue::from_str(value).map_err(|e| CodexError { status: 500, @@ -190,6 +212,14 @@ fn header_value(name: &str, value: &str) -> Result String { + let base_url = base_url.trim_end_matches('/'); + match base_url.strip_suffix("/responses") { + Some(api_root) => format!("{api_root}/alpha/search"), + None => format!("{base_url}/alpha/search"), + } +} + // --------------------------------------------------------------------------- // WebSocket request shaping // --------------------------------------------------------------------------- @@ -434,6 +464,74 @@ impl CodexHttpClient { .await } + pub async fn post_search( + &self, + body: &SearchRequest, + ctx: &RequestContext, + ) -> Result { + let mut auth = self.auth_manager.get_auth().await.map_err(|e| CodexError { + status: 401, + message: "Auth error".to_string(), + detail: Some(e.to_string()), + retry_after: None, + origin: CodexErrorOrigin::Auth, + })?; + let body_json = serde_json::to_string(body).map_err(|e| CodexError { + status: 500, + message: "Failed to serialize search request".to_string(), + detail: Some(e.to_string()), + retry_after: None, + origin: CodexErrorOrigin::Http, + })?; + let mut auth_refresh_attempted = false; + let mut retries = 0_u32; + + loop { + let response = self.attempt_post_search(&auth, &body_json, ctx).await?; + if response.status == 401 && !auth_refresh_attempted { + auth_refresh_attempted = true; + auth = self + .auth_manager + .force_refresh(&auth.access) + .await + .map_err(auth_refresh_error)?; + continue; + } + if should_retry_codex_status(response.status) + && retries < MAX_BUFFERED_TRANSPORT_RETRIES + { + let retry_after = response + .headers + .iter() + .find(|(key, _)| key.eq_ignore_ascii_case("retry-after")) + .map(|(_, value)| value.as_str()); + let delay = compute_backoff_delay(retries, retry_after); + if delay.exceeds_budget { + return Err(codex_status_error( + response, + crate::config::CodexTransport::Http, + )); + } + retries += 1; + sleep(delay.wait_ms).await; + continue; + } + if !(200..300).contains(&response.status) { + return Err(codex_status_error( + response, + crate::config::CodexTransport::Http, + )); + } + return serde_json::from_slice(&response.body).map_err(|e| CodexError { + status: 502, + message: "Failed to decode Codex search response".to_string(), + detail: Some(e.to_string()), + retry_after: None, + origin: CodexErrorOrigin::Http, + }); + } + } + async fn post_codex_with_transport( &self, body: &ResponsesRequest, @@ -1076,6 +1174,106 @@ impl CodexHttpClient { transport: ActualTransport::Http, }) } + + async fn attempt_post_search( + &self, + auth: &StoredAuth, + body_json: &str, + ctx: &RequestContext, + ) -> Result { + let url = search_endpoint(&self.base_url); + let headers = build_codex_search_headers(auth, ctx)?; + + if let Some(traffic) = ctx.traffic.as_deref() { + write_codex_http_request_capture(traffic, &url, &headers, body_json); + } + + let mut request = self.client.post(&url); + for (key, value) in headers.iter() { + request = request.header(key.as_str(), value.as_bytes()); + } + let started_at = Instant::now(); + let mut response = tokio::time::timeout( + Duration::from_millis(self.header_timeout_ms), + request.body(body_json.to_string()).send(), + ) + .await + .map_err(|_| CodexError { + status: 0, + message: format!( + "Timed out waiting {}ms for Codex search response headers", + self.header_timeout_ms + ), + detail: None, + retry_after: None, + origin: CodexErrorOrigin::Http, + })? + .map_err(|e| CodexError { + status: 0, + message: format!("Codex search network error: {e}"), + detail: None, + retry_after: None, + origin: CodexErrorOrigin::Http, + })?; + + let status = response.status().as_u16(); + let headers: Vec<(String, String)> = response + .headers() + .iter() + .map(|(key, value)| { + ( + key.to_string(), + value.to_str().unwrap_or_default().to_string(), + ) + }) + .collect(); + let mut body = Vec::new(); + let mut response_started = false; + loop { + let chunk = tokio::time::timeout( + Duration::from_millis(self.body_idle_timeout_ms), + response.chunk(), + ) + .await + .map_err(|_| CodexError { + status: 0, + message: format!( + "Timed out waiting {}ms for the next Codex search response body chunk", + self.body_idle_timeout_ms + ), + detail: Some("http_response_body".to_string()), + retry_after: None, + origin: CodexErrorOrigin::Http, + })? + .map_err(|e| CodexError { + status: 0, + message: format!("Transport error reading Codex search response body: {e}"), + detail: Some("http_response_body".to_string()), + retry_after: None, + origin: CodexErrorOrigin::Http, + })?; + let Some(chunk) = chunk else { + break; + }; + if !response_started { + if let Some(monitor) = ctx.monitor.as_ref() { + monitor.generation_started(&ctx.req_id); + } + response_started = true; + } + body.extend_from_slice(&chunk); + } + + if let Some(traffic) = ctx.traffic.as_deref() { + write_upstream_response_capture(traffic, status, started_at.elapsed(), &headers, &body); + } + + Ok(CodexResponse { + body, + status, + headers, + }) + } } fn write_codex_http_request_capture( @@ -1506,6 +1704,32 @@ mod tests { client } + async fn read_http_request(stream: &mut tokio::net::TcpStream) -> Vec { + let mut request = Vec::new(); + let mut chunk = [0_u8; 4096]; + loop { + let read = stream.read(&mut chunk).await.unwrap(); + assert!(read > 0, "request ended before its body was complete"); + request.extend_from_slice(&chunk[..read]); + let Some(header_end) = request.windows(4).position(|part| part == b"\r\n\r\n") else { + continue; + }; + let headers = String::from_utf8_lossy(&request[..header_end]); + let content_length = headers + .lines() + .find_map(|line| { + let (name, value) = line.split_once(':')?; + name.eq_ignore_ascii_case("content-length") + .then(|| value.trim().parse::().ok()) + .flatten() + }) + .unwrap_or(0); + if request.len() >= header_end + 4 + content_length { + return request; + } + } + } + #[tokio::test] async fn native_responses_replaces_auth_and_preserves_json_body() { let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); @@ -1678,6 +1902,77 @@ mod tests { assert_eq!(response.body, b"data: keep\n\n"); } + #[tokio::test] + async fn standalone_search_posts_json_to_alpha_endpoint() { + let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); + let addr = listener.local_addr().unwrap(); + let server = tokio::spawn(async move { + let (mut stream, _) = listener.accept().await.unwrap(); + let request = read_http_request(&mut stream).await; + let header_end = request + .windows(4) + .position(|part| part == b"\r\n\r\n") + .unwrap(); + let headers = String::from_utf8_lossy(&request[..header_end]); + assert!(headers.starts_with("POST /alpha/search HTTP/1.1")); + assert!( + headers + .to_ascii_lowercase() + .contains("accept: application/json") + ); + assert!(headers.contains("authorization: Bearer test")); + let body: serde_json::Value = + serde_json::from_slice(&request[header_end + 4..]).unwrap(); + assert_eq!(body["model"], "gpt-5.6-luna"); + assert!(body.get("reasoning").is_none()); + assert_eq!(body["commands"]["search_query"][0]["q"], "find Codex"); + + let response = serde_json::to_vec(&serde_json::json!({ + "encrypted_output": "opaque", + "output": "search output", + "results": [{ + "type": "text_result", + "ref_id": "turn0search0", + "url": "https://example.com", + "title": "Example" + }] + })) + .unwrap(); + let response_headers = format!( + "HTTP/1.1 200 OK\r\ncontent-type: application/json\r\ncontent-length: {}\r\nconnection: close\r\n\r\n", + response.len() + ); + stream.write_all(response_headers.as_bytes()).await.unwrap(); + stream.write_all(&response).await.unwrap(); + }); + + let client = authenticated_http_test_client(format!("http://{addr}/responses")); + let request = super::super::search::SearchRequest { + id: "session".to_string(), + model: "gpt-5.6-luna".to_string(), + reasoning: None, + input: None, + commands: super::super::search::SearchCommands { + search_query: vec![super::super::search::SearchQuery { + q: "find Codex".to_string(), + }], + }, + settings: super::super::search::SearchSettings { + filters: None, + allowed_callers: vec!["direct"], + external_web_access: true, + }, + max_output_tokens: 2_500, + }; + let response = client + .post_search(&request, &http_test_context()) + .await + .unwrap(); + server.await.unwrap(); + assert_eq!(response.output, "search output"); + assert_eq!(response.results.unwrap().len(), 1); + } + #[tokio::test] async fn auto_falls_back_to_http_after_statusful_websocket_handshake_failure() { let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); diff --git a/src/providers/codex/mod.rs b/src/providers/codex/mod.rs index 0717421e..7bbfe19e 100644 --- a/src/providers/codex/mod.rs +++ b/src/providers/codex/mod.rs @@ -7,6 +7,7 @@ pub mod count_tokens; pub(crate) mod events; pub mod native; pub mod request_summary; +pub mod search; pub mod translate; pub mod websocket; @@ -17,6 +18,7 @@ use axum::response::{IntoResponse, Response}; use bytes::Bytes; use http::StatusCode; use std::sync::Arc; +use std::time::Instant; use crate::anthropic::error::json_error; use crate::anthropic::schema::{CountTokensResponse, MessagesRequest}; @@ -123,6 +125,82 @@ impl Provider for CodexProvider { ), ); } + if search::is_standalone_search_request(&body) { + if let Some(monitor) = ctx.monitor.as_ref() { + monitor.model_resolved(&ctx.req_id, &resolved.model); + } + let (search_request, query) = match search::build_search_request( + &body, + &resolved.model, + ctx.session_id.as_deref(), + ) { + Ok(request) => request, + Err(error) => { + return json_error( + StatusCode::BAD_REQUEST, + "invalid_request_error", + error.to_string(), + ); + } + }; + let log = create_logger("codex"); + let started_at = Instant::now(); + log.info( + "codex_standalone_search_started", + Some(serde_json::Map::from_iter([ + ("reqId".to_string(), serde_json::json!(&ctx.req_id)), + ("model".to_string(), serde_json::json!(&resolved.model)), + ("stream".to_string(), serde_json::json!(want_stream)), + ])), + ); + if let Some(monitor) = ctx.monitor.as_ref() { + monitor.upstream_started(&ctx.req_id); + } + let search_response = match self.client.post_search(&search_request, &ctx).await { + Ok(response) => response, + Err(error) => { + log.warn( + "codex_standalone_search_failed", + Some(serde_json::Map::from_iter([ + ("reqId".to_string(), serde_json::json!(&ctx.req_id)), + ("model".to_string(), serde_json::json!(&resolved.model)), + ("status".to_string(), serde_json::json!(error.status)), + ( + "ms".to_string(), + serde_json::json!(started_at.elapsed().as_millis()), + ), + ])), + ); + return map_codex_error_to_response(&error); + } + }; + log.info( + "codex_standalone_search_completed", + Some(serde_json::Map::from_iter([ + ("reqId".to_string(), serde_json::json!(&ctx.req_id)), + ("model".to_string(), serde_json::json!(&resolved.model)), + ( + "resultCount".to_string(), + serde_json::json!(search_response.results.as_ref().map(Vec::len)), + ), + ( + "ms".to_string(), + serde_json::json!(started_at.elapsed().as_millis()), + ), + ])), + ); + if let Some(monitor) = ctx.monitor.as_ref() { + monitor.usage_updated(&ctx.req_id, Some(0), Some(0)); + } + return search::anthropic_search_response( + &search_response, + &query, + &message_id, + model, + want_stream, + ctx.traffic.as_deref(), + ); + } let use_responses_lite = apply_model_lane_for_request(&mut resolved.model, &body); if let Some(monitor) = ctx.monitor.as_ref() { monitor.model_resolved(&ctx.req_id, &resolved.model); diff --git a/src/providers/codex/search.rs b/src/providers/codex/search.rs new file mode 100644 index 00000000..9fd6c0c4 --- /dev/null +++ b/src/providers/codex/search.rs @@ -0,0 +1,539 @@ +use std::collections::HashSet; + +use serde::{Deserialize, Serialize}; +use serde_json::{Value, json}; + +use crate::anthropic::schema::MessagesRequest; +use crate::anthropic::sse::encode_sse_event; +use crate::traffic::TrafficCapture; + +use super::translate::web_search_compat::extract_web_search_results_from_text; + +const SEARCH_OUTPUT_TOKEN_BUDGET: u64 = 2_500; +const CLAUDE_SEARCH_PROMPT_PREFIX: &str = "Perform a web search for the query:"; + +#[derive(Debug, Clone, Serialize, PartialEq)] +pub struct SearchRequest { + pub id: String, + pub model: String, + #[serde(skip_serializing_if = "Option::is_none")] + pub reasoning: Option, + #[serde(skip_serializing_if = "Option::is_none")] + pub input: Option, + pub commands: SearchCommands, + pub settings: SearchSettings, + pub max_output_tokens: u64, +} + +#[derive(Debug, Clone, Serialize, PartialEq)] +pub struct SearchCommands { + pub search_query: Vec, +} + +#[derive(Debug, Clone, Serialize, PartialEq)] +pub struct SearchQuery { + pub q: String, +} + +#[derive(Debug, Clone, Serialize, PartialEq)] +pub struct SearchSettings { + #[serde(skip_serializing_if = "Option::is_none")] + pub filters: Option, + pub allowed_callers: Vec<&'static str>, + pub external_web_access: bool, +} + +#[derive(Debug, Clone, Serialize, PartialEq)] +pub struct SearchFilters { + #[serde(skip_serializing_if = "Option::is_none")] + pub allowed_domains: Option>, + #[serde(skip_serializing_if = "Option::is_none")] + pub blocked_domains: Option>, +} + +#[derive(Debug, Clone, Deserialize, PartialEq, Eq)] +pub struct SearchResponse { + pub encrypted_output: Option, + pub output: String, + #[serde(default)] + pub results: Option>, +} + +#[derive(Debug, Clone, PartialEq, Eq)] +struct SearchResult { + title: String, + url: String, +} + +pub fn is_standalone_search_request(req: &MessagesRequest) -> bool { + let Some(choice) = req.extra.get("tool_choice").and_then(Value::as_object) else { + return false; + }; + if choice.get("type").and_then(Value::as_str) != Some("tool") { + return false; + } + let Some(selected_name) = choice.get("name").and_then(Value::as_str) else { + return false; + }; + req.extra + .get("tools") + .and_then(Value::as_array) + .is_some_and(|tools| { + tools.iter().any(|tool| { + tool.get("type").and_then(Value::as_str) == Some("web_search_20250305") + && tool.get("name").and_then(Value::as_str) == Some(selected_name) + }) + }) +} + +pub fn build_search_request( + req: &MessagesRequest, + model: &str, + session_id: Option<&str>, +) -> Result<(SearchRequest, String), anyhow::Error> { + let query = extract_search_query(req) + .ok_or_else(|| anyhow::anyhow!("web_search request does not contain a text query"))?; + let input = search_input(req); + let filters = search_filters(req); + let id = session_id + .map(str::to_owned) + .unwrap_or_else(|| format!("search-{}", uuid::Uuid::new_v4())); + + Ok(( + SearchRequest { + id, + model: model.to_string(), + reasoning: None, + input, + commands: SearchCommands { + search_query: vec![SearchQuery { q: query.clone() }], + }, + settings: SearchSettings { + filters, + allowed_callers: vec!["direct"], + external_web_access: true, + }, + max_output_tokens: SEARCH_OUTPUT_TOKEN_BUDGET, + }, + query, + )) +} + +pub fn anthropic_search_response( + response: &SearchResponse, + query: &str, + message_id: &str, + model: &str, + stream: bool, + traffic: Option<&TrafficCapture>, +) -> axum::response::Response { + use axum::response::IntoResponse; + + let tool_use_id = format!("srvtoolu_ws_{}", uuid::Uuid::new_v4().simple()); + let results = search_results(response); + let content = response_content(response, query, &tool_use_id, &results); + let usage = json!({ + "input_tokens": 0, + "output_tokens": 0, + "cache_creation_input_tokens": 0, + "cache_read_input_tokens": 0, + "server_tool_use": {"web_search_requests": 1} + }); + + if !stream { + return ( + http::StatusCode::OK, + axum::Json(json!({ + "id": message_id, + "type": "message", + "role": "assistant", + "model": model, + "content": content, + "stop_reason": "end_turn", + "stop_sequence": null, + "usage": usage, + })), + ) + .into_response(); + } + + let mut body = Vec::new(); + emit( + &mut body, + traffic, + "message_start", + &json!({ + "type": "message_start", + "message": { + "id": message_id, + "type": "message", + "role": "assistant", + "model": model, + "content": [], + "stop_reason": null, + "stop_sequence": null, + "usage": {"input_tokens": 0, "output_tokens": 0} + } + }), + ); + emit( + &mut body, + traffic, + "content_block_start", + &json!({ + "type": "content_block_start", + "index": 0, + "content_block": { + "type": "server_tool_use", + "id": tool_use_id, + "name": "web_search", + "input": {} + } + }), + ); + emit( + &mut body, + traffic, + "content_block_delta", + &json!({ + "type": "content_block_delta", + "index": 0, + "delta": { + "type": "input_json_delta", + "partial_json": json!({"query": query}).to_string() + } + }), + ); + emit_block_stop(&mut body, traffic, 0); + emit( + &mut body, + traffic, + "content_block_start", + &json!({ + "type": "content_block_start", + "index": 1, + "content_block": { + "type": "web_search_tool_result", + "tool_use_id": tool_use_id, + "content": web_search_result_values(&results) + } + }), + ); + emit_block_stop(&mut body, traffic, 1); + if !response.output.is_empty() { + emit( + &mut body, + traffic, + "content_block_start", + &json!({ + "type": "content_block_start", + "index": 2, + "content_block": {"type": "text", "text": ""} + }), + ); + emit( + &mut body, + traffic, + "content_block_delta", + &json!({ + "type": "content_block_delta", + "index": 2, + "delta": {"type": "text_delta", "text": response.output} + }), + ); + emit_block_stop(&mut body, traffic, 2); + } + emit( + &mut body, + traffic, + "message_delta", + &json!({ + "type": "message_delta", + "delta": {"stop_reason": "end_turn", "stop_sequence": null}, + "usage": usage + }), + ); + emit( + &mut body, + traffic, + "message_stop", + &json!({"type": "message_stop"}), + ); + + let headers = [ + (http::header::CONTENT_TYPE, "text/event-stream"), + (http::header::CACHE_CONTROL, "no-cache"), + (http::header::CONNECTION, "keep-alive"), + ]; + (headers, body).into_response() +} + +fn extract_search_query(req: &MessagesRequest) -> Option { + req.messages + .iter() + .rev() + .filter(|message| message.role == "user") + .find_map(|message| { + let text = content_text(&message.content); + let text = text.trim(); + if text.is_empty() { + return None; + } + Some( + text.strip_prefix(CLAUDE_SEARCH_PROMPT_PREFIX) + .map(str::trim) + .filter(|query| !query.is_empty()) + .unwrap_or(text) + .to_string(), + ) + }) +} + +fn search_input(req: &MessagesRequest) -> Option { + let items: Vec = req + .messages + .iter() + .filter_map(|message| { + let text = content_text(&message.content); + (!text.is_empty()).then(|| { + json!({ + "type": "message", + "role": message.role, + "content": [{"type": "input_text", "text": text}] + }) + }) + }) + .collect(); + (!items.is_empty()).then_some(Value::Array(items)) +} + +fn content_text(content: &Value) -> String { + match content { + Value::String(text) => text.clone(), + Value::Array(blocks) => blocks + .iter() + .filter(|block| block.get("type").and_then(Value::as_str) == Some("text")) + .filter_map(|block| block.get("text").and_then(Value::as_str)) + .collect::>() + .join("\n"), + _ => String::new(), + } +} + +fn search_filters(req: &MessagesRequest) -> Option { + let tool = req + .extra + .get("tools")? + .as_array()? + .iter() + .find(|tool| tool.get("type").and_then(Value::as_str) == Some("web_search_20250305"))?; + let allowed_domains = string_array(tool.get("allowed_domains")); + let blocked_domains = string_array(tool.get("blocked_domains")); + (allowed_domains.is_some() || blocked_domains.is_some()).then_some(SearchFilters { + allowed_domains, + blocked_domains, + }) +} + +fn string_array(value: Option<&Value>) -> Option> { + let values = value?.as_array()?; + let values: Vec = values + .iter() + .filter_map(Value::as_str) + .map(str::to_owned) + .collect(); + (!values.is_empty()).then_some(values) +} + +fn search_results(response: &SearchResponse) -> Vec { + let mut results = Vec::new(); + let mut seen = HashSet::new(); + for result in response.results.iter().flatten() { + let Some(url) = result.get("url").and_then(Value::as_str) else { + continue; + }; + if !seen.insert(url.to_string()) { + continue; + } + let title = result + .get("title") + .and_then(Value::as_str) + .or_else(|| result.get("ref_id").and_then(Value::as_str)) + .unwrap_or(url); + results.push(SearchResult { + title: title.to_string(), + url: url.to_string(), + }); + } + if results.is_empty() { + for result in extract_web_search_results_from_text(&response.output) { + if seen.insert(result.url.clone()) { + results.push(SearchResult { + title: result.title, + url: result.url, + }); + } + } + } + results +} + +fn response_content( + response: &SearchResponse, + query: &str, + tool_use_id: &str, + results: &[SearchResult], +) -> Vec { + let mut content = vec![ + json!({ + "type": "server_tool_use", + "id": tool_use_id, + "name": "web_search", + "input": {"query": query} + }), + json!({ + "type": "web_search_tool_result", + "tool_use_id": tool_use_id, + "content": web_search_result_values(results) + }), + ]; + if !response.output.is_empty() { + content.push(json!({"type": "text", "text": response.output})); + } + content +} + +fn web_search_result_values(results: &[SearchResult]) -> Vec { + results + .iter() + .map(|result| { + json!({ + "type": "web_search_result", + "title": result.title, + "url": result.url, + }) + }) + .collect() +} + +fn emit_block_stop(out: &mut Vec, traffic: Option<&TrafficCapture>, index: usize) { + emit( + out, + traffic, + "content_block_stop", + &json!({"type": "content_block_stop", "index": index}), + ); +} + +fn emit(out: &mut Vec, traffic: Option<&TrafficCapture>, event: &str, data: &Value) { + if let Some(traffic) = traffic { + traffic.write_json_event( + "050-downstream-event", + &json!({"event": event, "data": data}), + ); + } + out.extend_from_slice(&encode_sse_event(Some(event), &data.to_string())); +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::anthropic::sse::parse_sse_events; + + fn request() -> MessagesRequest { + serde_json::from_value(json!({ + "model": "claude-haiku-4-5-20251001", + "max_tokens": 32000, + "stream": true, + "messages": [{ + "role": "user", + "content": [{ + "type": "text", + "text": "Perform a web search for the query: Codex standalone search" + }] + }], + "tools": [{ + "type": "web_search_20250305", + "name": "web_search", + "allowed_domains": ["openai.com"] + }], + "tool_choice": {"type": "tool", "name": "web_search"} + })) + .unwrap() + } + + #[test] + fn request_preserves_luna_and_omits_reasoning() { + let (request, query) = + build_search_request(&request(), "gpt-5.6-luna", Some("session-1")).unwrap(); + assert_eq!(request.model, "gpt-5.6-luna"); + assert_eq!(request.reasoning, None); + assert_eq!(request.id, "session-1"); + assert_eq!(request.max_output_tokens, 2_500); + assert_eq!(query, "Codex standalone search"); + assert_eq!(request.commands.search_query[0].q, query); + assert_eq!( + request.settings.filters.unwrap().allowed_domains, + Some(vec!["openai.com".to_string()]) + ); + } + + #[test] + fn only_forced_claude_search_uses_standalone_endpoint() { + let forced = request(); + assert!(is_standalone_search_request(&forced)); + + let mut automatic = forced.clone(); + automatic + .extra + .insert("tool_choice".to_string(), json!({"type": "auto"})); + assert!(!is_standalone_search_request(&automatic)); + } + + #[test] + fn streamed_response_matches_claude_server_tool_shape() { + let response = SearchResponse { + encrypted_output: Some("opaque".to_string()), + output: "See [OpenAI](https://openai.com).".to_string(), + results: Some(vec![json!({ + "type": "text_result", + "ref_id": "turn0search0", + "url": "https://openai.com", + "title": "OpenAI" + })]), + }; + let response = anthropic_search_response( + &response, + "Codex standalone search", + "msg_test", + "claude-haiku-4-5-20251001", + true, + None, + ); + let runtime = tokio::runtime::Runtime::new().unwrap(); + let body = runtime.block_on(axum::body::to_bytes(response.into_body(), usize::MAX)); + let events = parse_sse_events(&body.unwrap()); + let payloads: Vec = events + .iter() + .filter_map(|event| serde_json::from_str(&event.data).ok()) + .collect(); + assert!(payloads.iter().any(|payload| { + payload + .pointer("/content_block/type") + .and_then(Value::as_str) + == Some("server_tool_use") + })); + assert!(payloads.iter().any(|payload| { + payload + .pointer("/content_block/type") + .and_then(Value::as_str) + == Some("web_search_tool_result") + })); + assert!(payloads.iter().any(|payload| { + payload + .pointer("/usage/server_tool_use/web_search_requests") + .and_then(Value::as_u64) + == Some(1) + })); + } +} diff --git a/src/providers/codex/translate/web_search_compat.rs b/src/providers/codex/translate/web_search_compat.rs index eb44df17..d90d7866 100644 --- a/src/providers/codex/translate/web_search_compat.rs +++ b/src/providers/codex/translate/web_search_compat.rs @@ -37,7 +37,7 @@ pub fn server_tool_use_id_from_codex_web_search_id(id: &str) -> String { format!("srvtoolu_{suffix}") } -fn extract_web_search_results_from_text(text: &str) -> Vec { +pub(crate) fn extract_web_search_results_from_text(text: &str) -> Vec { let mut results: Vec = Vec::new(); let mut seen_urls: std::collections::HashSet = std::collections::HashSet::new(); @@ -62,7 +62,7 @@ fn extract_web_search_results_from_text(text: &str) -> Vec { } // Match bare URLs - let re2 = regex_lite::Regex::new(r"https?://[^\s<>()|]+").unwrap(); + let re2 = regex_lite::Regex::new(r#"https?://[^\s<>()|"'\\]+"#).unwrap(); for cap in re2.captures_iter(text) { let raw_url = cap.get(0).map(|m| m.as_str()).unwrap_or(""); let url = clean_url(raw_url); @@ -123,6 +123,9 @@ fn clean_url(value: &str) -> String { || out.ends_with(':') || out.ends_with('!') || out.ends_with('?') + || out.ends_with('"') + || out.ends_with('\'') + || out.ends_with('\\') { out.pop(); } @@ -176,6 +179,14 @@ mod tests { assert!(results.iter().any(|r| r.url == "https://other.com/page")); } + #[test] + fn extract_results_strips_escaped_quote_suffix() { + let text = r#"{\"url\":\"https://example.com/path\"}"#; + let results = extract_web_search_results_from_text(text); + assert_eq!(results.len(), 1); + assert_eq!(results[0].url, "https://example.com/path"); + } + #[test] fn build_compat_blocks() { let searches = vec![super::super::reducer::ReducerEvent::WebSearch { From 94e9a25008366b1a679366d42f0c8c305f023357 Mon Sep 17 00:00:00 2001 From: Raine Virta Date: Wed, 15 Jul 2026 09:38:18 +0300 Subject: [PATCH 2/2] fix standalone Codex search translation Standalone search requests failed when conversation context contained an assistant message because all content was serialized as input_text. Preserve the recent conversation tail using role-specific content types and bound assistant context to the same approximate budget used by Codex. Treat only structured Alpha Search result DTOs as search results so links in page content are not reported as independent sources. Estimate request and response token usage locally instead of reporting zero-token searches. --- docs/src/content/docs/providers/codex.md | 5 +- src/providers/codex/client.rs | 11 +- src/providers/codex/count_tokens.rs | 2 +- src/providers/codex/mod.rs | 5 +- src/providers/codex/search.rs | 205 ++++++++++++++++-- .../codex/translate/web_search_compat.rs | 15 +- 6 files changed, 201 insertions(+), 42 deletions(-) diff --git a/docs/src/content/docs/providers/codex.md b/docs/src/content/docs/providers/codex.md index 82547516..c6a9c391 100644 --- a/docs/src/content/docs/providers/codex.md +++ b/docs/src/content/docs/providers/codex.md @@ -40,8 +40,9 @@ Claude Code summary compaction requests are capped at low effort by default beca and preserves non-empty domain filters, so Luna searches do not require a Sol Responses turn. Automatic hosted-search requests remain on the full Responses API because the standalone endpoint cannot decide whether to invoke a tool. - Results map back to Anthropic `server_tool_use` and - `web_search_tool_result` blocks with search usage. + Structured result DTOs map back to Anthropic `server_tool_use` and + `web_search_tool_result` blocks, while standalone text output remains text. + The proxy locally estimates input and output tokens and reports search usage. - Top-level base64 user images map to `input_image`. - Supported base64 images nested in tool results also map to `input_image`. - Remote image URLs, malformed images, and unsupported tool-result image forms remain textual placeholders. diff --git a/src/providers/codex/client.rs b/src/providers/codex/client.rs index 56902f12..5792c80c 100644 --- a/src/providers/codex/client.rs +++ b/src/providers/codex/client.rs @@ -507,20 +507,14 @@ impl CodexHttpClient { .map(|(_, value)| value.as_str()); let delay = compute_backoff_delay(retries, retry_after); if delay.exceeds_budget { - return Err(codex_status_error( - response, - crate::config::CodexTransport::Http, - )); + return Err(codex_status_error(response)); } retries += 1; sleep(delay.wait_ms).await; continue; } if !(200..300).contains(&response.status) { - return Err(codex_status_error( - response, - crate::config::CodexTransport::Http, - )); + return Err(codex_status_error(response)); } return serde_json::from_slice(&response.body).map_err(|e| CodexError { status: 502, @@ -1272,6 +1266,7 @@ impl CodexHttpClient { body, status, headers, + transport: ActualTransport::Http, }) } } diff --git a/src/providers/codex/count_tokens.rs b/src/providers/codex/count_tokens.rs index 580e7d86..57051f83 100644 --- a/src/providers/codex/count_tokens.rs +++ b/src/providers/codex/count_tokens.rs @@ -116,7 +116,7 @@ fn count_tool_tokens(tools: &[ResponsesTool]) -> u64 { total } -fn approx_token_count(text: &str) -> u64 { +pub(crate) fn approx_token_count(text: &str) -> u64 { if text.is_empty() { return 0; } diff --git a/src/providers/codex/mod.rs b/src/providers/codex/mod.rs index 7bbfe19e..5bcb583a 100644 --- a/src/providers/codex/mod.rs +++ b/src/providers/codex/mod.rs @@ -189,8 +189,10 @@ impl Provider for CodexProvider { ), ])), ); + let input_tokens = search::search_request_input_tokens(&search_request); + let output_tokens = search::search_response_output_tokens(&search_response); if let Some(monitor) = ctx.monitor.as_ref() { - monitor.usage_updated(&ctx.req_id, Some(0), Some(0)); + monitor.usage_updated(&ctx.req_id, Some(input_tokens), Some(output_tokens)); } return search::anthropic_search_response( &search_response, @@ -198,6 +200,7 @@ impl Provider for CodexProvider { &message_id, model, want_stream, + input_tokens, ctx.traffic.as_deref(), ); } diff --git a/src/providers/codex/search.rs b/src/providers/codex/search.rs index 9fd6c0c4..8a787af2 100644 --- a/src/providers/codex/search.rs +++ b/src/providers/codex/search.rs @@ -7,9 +7,11 @@ use crate::anthropic::schema::MessagesRequest; use crate::anthropic::sse::encode_sse_event; use crate::traffic::TrafficCapture; -use super::translate::web_search_compat::extract_web_search_results_from_text; +use super::count_tokens::approx_token_count; const SEARCH_OUTPUT_TOKEN_BUDGET: u64 = 2_500; +const SEARCH_ASSISTANT_CONTEXT_TOKEN_BUDGET: u64 = 1_000; +const SEARCH_USER_CONTEXT_MESSAGES: usize = 2; const CLAUDE_SEARCH_PROMPT_PREFIX: &str = "Perform a web search for the query:"; #[derive(Debug, Clone, Serialize, PartialEq)] @@ -119,12 +121,53 @@ pub fn build_search_request( )) } +pub fn search_request_input_tokens(request: &SearchRequest) -> u64 { + let mut tokens = approx_token_count(&request.model); + tokens += request + .commands + .search_query + .iter() + .map(|query| approx_token_count(&query.q)) + .sum::(); + tokens += request.input.as_ref().map(value_text_tokens).unwrap_or(0); + if let Some(filters) = &request.settings.filters { + tokens += filters + .allowed_domains + .iter() + .flatten() + .chain(filters.blocked_domains.iter().flatten()) + .map(|domain| approx_token_count(domain)) + .sum::(); + } + tokens.max(1) +} + +pub fn search_response_output_tokens(response: &SearchResponse) -> u64 { + (approx_token_count(&response.output) + + response + .results + .as_ref() + .map(|results| results.iter().map(value_text_tokens).sum()) + .unwrap_or(0)) + .max(1) +} + +fn value_text_tokens(value: &Value) -> u64 { + match value { + Value::String(text) => approx_token_count(text), + Value::Array(values) => values.iter().map(value_text_tokens).sum(), + Value::Object(values) => values.values().map(value_text_tokens).sum(), + _ => 0, + } +} + pub fn anthropic_search_response( response: &SearchResponse, query: &str, message_id: &str, model: &str, stream: bool, + input_tokens: u64, traffic: Option<&TrafficCapture>, ) -> axum::response::Response { use axum::response::IntoResponse; @@ -132,9 +175,10 @@ pub fn anthropic_search_response( let tool_use_id = format!("srvtoolu_ws_{}", uuid::Uuid::new_v4().simple()); let results = search_results(response); let content = response_content(response, query, &tool_use_id, &results); + let output_tokens = search_response_output_tokens(response); let usage = json!({ - "input_tokens": 0, - "output_tokens": 0, + "input_tokens": input_tokens, + "output_tokens": output_tokens, "cache_creation_input_tokens": 0, "cache_read_input_tokens": 0, "server_tool_use": {"web_search_requests": 1} @@ -172,7 +216,7 @@ pub fn anthropic_search_response( "content": [], "stop_reason": null, "stop_sequence": null, - "usage": {"input_tokens": 0, "output_tokens": 0} + "usage": {"input_tokens": input_tokens, "output_tokens": 0} } }), ); @@ -290,16 +334,47 @@ fn extract_search_query(req: &MessagesRequest) -> Option { } fn search_input(req: &MessagesRequest) -> Option { - let items: Vec = req + let mut messages: Vec<(&str, String)> = req .messages .iter() + .filter(|message| matches!(message.role.as_str(), "user" | "assistant")) .filter_map(|message| { let text = content_text(&message.content); + (!text.is_empty()).then_some((message.role.as_str(), text)) + }) + .collect(); + let latest_user = messages.iter().rposition(|(role, _)| *role == "user")?; + messages.truncate(latest_user + 1); + let first_user = messages + .iter() + .enumerate() + .rev() + .filter(|(_, (role, _))| *role == "user") + .take(SEARCH_USER_CONTEXT_MESSAGES) + .last() + .map(|(index, _)| index) + .unwrap_or(latest_user); + messages.drain(..first_user); + + let mut assistant_budget = SEARCH_ASSISTANT_CONTEXT_TOKEN_BUDGET; + let items: Vec = messages + .into_iter() + .filter_map(|(role, text)| { + let (content_type, text) = if role == "assistant" { + if assistant_budget == 0 { + return None; + } + let text = truncate_to_approx_tokens(&text, assistant_budget); + assistant_budget = assistant_budget.saturating_sub(approx_token_count(&text)); + ("output_text", text) + } else { + ("input_text", text) + }; (!text.is_empty()).then(|| { json!({ "type": "message", - "role": message.role, - "content": [{"type": "input_text", "text": text}] + "role": role, + "content": [{"type": content_type, "text": text}] }) }) }) @@ -307,6 +382,33 @@ fn search_input(req: &MessagesRequest) -> Option { (!items.is_empty()).then_some(Value::Array(items)) } +fn truncate_to_approx_tokens(text: &str, max_tokens: u64) -> String { + if approx_token_count(text) <= max_tokens { + return text.to_string(); + } + + let mut tokens = 0_u64; + let mut in_word = false; + let mut end = 0; + for (index, ch) in text.char_indices() { + let word_char = ch.is_alphanumeric() || ch == '-' || ch == '_'; + let starts_token = if word_char { + !in_word + } else { + !ch.is_whitespace() + }; + if starts_token { + if tokens == max_tokens { + break; + } + tokens += 1; + } + in_word = word_char; + end = index + ch.len_utf8(); + } + text[..end].to_string() +} + fn content_text(content: &Value) -> String { match content { Value::String(text) => text.clone(), @@ -365,16 +467,6 @@ fn search_results(response: &SearchResponse) -> Vec { url: url.to_string(), }); } - if results.is_empty() { - for result in extract_web_search_results_from_text(&response.output) { - if seen.insert(result.url.clone()) { - results.push(SearchResult { - title: result.title, - url: result.url, - }); - } - } - } results } @@ -490,6 +582,84 @@ mod tests { assert!(!is_standalone_search_request(&automatic)); } + #[test] + fn search_input_uses_role_specific_content_and_recent_context() { + let mut req = request(); + req.messages = serde_json::from_value(json!([ + {"role": "user", "content": "old user"}, + {"role": "assistant", "content": "old assistant"}, + {"role": "user", "content": "previous user"}, + {"role": "assistant", "content": "previous assistant"}, + { + "role": "user", + "content": "Perform a web search for the query: current query" + }, + {"role": "assistant", "content": "content after latest user"} + ])) + .unwrap(); + + let (search, _) = build_search_request(&req, "gpt-5.6-luna", None).unwrap(); + let input = search.input.unwrap(); + let items = input.as_array().unwrap(); + assert_eq!(items.len(), 3); + assert_eq!(items[0]["content"][0]["text"], "previous user"); + assert_eq!(items[0]["content"][0]["type"], "input_text"); + assert_eq!(items[1]["content"][0]["text"], "previous assistant"); + assert_eq!(items[1]["content"][0]["type"], "output_text"); + assert_eq!(items[2]["content"][0]["type"], "input_text"); + } + + #[test] + fn search_input_bounds_assistant_context() { + let mut req = request(); + let long_assistant = (0..2_000) + .map(|index| format!("word{index}")) + .collect::>() + .join(" "); + req.messages = serde_json::from_value(json!([ + {"role": "user", "content": "previous user"}, + {"role": "assistant", "content": long_assistant}, + { + "role": "user", + "content": "Perform a web search for the query: current query" + } + ])) + .unwrap(); + + let (search, _) = build_search_request(&req, "gpt-5.6-luna", None).unwrap(); + let assistant = search.input.unwrap()[1]["content"][0]["text"] + .as_str() + .unwrap() + .to_string(); + assert!(approx_token_count(&assistant) <= SEARCH_ASSISTANT_CONTEXT_TOKEN_BUDGET); + assert!(!assistant.contains("word1999")); + } + + #[test] + fn missing_structured_results_does_not_infer_urls_from_output() { + let response = SearchResponse { + encrypted_output: None, + output: "Result from https://github.com with an embedded https://example.com link" + .to_string(), + results: None, + }; + + assert!(search_results(&response).is_empty()); + } + + #[test] + fn standalone_usage_estimates_are_nonzero() { + let (request, _) = build_search_request(&request(), "gpt-5.6-luna", None).unwrap(); + let response = SearchResponse { + encrypted_output: None, + output: "search output".to_string(), + results: None, + }; + + assert!(search_request_input_tokens(&request) > 0); + assert!(search_response_output_tokens(&response) > 0); + } + #[test] fn streamed_response_matches_claude_server_tool_shape() { let response = SearchResponse { @@ -508,6 +678,7 @@ mod tests { "msg_test", "claude-haiku-4-5-20251001", true, + 12, None, ); let runtime = tokio::runtime::Runtime::new().unwrap(); diff --git a/src/providers/codex/translate/web_search_compat.rs b/src/providers/codex/translate/web_search_compat.rs index d90d7866..eb44df17 100644 --- a/src/providers/codex/translate/web_search_compat.rs +++ b/src/providers/codex/translate/web_search_compat.rs @@ -37,7 +37,7 @@ pub fn server_tool_use_id_from_codex_web_search_id(id: &str) -> String { format!("srvtoolu_{suffix}") } -pub(crate) fn extract_web_search_results_from_text(text: &str) -> Vec { +fn extract_web_search_results_from_text(text: &str) -> Vec { let mut results: Vec = Vec::new(); let mut seen_urls: std::collections::HashSet = std::collections::HashSet::new(); @@ -62,7 +62,7 @@ pub(crate) fn extract_web_search_results_from_text(text: &str) -> Vec()|"'\\]+"#).unwrap(); + let re2 = regex_lite::Regex::new(r"https?://[^\s<>()|]+").unwrap(); for cap in re2.captures_iter(text) { let raw_url = cap.get(0).map(|m| m.as_str()).unwrap_or(""); let url = clean_url(raw_url); @@ -123,9 +123,6 @@ fn clean_url(value: &str) -> String { || out.ends_with(':') || out.ends_with('!') || out.ends_with('?') - || out.ends_with('"') - || out.ends_with('\'') - || out.ends_with('\\') { out.pop(); } @@ -179,14 +176,6 @@ mod tests { assert!(results.iter().any(|r| r.url == "https://other.com/page")); } - #[test] - fn extract_results_strips_escaped_quote_suffix() { - let text = r#"{\"url\":\"https://example.com/path\"}"#; - let results = extract_web_search_results_from_text(text); - assert_eq!(results.len(), 1); - assert_eq!(results[0].url, "https://example.com/path"); - } - #[test] fn build_compat_blocks() { let searches = vec![super::super::reducer::ReducerEvent::WebSearch {