#![forbid(unsafe_code)]
#![doc = include_str!("../Documentation.md")]
#[cfg(not(unix))]
compile_error!("kcode-k1-codex-websearch requires Unix process groups");
#[cfg(test)]
use nix::sys::signal::kill;
use nix::{
sys::signal::{Signal, killpg},
unistd::Pid,
};
use serde_json::Value;
use std::{
ffi::OsString,
io,
path::PathBuf,
time::{Duration, Instant},
};
use tokio::{
io::{AsyncRead, AsyncReadExt, AsyncWriteExt},
process::Command,
sync::oneshot,
};
const ARG_ERROR: &str =
"WebSearch failed: model or reasoning_effort is not representable as a process argument";
const DEADLINE_ERROR: &str = "WebSearch failed: deadline exceeded";
const COMPAT_ERROR: &str = "WebSearch failed: Codex compatibility check failed";
const EXEC_ERROR: &str = "WebSearch failed: Codex execution failed";
const JSON_ERROR: &str = "WebSearch failed: invalid Codex JSONL output";
const INCOMPLETE_ERROR: &str = "WebSearch failed: Codex response was incomplete";
const MESSAGE_ERROR: &str = "WebSearch failed: Codex returned no agent message";
const CLEANUP_RESERVE: Duration = Duration::from_secs(1);
const HELP_LIMIT: Duration = Duration::from_secs(2);
const CONFIG: &[&str] = &[
r#"web_search="live""#,
"tools.web_search=true",
"tools.view_image=false",
"apps._default.enabled=false",
"agents.enabled=false",
"features.apps=false",
"features.code_mode.enabled=false",
"features.goals=false",
"features.hooks=false",
"features.memories=false",
"features.multi_agent=false",
"features.remote_plugin=false",
"features.shell_snapshot=false",
"features.shell_tool=false",
"features.skill_mcp_dependency_install=false",
"features.unified_exec=false",
"memories.generate_memories=false",
"memories.use_memories=false",
r#"history.persistence="none""#,
"check_for_update_on_startup=false",
"feedback.enabled=false",
"analytics.enabled=false",
"allow_login_shell=false",
"skills.config=[]",
"mcp_servers={}",
"plugins={}",
"marketplaces={}",
"hooks={}",
];
const REQUIRED_FLAGS: &[&str] = &[
"--search",
"--ephemeral",
"--ignore-user-config",
"--ignore-rules",
"--json",
"--sandbox",
"--ask-for-approval",
"--skip-git-repo-check",
"--no-daemon",
"--strict-config",
"--model",
"--config",
];
#[derive(Clone)]
pub struct Runner {
executable: PathBuf,
}
pub struct Request {
pub query: String,
pub model: String,
pub reasoning_effort: String,
pub deadline: Instant,
}
impl Runner {
pub fn new(executable: PathBuf) -> Self {
Self { executable }
}
pub async fn run(&self, request: Request) -> Result<String, String> {
if request.model.contains('\0') || request.reasoning_effort.contains('\0') {
return Err(ARG_ERROR.to_owned());
}
let cutoff = request
.deadline
.checked_sub(CLEANUP_RESERVE)
.unwrap_or(request.deadline);
if Instant::now() >= cutoff {
return Err(DEADLINE_ERROR.to_owned());
}
let mut help = self.check_help(&["--help"], cutoff).await?;
help.extend(self.check_help(&["exec", "--help"], cutoff).await?);
let help = String::from_utf8_lossy(&help);
if REQUIRED_FLAGS.iter().any(|flag| !help.contains(flag)) {
return Err(COMPAT_ERROR.to_owned());
}
if Instant::now() >= cutoff {
return Err(DEADLINE_ERROR.to_owned());
}
let output = run_process(
self.executable.clone(),
exec_args(&request.model, &request.reasoning_effort),
request.query.into_bytes(),
cutoff,
)
.await
.map_err(|error| match error {
ProcError::Timeout => DEADLINE_ERROR.to_owned(),
ProcError::Io | ProcError::Cancelled => EXEC_ERROR.to_owned(),
})?;
if !output.success {
return Err(EXEC_ERROR.to_owned());
}
parse_jsonl(&output.stdout)
}
async fn check_help(&self, args: &[&str], cutoff: Instant) -> Result<Vec<u8>, String> {
if Instant::now() >= cutoff {
return Err(DEADLINE_ERROR.to_owned());
}
let until = cutoff.min(Instant::now() + HELP_LIMIT);
let output = run_process(
self.executable.clone(),
args.iter().map(OsString::from).collect(),
Vec::new(),
until,
)
.await
.map_err(|error| match error {
ProcError::Timeout if Instant::now() >= cutoff => DEADLINE_ERROR.to_owned(),
ProcError::Timeout | ProcError::Io | ProcError::Cancelled => COMPAT_ERROR.to_owned(),
})?;
if !output.success {
return Err(COMPAT_ERROR.to_owned());
}
let mut text = output.stdout;
text.extend(output.stderr);
Ok(text)
}
}
fn exec_args(model: &str, effort: &str) -> Vec<OsString> {
let mut args = [
"exec",
"--search",
"--ephemeral",
"--ignore-user-config",
"--ignore-rules",
"--json",
"--sandbox",
"read-only",
"--ask-for-approval",
"never",
"--skip-git-repo-check",
"--no-daemon",
"--strict-config",
"--model",
model,
]
.into_iter()
.map(OsString::from)
.collect::<Vec<_>>();
args.push("-c".into());
args.push(format!("model_reasoning_effort={}", toml_quote(effort)).into());
for config in CONFIG {
args.push("-c".into());
args.push((*config).into());
}
args.push("-".into());
args
}
fn toml_quote(value: &str) -> String {
let mut output = String::from("\"");
for character in value.chars() {
match character {
'"' => output.push_str("\\\""),
'\\' => output.push_str("\\\\"),
'\u{8}' => output.push_str("\\b"),
'\t' => output.push_str("\\t"),
'\n' => output.push_str("\\n"),
'\u{c}' => output.push_str("\\f"),
'\r' => output.push_str("\\r"),
value if value.is_control() => output.push_str(&format!("\\u{:04X}", value as u32)),
value => output.push(value),
}
}
output.push('"');
output
}
fn parse_jsonl(bytes: &[u8]) -> Result<String, String> {
let mut turn_completed = false;
let mut turn_failed = false;
let mut message = None;
for raw in bytes.split(|byte| *byte == b'\n') {
let line = trim_ascii(raw);
if line.is_empty() {
continue;
}
let value: Value = serde_json::from_slice(line).map_err(|_| JSON_ERROR.to_owned())?;
match value.get("type").and_then(Value::as_str) {
Some("turn.completed") => turn_completed = true,
Some("turn.failed") => turn_failed = true,
Some("item.completed")
if value.pointer("/item/type").and_then(Value::as_str) == Some("agent_message") =>
{
message = Some(
value
.pointer("/item/text")
.and_then(Value::as_str)
.ok_or_else(|| JSON_ERROR.to_owned())?
.to_owned(),
);
}
_ => {}
}
}
if turn_failed {
return Err(EXEC_ERROR.to_owned());
}
if !turn_completed {
return Err(INCOMPLETE_ERROR.to_owned());
}
message.ok_or_else(|| MESSAGE_ERROR.to_owned())
}
fn trim_ascii(mut bytes: &[u8]) -> &[u8] {
while bytes.first().is_some_and(u8::is_ascii_whitespace) {
bytes = &bytes[1..];
}
while bytes.last().is_some_and(u8::is_ascii_whitespace) {
bytes = &bytes[..bytes.len() - 1];
}
bytes
}
struct ProcessOutput {
success: bool,
stdout: Vec<u8>,
stderr: Vec<u8>,
}
enum ProcError {
Io,
Timeout,
Cancelled,
}
struct CancelGuard(Option<oneshot::Sender<()>>);
impl Drop for CancelGuard {
fn drop(&mut self) {
if let Some(sender) = self.0.take() {
let _ = sender.send(());
}
}
}
async fn drain<R: AsyncRead + Unpin>(mut reader: R) -> io::Result<Vec<u8>> {
let mut bytes = Vec::new();
reader.read_to_end(&mut bytes).await?;
Ok(bytes)
}
async fn run_process(
executable: PathBuf,
args: Vec<OsString>,
input: Vec<u8>,
until: Instant,
) -> Result<ProcessOutput, ProcError> {
if Instant::now() >= until {
return Err(ProcError::Timeout);
}
let directory = tempfile::tempdir().map_err(|_| ProcError::Io)?;
let mut command = Command::new(executable);
command
.args(args)
.current_dir(directory.path())
.stdin(std::process::Stdio::piped())
.stdout(std::process::Stdio::piped())
.stderr(std::process::Stdio::piped())
.kill_on_drop(true)
.process_group(0);
let mut child = command.spawn().map_err(|_| ProcError::Io)?;
let pid = Pid::from_raw(child.id().ok_or(ProcError::Io)? as i32);
let mut stdin = child.stdin.take().ok_or(ProcError::Io)?;
let stdout = child.stdout.take().ok_or(ProcError::Io)?;
let stderr = child.stderr.take().ok_or(ProcError::Io)?;
let (cancel_tx, mut cancel_rx) = oneshot::channel();
let worker = tokio::spawn(async move {
let _directory = directory;
let complete = async move {
let write = async move {
stdin.write_all(&input).await?;
stdin.shutdown().await
};
let (status, written, stdout, stderr) =
tokio::join!(child.wait(), write, drain(stdout), drain(stderr));
written.map_err(|_| ProcError::Io)?;
Ok::<_, ProcError>(ProcessOutput {
success: status.map_err(|_| ProcError::Io)?.success(),
stdout: stdout.map_err(|_| ProcError::Io)?,
stderr: stderr.map_err(|_| ProcError::Io)?,
})
};
tokio::pin!(complete);
let timer = tokio::time::sleep_until(tokio::time::Instant::from_std(until));
tokio::pin!(timer);
tokio::select! {
result = &mut complete => result,
_ = &mut timer => {
let _ = killpg(pid, Signal::SIGKILL);
let _ = (&mut complete).await;
Err(ProcError::Timeout)
}
_ = &mut cancel_rx => {
let _ = killpg(pid, Signal::SIGKILL);
let _ = (&mut complete).await;
Err(ProcError::Cancelled)
}
}
});
let mut guard = CancelGuard(Some(cancel_tx));
let result = worker.await.map_err(|_| ProcError::Io)?;
guard.0.take();
result
}
#[cfg(test)]
mod tests {
use super::*;
use std::{fs, os::unix::fs::PermissionsExt, path::Path};
struct Fake {
_directory: tempfile::TempDir,
executable: PathBuf,
args: PathBuf,
input: PathBuf,
pids: PathBuf,
}
fn quote(path: &Path) -> String {
format!("'{}'", path.display().to_string().replace('\'', "'\"'\"'"))
}
fn fake(body: &str, exit: i32) -> Fake {
let directory = tempfile::tempdir().unwrap();
let executable = directory.path().join("codex");
let args = directory.path().join("args");
let input = directory.path().join("input");
let pids = directory.path().join("pids");
let body = body.replace("@PIDS@", "e(&pids));
let script = format!(
"#!/bin/sh\nif [ \"$1\" = \"--help\" ] || {{ [ \"$1\" = exec ] && [ \"$2\" = \"--help\" ]; }}; then\nprintf '%s\\n' '{}'\nexit 0\nfi\nprintf '%s\\n' \"$@\" > {}\ncat > {}\n{}\nexit {}\n",
REQUIRED_FLAGS.join(" "),
quote(&args),
quote(&input),
body,
exit,
);
fs::write(&executable, script).unwrap();
fs::set_permissions(&executable, fs::Permissions::from_mode(0o755)).unwrap();
Fake {
_directory: directory,
executable,
args,
input,
pids,
}
}
fn request(query: &str, duration: Duration) -> Request {
Request {
query: query.into(),
model: "model-x".into(),
reasoning_effort: "high".into(),
deadline: Instant::now() + duration,
}
}
#[tokio::test]
async fn invocation_stdin_and_last_message_are_exact() {
let fake = fake(
r#"printf '%s\n' \
'{"type":"future.event","x":1}' \
'{"type":"item.completed","item":{"type":"agent_message","text":"old"}}' \
'{"type":"item.completed","item":{"type":"agent_message","text":"最後\nline"}}' \
'{"type":"turn.completed","extra":true}'"#,
0,
);
let query = "héllo\n世界\0tail";
let answer = Runner::new(fake.executable.clone())
.run(request(query, Duration::from_secs(5)))
.await
.unwrap();
assert_eq!(answer, "最後\nline");
assert_eq!(fs::read(&fake.input).unwrap(), query.as_bytes());
let text = fs::read_to_string(&fake.args).unwrap();
let got = text.lines().collect::<Vec<_>>();
assert_eq!(
&got[..15],
[
"exec",
"--search",
"--ephemeral",
"--ignore-user-config",
"--ignore-rules",
"--json",
"--sandbox",
"read-only",
"--ask-for-approval",
"never",
"--skip-git-repo-check",
"--no-daemon",
"--strict-config",
"--model",
"model-x"
]
);
assert_eq!(&got[15..17], ["-c", r#"model_reasoning_effort="high""#]);
for (index, config) in CONFIG.iter().enumerate() {
assert_eq!(&got[17 + index * 2..19 + index * 2], ["-c", *config]);
}
assert_eq!(got.last(), Some(&"-"));
assert_eq!(toml_quote("a\"\n\u{7f}"), "\"a\\\"\\n\\u007F\"");
}
#[tokio::test]
async fn rejects_failures_and_nul_arguments() {
let cases = [
("printf '%s\\n' '{'", 0),
(r#"printf '%s\n' '{"type":"turn.failed"}'"#, 0),
(r#"printf '%s\n' '{"type":"turn.completed"}'"#, 7),
(r#"printf '%s\n' '{"type":"turn.completed"}'"#, 0),
(
r#"printf '%s\n' '{"type":"item.completed","item":{"type":"agent_message","text":"x"}}'"#,
0,
),
];
for (body, status) in cases {
let fake = fake(body, status);
assert!(
Runner::new(fake.executable)
.run(request("q", Duration::from_secs(5)))
.await
.is_err()
);
}
for model in [true, false] {
let runner = Runner::new("/not/executed".into());
let mut request = request("q", Duration::from_secs(2));
if model {
request.model.push('\0');
} else {
request.reasoning_effort.push('\0');
}
assert_eq!(runner.run(request).await.unwrap_err(), ARG_ERROR);
}
}
async fn wait_for_file(path: &Path) {
for _ in 0..200 {
if path.exists() {
return;
}
tokio::time::sleep(Duration::from_millis(10)).await;
}
panic!("fake process did not start");
}
fn process_is_live(pid: i32) -> bool {
if kill(Pid::from_raw(pid), None).is_err() {
return false;
}
let Ok(stat) = fs::read_to_string(format!("/proc/{pid}/stat")) else {
return true;
};
let Some(tail) = stat.rsplit_once(") ").map(|(_, tail)| tail) else {
return true;
};
!matches!(tail.as_bytes().first(), Some(b'Z' | b'X'))
}
async fn wait_for_processes_to_stop(path: &Path) {
let pids = fs::read_to_string(path)
.unwrap()
.split_whitespace()
.map(|value| value.parse::<i32>().unwrap())
.collect::<Vec<_>>();
assert!(pids.len() >= 2);
for _ in 0..200 {
if pids.iter().all(|pid| !process_is_live(*pid)) {
return;
}
tokio::time::sleep(Duration::from_millis(10)).await;
}
panic!("process-group member survived");
}
#[tokio::test]
async fn deadline_and_cancellation_kill_descendants() {
let deadline = fake(r#"sleep 30 & echo "$$ $!" > @PIDS@; wait"#, 0);
let error = Runner::new(deadline.executable.clone())
.run(request("q", Duration::from_millis(1600)))
.await
.unwrap_err();
assert_eq!(error, DEADLINE_ERROR);
wait_for_processes_to_stop(&deadline.pids).await;
let cancelled = fake(r#"sleep 30 & echo "$$ $!" > @PIDS@; wait"#, 0);
let runner = Runner::new(cancelled.executable.clone());
let task =
tokio::spawn(async move { runner.run(request("q", Duration::from_secs(10))).await });
wait_for_file(&cancelled.pids).await;
task.abort();
let _ = task.await;
wait_for_processes_to_stop(&cancelled.pids).await;
}
}