diff --git a/README.md b/README.md index c73cc0d..1d9108e 100644 --- a/README.md +++ b/README.md @@ -279,13 +279,17 @@ 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. 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/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/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 ead8d11..011f94a 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,85 @@ 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()), + ), + ])), + ); + 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(input_tokens), Some(output_tokens)); + } + return search::anthropic_search_response( + &search_response, + &query, + &message_id, + model, + want_stream, + input_tokens, + 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..8a787af --- /dev/null +++ b/src/providers/codex/search.rs @@ -0,0 +1,710 @@ +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::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)] +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 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; + + 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": input_tokens, + "output_tokens": output_tokens, + "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": input_tokens, "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 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": role, + "content": [{"type": content_type, "text": text}] + }) + }) + }) + .collect(); + (!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(), + 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(), + }); + } + 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 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 { + 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, + 12, + 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) + })); + } +}