use std::collections::HashSet;
use serde_json::{json, Value};
use url::Url;
use {
crate::credential_store::AppCredentialStore,
rho_providers::auth::codex_oauth::{chatgpt_plan_from_id_token, ChatGptPlan},
rho_providers::credentials::CodexTokens,
rho_providers::providers::openai::auth::{refresh_codex_token, CodexAuthSource},
rho_tools::tool::ToolError,
};
use super::{
normalize_domain_filters, read_bounded_text, recency_label, search_error, SearchBackendConfig,
SearchItem,
};
const OPENAI_SEARCH_MODEL: &str = "gpt-5.6-luna";
const CODEX_SEARCH_MODEL: &str = "gpt-5.6-terra";
#[derive(Clone)]
pub(super) enum OpenAiSearchAuth {
Codex {
tokens: CodexTokens,
source: CodexAuthSource,
},
ApiKey(String),
}
impl OpenAiSearchAuth {
fn model(&self) -> &'static str {
match self {
Self::Codex { tokens, .. } => match tokens
.id_token
.as_deref()
.map(chatgpt_plan_from_id_token)
.unwrap_or(ChatGptPlan::Unknown)
{
ChatGptPlan::Free | ChatGptPlan::Go | ChatGptPlan::Unknown => CODEX_SEARCH_MODEL,
ChatGptPlan::Plus
| ChatGptPlan::Pro
| ChatGptPlan::ProLite
| ChatGptPlan::Team
| ChatGptPlan::SelfServeBusinessUsageBased
| ChatGptPlan::Business
| ChatGptPlan::EnterpriseCbpUsageBased
| ChatGptPlan::Enterprise
| ChatGptPlan::Edu => OPENAI_SEARCH_MODEL,
},
Self::ApiKey(_) => OPENAI_SEARCH_MODEL,
}
}
fn bearer_token(&self) -> &str {
match self {
Self::Codex { tokens, .. } => &tokens.access_token,
Self::ApiKey(key) => key,
}
}
}
pub(super) async fn search(
client: &reqwest::Client,
query: &str,
num_results: usize,
recency_filter: Option<&str>,
domain_filter: Option<&[String]>,
config: &SearchBackendConfig,
auth: &OpenAiSearchAuth,
) -> Result<Vec<SearchItem>, ToolError> {
let mut auth = auth.clone();
let mut secrets = vec![auth.bearer_token().to_string()];
if let OpenAiSearchAuth::Codex { tokens, .. } = &auth {
secrets.extend(
tokens
.refresh_token
.iter()
.chain(tokens.id_token.iter())
.cloned(),
);
}
let body = openai_search_body(&auth, query, num_results, recency_filter, domain_filter);
let (mut status, mut text) =
send_openai_search_request(client, config, &auth, &body, &secret_refs(&secrets)).await?;
if status == reqwest::StatusCode::UNAUTHORIZED {
if let OpenAiSearchAuth::Codex {
tokens,
source: CodexAuthSource::Store,
} = &auth
{
if let Some(refresh_token) = tokens.refresh_token.as_deref() {
let store = AppCredentialStore;
let refreshed = refresh_codex_token(
client,
&store,
refresh_token,
CodexAuthSource::Store,
tokens,
)
.await
.map_err(|err| {
search_error(
format!("OpenAI web search token refresh failed: {err}"),
secret_refs(&secrets).as_slice(),
)
})?;
secrets.push(refreshed.access_token.clone());
auth = OpenAiSearchAuth::Codex {
tokens: refreshed,
source: CodexAuthSource::Store,
};
let body =
openai_search_body(&auth, query, num_results, recency_filter, domain_filter);
let retried = send_openai_search_request(
client,
config,
&auth,
&body,
&secret_refs(&secrets),
)
.await?;
status = retried.0;
text = retried.1;
}
}
}
if !status.is_success() {
return Err(search_error(
format!(
"OpenAI web search failed: HTTP {status}: {}",
super::redact_secrets(&text, &secret_refs(&secrets))
.chars()
.take(300)
.collect::<String>()
),
&secret_refs(&secrets),
));
}
let output = parse_openai_search_output(&text)?;
Ok(extract_openai_search_results(&output, num_results))
}
fn secret_refs(secrets: &[String]) -> Vec<&str> {
secrets.iter().map(String::as_str).collect()
}
fn openai_search_body(
auth: &OpenAiSearchAuth,
query: &str,
num_results: usize,
recency_filter: Option<&str>,
domain_filter: Option<&[String]>,
) -> Value {
json!({
"model": auth.model(),
"instructions": openai_search_instructions(num_results, recency_filter, domain_filter),
"input": [{"role": "user", "content": [{"type": "input_text", "text": query}]}],
"tools": [openai_web_search_tool(domain_filter)],
"include": ["web_search_call.action.sources"],
"store": false,
"stream": true,
"tool_choice": "required",
"parallel_tool_calls": true,
})
}
async fn send_openai_search_request(
client: &reqwest::Client,
config: &SearchBackendConfig,
auth: &OpenAiSearchAuth,
body: &Value,
secrets: &[&str],
) -> Result<(reqwest::StatusCode, String), ToolError> {
let mut request = client
.post(config.destination_url("responses")?)
.bearer_auth(auth.bearer_token())
.header("Content-Type", "application/json");
if let OpenAiSearchAuth::Codex { tokens, .. } = auth {
request = request
.header("User-Agent", "codex-cli")
.header("originator", "codex_cli_rs");
if let Some(account_id) = tokens.account_id.as_deref() {
request = request.header("ChatGPT-Account-ID", account_id);
}
}
let response =
request.json(body).send().await.map_err(|err| {
search_error(format!("OpenAI web search request failed: {err}"), secrets)
})?;
read_bounded_text(response, secrets).await
}
fn openai_search_instructions(
num_results: usize,
recency_filter: Option<&str>,
domain_filter: Option<&[String]>,
) -> String {
let mut lines = vec![
"Search the web and return concise source-backed results.".to_string(),
"Prefer clickable source citations when possible.".to_string(),
format!("Prefer around {} distinct sources.", num_results.min(20)),
];
if let Some(recency) = recency_filter.and_then(recency_label) {
lines.push(format!("Prefer sources from the {recency}."));
}
let filters = normalize_domain_filters(domain_filter);
if !filters.allowed.is_empty() {
lines.push(format!(
"Only use sources from: {}.",
filters.allowed.join(", ")
));
}
if !filters.blocked.is_empty() {
lines.push(format!(
"Do not use sources from: {}.",
filters.blocked.join(", ")
));
}
lines.join(" ")
}
fn openai_web_search_tool(domain_filter: Option<&[String]>) -> Value {
let filters = normalize_domain_filters(domain_filter);
let mut tool = serde_json::Map::from_iter([("type".into(), json!("web_search"))]);
if !filters.allowed.is_empty() || !filters.blocked.is_empty() {
tool.insert(
"filters".into(),
json!({
"allowed_domains": filters.allowed,
"blocked_domains": filters.blocked,
}),
);
}
Value::Object(tool)
}
fn parse_openai_search_output(text: &str) -> Result<Vec<Value>, ToolError> {
let trimmed = text.trim();
if trimmed.starts_with('{') || trimmed.starts_with('[') {
let value: Value = serde_json::from_str(trimmed).map_err(|err| {
ToolError::Message(format!("OpenAI web search returned invalid JSON: {err}"))
})?;
return Ok(value
.get("output")
.and_then(Value::as_array)
.cloned()
.unwrap_or_default());
}
let mut output = Vec::new();
let mut completed_output = None;
for line in text.lines() {
let Some(data) = line.strip_prefix("data: ") else {
continue;
};
let data = data.trim();
if data.is_empty() || data == "[DONE]" {
continue;
}
let Ok(value) = serde_json::from_str::<Value>(data) else {
continue;
};
if value.get("type").and_then(Value::as_str) == Some("response.output_item.done") {
if let Some(item) = value.get("item") {
output.push(item.clone());
}
}
if matches!(
value.get("type").and_then(Value::as_str),
Some("response.done" | "response.completed")
) {
completed_output = value
.get("response")
.and_then(|response| response.get("output"))
.and_then(Value::as_array)
.cloned();
}
}
Ok(completed_output.unwrap_or(output))
}
fn extract_openai_search_results(output: &[Value], num_results: usize) -> Vec<SearchItem> {
let mut results = Vec::new();
let mut seen = HashSet::new();
for item in output {
if item.get("type").and_then(Value::as_str) != Some("message") {
continue;
}
for part in item
.get("content")
.and_then(Value::as_array)
.into_iter()
.flatten()
{
let text = part.get("text").and_then(Value::as_str).unwrap_or_default();
for annotation in part
.get("annotations")
.and_then(Value::as_array)
.into_iter()
.flatten()
{
if annotation.get("type").and_then(Value::as_str) != Some("url_citation") {
continue;
}
add_openai_search_result(
&mut results,
&mut seen,
annotation.get("url").and_then(Value::as_str),
annotation.get("title").and_then(Value::as_str),
citation_snippet(
text,
annotation.get("start_index").and_then(Value::as_u64),
annotation.get("end_index").and_then(Value::as_u64),
),
);
}
}
}
for item in output {
if item.get("type").and_then(Value::as_str) != Some("web_search_call") {
continue;
}
for group in [
item.get("action")
.and_then(|action| action.get("sources"))
.and_then(Value::as_array),
item.get("sources").and_then(Value::as_array),
item.get("results").and_then(Value::as_array),
]
.into_iter()
.flatten()
{
for source in group {
add_openai_search_result(
&mut results,
&mut seen,
source
.get("url")
.or_else(|| source.get("source_website_url"))
.and_then(Value::as_str),
source
.get("title")
.or_else(|| source.get("caption"))
.and_then(Value::as_str),
String::new(),
);
}
}
}
results.truncate(num_results.min(20));
results
}
fn add_openai_search_result(
results: &mut Vec<SearchItem>,
seen: &mut HashSet<String>,
url: Option<&str>,
title: Option<&str>,
snippet: String,
) {
let Some(url) = url.filter(|url| !url.trim().is_empty()) else {
return;
};
let url = clean_openai_source_url(url);
if !seen.insert(url.clone()) {
return;
}
results.push(SearchItem {
title: title
.filter(|title| !title.trim().is_empty())
.map(str::to_string)
.or_else(|| Some(url.clone())),
url: Some(url),
snippet,
});
}
fn clean_openai_source_url(raw_url: &str) -> String {
Url::parse(raw_url)
.map(|mut url| {
let query = url
.query_pairs()
.filter(|(key, value)| !(key == "utm_source" && value == "openai"))
.map(|(key, value)| (key.into_owned(), value.into_owned()))
.collect::<Vec<_>>();
url.set_query(None);
if !query.is_empty() {
url.query_pairs_mut().extend_pairs(query);
}
url.to_string()
})
.unwrap_or_else(|_| raw_url.replace("?utm_source=openai", ""))
}
fn citation_snippet(text: &str, start: Option<u64>, end: Option<u64>) -> String {
let (Some(start), Some(end)) = (start, end) else {
return String::new();
};
let start = start as usize;
let end = end.max(start as u64) as usize;
let before = start.saturating_sub(100);
let after = end.saturating_add(100);
text.chars()
.skip(before)
.take(after.saturating_sub(before))
.collect::<String>()
.replace(['[', ']', '(', ')'], "")
.trim()
.chars()
.take(300)
.collect()
}