From 66010db5b78c09e08ca45a30ab7890bc5881134c 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 --- README.md | 16 +- src/providers/codex/client.rs | 295 ++++++++++ src/providers/codex/mod.rs | 79 +++ src/providers/codex/search.rs | 539 ++++++++++++++++++ .../codex/translate/web_search_compat.rs | 15 +- 5 files changed, 935 insertions(+), 9 deletions(-) create mode 100644 src/providers/codex/search.rs diff --git a/README.md b/README.md index c73cc0d..6aae8af 100644 --- a/README.md +++ b/README.md @@ -279,13 +279,15 @@ so simple prompts can emit no thinking block. Set `codex.reasoningSummary` / `CCP_CODEX_REASONING_SUMMARY` to `off` or `none` to suppress summaries while keeping `reasoning.effort` and encrypted continuation content. -Claude Code's hosted `web_search_20250305` tool is translated to Codex's native -Responses `web_search` tool with live external web access and non-empty native -domain filters. Forced searches use Codex's required `allowed_tools` form so -structured filters remain active. Codex does not expose a hosted-search limit, -so the Claude tool's `max_uses` value is not enforced. Codex hosted search calls -are emitted back to Claude Code as Anthropic `server_tool_use` and -`web_search_tool_result` blocks with +Claude Code's forced `web_search_20250305` subrequest is translated to Codex's +standalone `/alpha/search` endpoint, matching the current Codex CLI. The proxy +passes the resolved request model through unchanged, sends no search reasoning, +and preserves non-empty domain filters, so a Haiku/Luna search no longer needs +an extra Sol Responses turn. Automatic hosted-search requests stay on the full +Responses API because the standalone endpoint cannot decide whether to invoke a +tool. Codex does not expose a standalone-search limit, so the Claude tool's +`max_uses` value is not enforced. Search results are emitted back to Claude Code +as Anthropic `server_tool_use` and `web_search_tool_result` blocks with `usage.server_tool_use.web_search_requests` so Claude Code can account for completed searches. diff --git a/src/providers/codex/client.rs b/src/providers/codex/client.rs index 4995245..771de4f 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; // --------------------------------------------------------------------------- @@ -155,6 +156,27 @@ pub fn build_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, @@ -165,6 +187,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 // --------------------------------------------------------------------------- @@ -296,6 +326,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, @@ -934,6 +1032,106 @@ impl CodexHttpClient { headers, }) } + + 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( @@ -1358,6 +1556,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 buffered_http_retries_retryable_status() { let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); @@ -1395,6 +1619,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 ead8d11..e6e4a4c 100644 --- a/src/providers/codex/mod.rs +++ b/src/providers/codex/mod.rs @@ -4,6 +4,7 @@ pub mod continuation; pub mod count_tokens; pub(crate) mod events; pub mod request_summary; +pub mod search; pub mod translate; pub mod websocket; @@ -14,10 +15,12 @@ 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}; use crate::config; +use crate::logging::create_logger; use crate::monitor::usage_from_anthropic_sse; use crate::provider::{CliHandlers, Provider, RequestContext}; use crate::registry; @@ -104,6 +107,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 0000000..9fd6c0c --- /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 eb44df1..d90d786 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 00777a15be01ef4d36aceb50e0e35238b42c6393 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. --- README.md | 6 +- 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 +- 5 files changed, 199 insertions(+), 34 deletions(-) diff --git a/README.md b/README.md index 6aae8af..1d9108e 100644 --- a/README.md +++ b/README.md @@ -286,8 +286,10 @@ and preserves non-empty domain filters, so a Haiku/Luna search no longer needs an extra Sol Responses turn. Automatic hosted-search requests stay on the full Responses API because the standalone endpoint cannot decide whether to invoke a tool. Codex does not expose a standalone-search limit, so the Claude tool's -`max_uses` value is not enforced. Search results are emitted back to Claude Code -as Anthropic `server_tool_use` and `web_search_tool_result` blocks with +`max_uses` value is not enforced. Structured result DTOs from Codex are emitted +as Anthropic `server_tool_use` and `web_search_tool_result` blocks, while the +standalone tool output is returned as text. The proxy locally estimates input +and output tokens and reports `usage.server_tool_use.web_search_requests` so Claude Code can account for completed searches. diff --git a/src/providers/codex/count_tokens.rs b/src/providers/codex/count_tokens.rs index 4ef09a7..93420d4 100644 --- a/src/providers/codex/count_tokens.rs +++ b/src/providers/codex/count_tokens.rs @@ -94,7 +94,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 e6e4a4c..011f94a 100644 --- a/src/providers/codex/mod.rs +++ b/src/providers/codex/mod.rs @@ -171,8 +171,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, @@ -180,6 +182,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 9fd6c0c..8a787af 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 d90d786..eb44df1 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 {