use super::*;
use crate::providers::{ChatMessage, sse::sse_data};
use serde_json::json;
use std::sync::{Mutex, atomic::AtomicU64};
const MAX_SEARCH_RESPONSE_BYTES: usize = 2 * 1024 * 1024;
impl<T: HttpTransport + Send + Sync> OpenAiCodexProvider<T> {
pub(crate) fn web_search(
&self,
query: &str,
limit: usize,
cancellation: &AgentCancellation,
) -> anyhow::Result<Value> {
cancellation.check()?;
let transport = SearchTransport {
inner: &self.transport,
response: Mutex::new(String::new()),
};
let provider = OpenAiCodexProvider {
model: self.model.clone(),
access_token: self.access_token.clone(),
account_id: self.account_id.clone(),
transport,
auth_refresh: self.auth_refresh.clone(),
service_tier: self.service_tier.clone(),
};
let request = ProviderRequest::new_without_tools(
&self.model,
vec![
ChatMessage::system(format!(
"You are a concise web search assistant. Search the web and preserve URL citations. Aim for {limit} relevant sources."
)),
ChatMessage::user(query),
],
);
provider.stream_cancellable(request, cancellation, &mut |_| Ok(()))?;
cancellation.check()?;
let response = provider
.transport
.response
.into_inner()
.map_err(|_| anyhow::anyhow!("Codex search response lock poisoned"))?;
parse_search_response(&response, || cancellation.check())
}
}
struct SearchTransport<'a, T> {
inner: &'a T,
response: Mutex<String>,
}
impl<T: HttpTransport> HttpTransport for SearchTransport<'_, T> {
#[cfg(test)]
fn stream_json(
&self,
request: HttpRequest,
on_chunk: &mut dyn FnMut(&str) -> anyhow::Result<()>,
) -> anyhow::Result<()> {
self.stream_json_cancellable(request, &AgentCancellation::default(), on_chunk)
}
#[cfg(test)]
fn stream_json_cancellable(
&self,
request: HttpRequest,
cancellation: &AgentCancellation,
on_chunk: &mut dyn FnMut(&str) -> anyhow::Result<()>,
) -> anyhow::Result<()> {
self.stream_json_cancellable_with_semantic_deadline(
request,
cancellation,
&AtomicU64::new(0),
on_chunk,
)
}
fn stream_json_cancellable_with_semantic_deadline(
&self,
request: HttpRequest,
cancellation: &AgentCancellation,
semantic_deadline: &AtomicU64,
on_chunk: &mut dyn FnMut(&str) -> anyhow::Result<()>,
) -> anyhow::Result<()> {
self.stream_json_cancellable_with_response_metadata(
request,
cancellation,
semantic_deadline,
&mut |_| {},
on_chunk,
)
}
fn stream_json_cancellable_with_response_metadata(
&self,
mut request: HttpRequest,
cancellation: &AgentCancellation,
semantic_deadline: &AtomicU64,
on_response_id: &mut dyn FnMut(Option<String>),
on_chunk: &mut dyn FnMut(&str) -> anyhow::Result<()>,
) -> anyhow::Result<()> {
cancellation.check()?;
request.body["tools"] = json!([{"type":"web_search", "external_web_access":true, "search_context_size":"medium"}]);
request.body["tool_choice"] = json!("required");
request.body["parallel_tool_calls"] = json!(true);
request.body["include"] = json!([]);
let mut response = self
.response
.lock()
.map_err(|_| anyhow::anyhow!("Codex search response lock poisoned"))?;
response.clear();
self.inner.stream_json_cancellable_with_response_metadata(
request,
cancellation,
semantic_deadline,
on_response_id,
&mut |chunk| {
cancellation.check()?;
anyhow::ensure!(
response.len().saturating_add(chunk.len()) <= MAX_SEARCH_RESPONSE_BYTES,
"Codex search response exceeds byte limit"
);
response.push_str(chunk);
on_chunk(chunk)
},
)
}
}
fn next_search_event_boundary(
stream: &str,
check_cancellation: &mut impl FnMut() -> anyhow::Result<()>,
) -> anyhow::Result<Option<(usize, usize)>> {
let bytes = stream.as_bytes();
for (index, byte) in bytes.iter().enumerate() {
if index % 4096 == 0 {
check_cancellation()?;
}
if *byte != b'\n' {
continue;
}
let boundary_end = match bytes.get(index + 1..) {
Some([b'\n', ..]) => index + 2,
Some([b'\r', b'\n', ..]) => index + 3,
_ => continue,
};
let boundary_start = if index > 0 && bytes[index - 1] == b'\r' {
index - 1
} else {
index
};
return Ok(Some((boundary_start, boundary_end - boundary_start)));
}
Ok(None)
}
fn parse_search_response(
mut stream: &str,
mut check_cancellation: impl FnMut() -> anyhow::Result<()>,
) -> anyhow::Result<Value> {
check_cancellation()?;
let mut deltas = String::new();
let mut messages = Vec::new();
let mut completed = false;
while let Some((end, boundary)) = next_search_event_boundary(stream, &mut check_cancellation)? {
check_cancellation()?;
let event = &stream[..end];
stream = &stream[end + boundary..];
let Some(data) = sse_data(event) else {
continue;
};
if data == "[DONE]" {
continue;
}
let value: Value = serde_json::from_str(&data)
.map_err(|_| anyhow::anyhow!("Codex search returned malformed SSE JSON"))?;
match value["type"].as_str().unwrap_or_default() {
"response.output_text.delta" => {
if let Some(delta) = value["delta"].as_str() {
deltas.push_str(delta);
}
}
"response.output_item.done" => collect_message(&value["item"], &mut messages),
"response.completed" => {
anyhow::ensure!(
value["response"]["status"]
.as_str()
.is_none_or(|status| status == "completed"),
"Codex search did not complete successfully"
);
completed = true;
if let Some(output) = value["response"]["output"].as_array() {
let mut final_messages = Vec::new();
for item in output {
check_cancellation()?;
collect_message(item, &mut final_messages);
}
if !final_messages.is_empty() {
messages = final_messages;
}
}
}
"response.failed" | "response.incomplete" | "error" => {
anyhow::bail!("Codex search failed or was incomplete")
}
_ => {}
}
}
check_cancellation()?;
anyhow::ensure!(completed, "Codex search ended without response.completed");
let synthesis = search_synthesis(messages, deltas)?;
check_cancellation()?;
Ok(synthesis)
}
fn search_synthesis(messages: Vec<Value>, deltas: String) -> anyhow::Result<Value> {
let mut text = Vec::new();
let mut citations = Vec::new();
for message in messages {
let Some(content) = message["content"].as_array() else {
continue;
};
for part in content {
if let Some(part_text) = part["text"].as_str() {
text.push(part_text.to_owned());
}
collect_citations(part, &mut citations);
}
}
let synthesis = if text.is_empty() {
deltas
} else {
text.join("\n")
};
anyhow::ensure!(
!synthesis.trim().is_empty(),
"Codex search returned no synthesis"
);
Ok(json!({
"backend":"openai-codex",
"kind":"search_synthesis",
"synthesis":synthesis,
"citations":citations,
"note":"Model-generated synthesis, not source-page extracts. Citation indices refer to their original content parts. No cached page references; open citation URLs separately with web open (no API key required). Source publication/crawl/fetch timestamps are not supplied."
}))
}
fn collect_citations(part: &Value, citations: &mut Vec<Value>) {
let Some(annotations) = part["annotations"].as_array() else {
return;
};
for annotation in annotations {
if annotation["type"] != "url_citation" || annotation["url"].as_str().is_none() {
continue;
}
let citation = json!({
"url":annotation["url"], "title":annotation["title"],
"start_index":annotation["start_index"], "end_index":annotation["end_index"]
});
if !citations.contains(&citation) {
citations.push(citation);
}
}
}
fn collect_message(item: &Value, messages: &mut Vec<Value>) {
if item["type"] == "message" && item["role"] == "assistant" && !messages.contains(item) {
messages.push(item.clone());
}
}
#[cfg(test)]
mod tests;