1use std::collections::HashSet;
2
3use serde::{Deserialize, Serialize};
4use serde_json::{Value, json};
5
6use crate::anthropic::schema::MessagesRequest;
7use crate::anthropic::sse::encode_sse_event;
8use crate::traffic::TrafficCapture;
9
10use super::count_tokens::{approx_token_count, truncate_to_token_budget};
11
12const SEARCH_OUTPUT_TOKEN_BUDGET: u64 = 2_500;
13const SEARCH_ASSISTANT_CONTEXT_TOKEN_BUDGET: u64 = 1_000;
14const SEARCH_USER_CONTEXT_MESSAGES: usize = 2;
15const CLAUDE_SEARCH_PROMPT_PREFIX: &str = "Perform a web search for the query:";
16
17#[derive(Debug, Clone, Serialize, PartialEq)]
18pub struct SearchRequest {
19 pub id: String,
20 pub model: String,
21 #[serde(skip_serializing_if = "Option::is_none")]
22 pub reasoning: Option<Value>,
23 #[serde(skip_serializing_if = "Option::is_none")]
24 pub input: Option<Value>,
25 pub commands: SearchCommands,
26 pub settings: SearchSettings,
27 pub max_output_tokens: u64,
28}
29
30#[derive(Debug, Clone, Serialize, PartialEq)]
31pub struct SearchCommands {
32 pub search_query: Vec<SearchQuery>,
33}
34
35#[derive(Debug, Clone, Serialize, PartialEq)]
36pub struct SearchQuery {
37 pub q: String,
38}
39
40#[derive(Debug, Clone, Serialize, PartialEq)]
41pub struct SearchSettings {
42 #[serde(skip_serializing_if = "Option::is_none")]
43 pub filters: Option<SearchFilters>,
44 pub allowed_callers: Vec<&'static str>,
45 pub external_web_access: bool,
46}
47
48#[derive(Debug, Clone, Serialize, PartialEq)]
49pub struct SearchFilters {
50 #[serde(skip_serializing_if = "Option::is_none")]
51 pub allowed_domains: Option<Vec<String>>,
52 #[serde(skip_serializing_if = "Option::is_none")]
53 pub blocked_domains: Option<Vec<String>>,
54}
55
56#[derive(Debug, Clone, Deserialize, PartialEq, Eq)]
57pub struct SearchResponse {
58 pub encrypted_output: Option<String>,
59 pub output: String,
60 #[serde(default)]
61 pub results: Option<Vec<Value>>,
62}
63
64#[derive(Debug, Clone, PartialEq, Eq)]
65struct SearchResult {
66 title: String,
67 url: String,
68}
69
70pub fn is_standalone_search_request(req: &MessagesRequest) -> bool {
71 let Some(choice) = req.extra.get("tool_choice").and_then(Value::as_object) else {
72 return false;
73 };
74 if choice.get("type").and_then(Value::as_str) != Some("tool") {
75 return false;
76 }
77 let Some(selected_name) = choice.get("name").and_then(Value::as_str) else {
78 return false;
79 };
80 req.extra
81 .get("tools")
82 .and_then(Value::as_array)
83 .is_some_and(|tools| {
84 tools.iter().any(|tool| {
85 tool.get("type").and_then(Value::as_str) == Some("web_search_20250305")
86 && tool.get("name").and_then(Value::as_str) == Some(selected_name)
87 })
88 })
89}
90
91pub fn build_search_request(
92 req: &MessagesRequest,
93 model: &str,
94 session_id: Option<&str>,
95) -> Result<(SearchRequest, String), anyhow::Error> {
96 let query = extract_search_query(req)
97 .ok_or_else(|| anyhow::anyhow!("web_search request does not contain a text query"))?;
98 let input = search_input(req);
99 let filters = search_filters(req);
100 let id = session_id
101 .map(str::to_owned)
102 .unwrap_or_else(|| format!("search-{}", uuid::Uuid::new_v4()));
103
104 Ok((
105 SearchRequest {
106 id,
107 model: model.to_string(),
108 reasoning: None,
109 input,
110 commands: SearchCommands {
111 search_query: vec![SearchQuery { q: query.clone() }],
112 },
113 settings: SearchSettings {
114 filters,
115 allowed_callers: vec!["direct"],
116 external_web_access: true,
117 },
118 max_output_tokens: SEARCH_OUTPUT_TOKEN_BUDGET,
119 },
120 query,
121 ))
122}
123
124pub fn search_request_input_tokens(request: &SearchRequest) -> u64 {
125 let mut tokens = approx_token_count(&request.model);
126 tokens += request
127 .commands
128 .search_query
129 .iter()
130 .map(|query| approx_token_count(&query.q))
131 .sum::<u64>();
132 tokens += request.input.as_ref().map(value_text_tokens).unwrap_or(0);
133 if let Some(filters) = &request.settings.filters {
134 tokens += filters
135 .allowed_domains
136 .iter()
137 .flatten()
138 .chain(filters.blocked_domains.iter().flatten())
139 .map(|domain| approx_token_count(domain))
140 .sum::<u64>();
141 }
142 tokens.max(1)
143}
144
145pub fn search_response_output_tokens(response: &SearchResponse) -> u64 {
146 (approx_token_count(&response.output)
147 + response
148 .results
149 .as_ref()
150 .map(|results| results.iter().map(value_text_tokens).sum())
151 .unwrap_or(0))
152 .max(1)
153}
154
155fn value_text_tokens(value: &Value) -> u64 {
156 match value {
157 Value::String(text) => approx_token_count(text),
158 Value::Array(values) => values.iter().map(value_text_tokens).sum(),
159 Value::Object(values) => values.values().map(value_text_tokens).sum(),
160 _ => 0,
161 }
162}
163
164pub fn anthropic_search_response(
165 response: &SearchResponse,
166 query: &str,
167 message_id: &str,
168 model: &str,
169 stream: bool,
170 input_tokens: u64,
171 traffic: Option<&TrafficCapture>,
172) -> axum::response::Response {
173 use axum::response::IntoResponse;
174
175 let tool_use_id = format!("srvtoolu_ws_{}", uuid::Uuid::new_v4().simple());
176 let results = search_results(response);
177 let content = response_content(response, query, &tool_use_id, &results);
178 let output_tokens = search_response_output_tokens(response);
179 let usage = json!({
180 "input_tokens": input_tokens,
181 "output_tokens": output_tokens,
182 "cache_creation_input_tokens": 0,
183 "cache_read_input_tokens": 0,
184 "server_tool_use": {"web_search_requests": 1}
185 });
186
187 if !stream {
188 return (
189 http::StatusCode::OK,
190 axum::Json(json!({
191 "id": message_id,
192 "type": "message",
193 "role": "assistant",
194 "model": model,
195 "content": content,
196 "stop_reason": "end_turn",
197 "stop_sequence": null,
198 "usage": usage,
199 })),
200 )
201 .into_response();
202 }
203
204 let mut body = Vec::new();
205 emit(
206 &mut body,
207 traffic,
208 "message_start",
209 &json!({
210 "type": "message_start",
211 "message": {
212 "id": message_id,
213 "type": "message",
214 "role": "assistant",
215 "model": model,
216 "content": [],
217 "stop_reason": null,
218 "stop_sequence": null,
219 "usage": {"input_tokens": input_tokens, "output_tokens": 0}
220 }
221 }),
222 );
223 emit(
224 &mut body,
225 traffic,
226 "content_block_start",
227 &json!({
228 "type": "content_block_start",
229 "index": 0,
230 "content_block": {
231 "type": "server_tool_use",
232 "id": tool_use_id,
233 "name": "web_search",
234 "input": {}
235 }
236 }),
237 );
238 emit(
239 &mut body,
240 traffic,
241 "content_block_delta",
242 &json!({
243 "type": "content_block_delta",
244 "index": 0,
245 "delta": {
246 "type": "input_json_delta",
247 "partial_json": json!({"query": query}).to_string()
248 }
249 }),
250 );
251 emit_block_stop(&mut body, traffic, 0);
252 emit(
253 &mut body,
254 traffic,
255 "content_block_start",
256 &json!({
257 "type": "content_block_start",
258 "index": 1,
259 "content_block": {
260 "type": "web_search_tool_result",
261 "tool_use_id": tool_use_id,
262 "content": web_search_result_values(&results)
263 }
264 }),
265 );
266 emit_block_stop(&mut body, traffic, 1);
267 if !response.output.is_empty() {
268 emit(
269 &mut body,
270 traffic,
271 "content_block_start",
272 &json!({
273 "type": "content_block_start",
274 "index": 2,
275 "content_block": {"type": "text", "text": ""}
276 }),
277 );
278 emit(
279 &mut body,
280 traffic,
281 "content_block_delta",
282 &json!({
283 "type": "content_block_delta",
284 "index": 2,
285 "delta": {"type": "text_delta", "text": response.output}
286 }),
287 );
288 emit_block_stop(&mut body, traffic, 2);
289 }
290 emit(
291 &mut body,
292 traffic,
293 "message_delta",
294 &json!({
295 "type": "message_delta",
296 "delta": {"stop_reason": "end_turn", "stop_sequence": null},
297 "usage": usage
298 }),
299 );
300 emit(
301 &mut body,
302 traffic,
303 "message_stop",
304 &json!({"type": "message_stop"}),
305 );
306
307 let headers = [
308 (http::header::CONTENT_TYPE, "text/event-stream"),
309 (http::header::CACHE_CONTROL, "no-cache"),
310 (http::header::CONNECTION, "keep-alive"),
311 ];
312 (headers, body).into_response()
313}
314
315fn extract_search_query(req: &MessagesRequest) -> Option<String> {
316 req.messages
317 .iter()
318 .rev()
319 .filter(|message| message.role == "user")
320 .find_map(|message| {
321 let text = content_text(&message.content);
322 let text = text.trim();
323 if text.is_empty() {
324 return None;
325 }
326 Some(
327 text.strip_prefix(CLAUDE_SEARCH_PROMPT_PREFIX)
328 .map(str::trim)
329 .filter(|query| !query.is_empty())
330 .unwrap_or(text)
331 .to_string(),
332 )
333 })
334}
335
336fn search_input(req: &MessagesRequest) -> Option<Value> {
337 let mut messages: Vec<(&str, String)> = req
338 .messages
339 .iter()
340 .filter(|message| matches!(message.role.as_str(), "user" | "assistant"))
341 .filter_map(|message| {
342 let text = content_text(&message.content);
343 (!text.is_empty()).then_some((message.role.as_str(), text))
344 })
345 .collect();
346 let latest_user = messages.iter().rposition(|(role, _)| *role == "user")?;
347 messages.truncate(latest_user + 1);
348 let first_user = messages
349 .iter()
350 .enumerate()
351 .rev()
352 .filter(|(_, (role, _))| *role == "user")
353 .take(SEARCH_USER_CONTEXT_MESSAGES)
354 .last()
355 .map(|(index, _)| index)
356 .unwrap_or(latest_user);
357 messages.drain(..first_user);
358
359 let mut assistant_budget = SEARCH_ASSISTANT_CONTEXT_TOKEN_BUDGET;
360 let items: Vec<Value> = messages
361 .into_iter()
362 .filter_map(|(role, text)| {
363 let (content_type, text) = if role == "assistant" {
364 if assistant_budget == 0 {
365 return None;
366 }
367 let text = truncate_to_token_budget(&text, assistant_budget);
368 assistant_budget = assistant_budget.saturating_sub(approx_token_count(&text));
369 ("output_text", text)
370 } else {
371 ("input_text", text)
372 };
373 (!text.is_empty()).then(|| {
374 json!({
375 "type": "message",
376 "role": role,
377 "content": [{"type": content_type, "text": text}]
378 })
379 })
380 })
381 .collect();
382 (!items.is_empty()).then_some(Value::Array(items))
383}
384
385fn content_text(content: &Value) -> String {
386 match content {
387 Value::String(text) => text.clone(),
388 Value::Array(blocks) => blocks
389 .iter()
390 .filter(|block| block.get("type").and_then(Value::as_str) == Some("text"))
391 .filter_map(|block| block.get("text").and_then(Value::as_str))
392 .collect::<Vec<_>>()
393 .join("\n"),
394 _ => String::new(),
395 }
396}
397
398fn search_filters(req: &MessagesRequest) -> Option<SearchFilters> {
399 let tool = req
400 .extra
401 .get("tools")?
402 .as_array()?
403 .iter()
404 .find(|tool| tool.get("type").and_then(Value::as_str) == Some("web_search_20250305"))?;
405 let allowed_domains = string_array(tool.get("allowed_domains"));
406 let blocked_domains = string_array(tool.get("blocked_domains"));
407 (allowed_domains.is_some() || blocked_domains.is_some()).then_some(SearchFilters {
408 allowed_domains,
409 blocked_domains,
410 })
411}
412
413fn string_array(value: Option<&Value>) -> Option<Vec<String>> {
414 let values = value?.as_array()?;
415 let values: Vec<String> = values
416 .iter()
417 .filter_map(Value::as_str)
418 .map(str::to_owned)
419 .collect();
420 (!values.is_empty()).then_some(values)
421}
422
423fn search_results(response: &SearchResponse) -> Vec<SearchResult> {
424 let mut results = Vec::new();
425 let mut seen = HashSet::new();
426 for result in response.results.iter().flatten() {
427 let Some(url) = result.get("url").and_then(Value::as_str) else {
428 continue;
429 };
430 if !seen.insert(url.to_string()) {
431 continue;
432 }
433 let title = result
434 .get("title")
435 .and_then(Value::as_str)
436 .or_else(|| result.get("ref_id").and_then(Value::as_str))
437 .unwrap_or(url);
438 results.push(SearchResult {
439 title: title.to_string(),
440 url: url.to_string(),
441 });
442 }
443 results
444}
445
446fn response_content(
447 response: &SearchResponse,
448 query: &str,
449 tool_use_id: &str,
450 results: &[SearchResult],
451) -> Vec<Value> {
452 let mut content = vec![
453 json!({
454 "type": "server_tool_use",
455 "id": tool_use_id,
456 "name": "web_search",
457 "input": {"query": query}
458 }),
459 json!({
460 "type": "web_search_tool_result",
461 "tool_use_id": tool_use_id,
462 "content": web_search_result_values(results)
463 }),
464 ];
465 if !response.output.is_empty() {
466 content.push(json!({"type": "text", "text": response.output}));
467 }
468 content
469}
470
471fn web_search_result_values(results: &[SearchResult]) -> Vec<Value> {
472 results
473 .iter()
474 .map(|result| {
475 json!({
476 "type": "web_search_result",
477 "title": result.title,
478 "url": result.url,
479 })
480 })
481 .collect()
482}
483
484fn emit_block_stop(out: &mut Vec<u8>, traffic: Option<&TrafficCapture>, index: usize) {
485 emit(
486 out,
487 traffic,
488 "content_block_stop",
489 &json!({"type": "content_block_stop", "index": index}),
490 );
491}
492
493fn emit(out: &mut Vec<u8>, traffic: Option<&TrafficCapture>, event: &str, data: &Value) {
494 if let Some(traffic) = traffic {
495 traffic.write_json_event(
496 "050-downstream-event",
497 &json!({"event": event, "data": data}),
498 );
499 }
500 out.extend_from_slice(&encode_sse_event(Some(event), &data.to_string()));
501}
502
503#[cfg(test)]
504mod tests {
505 use super::*;
506 use crate::anthropic::sse::parse_sse_events;
507
508 fn request() -> MessagesRequest {
509 serde_json::from_value(json!({
510 "model": "claude-haiku-4-5-20251001",
511 "max_tokens": 32000,
512 "stream": true,
513 "messages": [{
514 "role": "user",
515 "content": [{
516 "type": "text",
517 "text": "Perform a web search for the query: Codex standalone search"
518 }]
519 }],
520 "tools": [{
521 "type": "web_search_20250305",
522 "name": "web_search",
523 "allowed_domains": ["openai.com"]
524 }],
525 "tool_choice": {"type": "tool", "name": "web_search"}
526 }))
527 .unwrap()
528 }
529
530 #[test]
531 fn request_preserves_luna_and_omits_reasoning() {
532 let (request, query) =
533 build_search_request(&request(), "gpt-5.6-luna", Some("session-1")).unwrap();
534 assert_eq!(request.model, "gpt-5.6-luna");
535 assert_eq!(request.reasoning, None);
536 assert_eq!(request.id, "session-1");
537 assert_eq!(request.max_output_tokens, 2_500);
538 assert_eq!(query, "Codex standalone search");
539 assert_eq!(request.commands.search_query[0].q, query);
540 assert_eq!(
541 request.settings.filters.unwrap().allowed_domains,
542 Some(vec!["openai.com".to_string()])
543 );
544 }
545
546 #[test]
547 fn only_forced_claude_search_uses_standalone_endpoint() {
548 let forced = request();
549 assert!(is_standalone_search_request(&forced));
550
551 let mut automatic = forced.clone();
552 automatic
553 .extra
554 .insert("tool_choice".to_string(), json!({"type": "auto"}));
555 assert!(!is_standalone_search_request(&automatic));
556 }
557
558 #[test]
559 fn search_input_uses_role_specific_content_and_recent_context() {
560 let mut req = request();
561 req.messages = serde_json::from_value(json!([
562 {"role": "user", "content": "old user"},
563 {"role": "assistant", "content": "old assistant"},
564 {"role": "user", "content": "previous user"},
565 {"role": "assistant", "content": "previous assistant"},
566 {
567 "role": "user",
568 "content": "Perform a web search for the query: current query"
569 },
570 {"role": "assistant", "content": "content after latest user"}
571 ]))
572 .unwrap();
573
574 let (search, _) = build_search_request(&req, "gpt-5.6-luna", None).unwrap();
575 let input = search.input.unwrap();
576 let items = input.as_array().unwrap();
577 assert_eq!(items.len(), 3);
578 assert_eq!(items[0]["content"][0]["text"], "previous user");
579 assert_eq!(items[0]["content"][0]["type"], "input_text");
580 assert_eq!(items[1]["content"][0]["text"], "previous assistant");
581 assert_eq!(items[1]["content"][0]["type"], "output_text");
582 assert_eq!(items[2]["content"][0]["type"], "input_text");
583 }
584
585 #[test]
586 fn search_input_bounds_assistant_context() {
587 let mut req = request();
588 let long_assistant = (0..2_000)
589 .map(|index| format!("word{index}"))
590 .collect::<Vec<_>>()
591 .join(" ");
592 req.messages = serde_json::from_value(json!([
593 {"role": "user", "content": "previous user"},
594 {"role": "assistant", "content": long_assistant},
595 {
596 "role": "user",
597 "content": "Perform a web search for the query: current query"
598 }
599 ]))
600 .unwrap();
601
602 let (search, _) = build_search_request(&req, "gpt-5.6-luna", None).unwrap();
603 let assistant = search.input.unwrap()[1]["content"][0]["text"]
604 .as_str()
605 .unwrap()
606 .to_string();
607 assert!(approx_token_count(&assistant) <= SEARCH_ASSISTANT_CONTEXT_TOKEN_BUDGET);
608 assert!(!assistant.contains("word1999"));
609 }
610
611 #[test]
612 fn missing_structured_results_does_not_infer_urls_from_output() {
613 let response = SearchResponse {
614 encrypted_output: None,
615 output: "Result from https://github.com with an embedded https://example.com link"
616 .to_string(),
617 results: None,
618 };
619
620 assert!(search_results(&response).is_empty());
621 }
622
623 #[test]
624 fn standalone_usage_estimates_are_nonzero() {
625 let (request, _) = build_search_request(&request(), "gpt-5.6-luna", None).unwrap();
626 let response = SearchResponse {
627 encrypted_output: None,
628 output: "search output".to_string(),
629 results: None,
630 };
631
632 assert!(search_request_input_tokens(&request) > 0);
633 assert!(search_response_output_tokens(&response) > 0);
634 }
635
636 #[test]
637 fn streamed_response_matches_claude_server_tool_shape() {
638 let response = SearchResponse {
639 encrypted_output: Some("opaque".to_string()),
640 output: "See [OpenAI](https://openai.com).".to_string(),
641 results: Some(vec![json!({
642 "type": "text_result",
643 "ref_id": "turn0search0",
644 "url": "https://openai.com",
645 "title": "OpenAI"
646 })]),
647 };
648 let response = anthropic_search_response(
649 &response,
650 "Codex standalone search",
651 "msg_test",
652 "claude-haiku-4-5-20251001",
653 true,
654 12,
655 None,
656 );
657 let runtime = tokio::runtime::Runtime::new().unwrap();
658 let body = runtime.block_on(axum::body::to_bytes(response.into_body(), usize::MAX));
659 let events = parse_sse_events(&body.unwrap());
660 let payloads: Vec<Value> = events
661 .iter()
662 .filter_map(|event| serde_json::from_str(&event.data).ok())
663 .collect();
664 assert!(payloads.iter().any(|payload| {
665 payload
666 .pointer("/content_block/type")
667 .and_then(Value::as_str)
668 == Some("server_tool_use")
669 }));
670 assert!(payloads.iter().any(|payload| {
671 payload
672 .pointer("/content_block/type")
673 .and_then(Value::as_str)
674 == Some("web_search_tool_result")
675 }));
676 assert!(payloads.iter().any(|payload| {
677 payload
678 .pointer("/usage/server_tool_use/web_search_requests")
679 .and_then(Value::as_u64)
680 == Some(1)
681 }));
682 }
683}