kcode_k1_codex_websearch/
lib.rs

1#![forbid(unsafe_code)]
2#![doc = include_str!("../Documentation.md")]
3#[cfg(not(unix))]
4compile_error!("kcode-k1-codex-websearch requires Unix process groups");
5
6use nix::{
7    sys::signal::{Signal, killpg},
8    unistd::Pid,
9};
10use serde_json::Value;
11use std::{ffi::OsString, io, path::PathBuf, time::Instant};
12use tokio::{
13    io::{AsyncRead, AsyncReadExt, AsyncWriteExt},
14    process::Command,
15    sync::oneshot,
16};
17
18const ARG_ERROR: &str =
19    "WebSearch failed: model or reasoning_effort is not representable as a process argument";
20const PREFLIGHT: &str =
21    "WebSearch failed: Codex compatibility check failed in kcode-k1-codex-websearch preflight";
22const EXEC_IO_ERROR: &str = "WebSearch failed: kcode-k1-codex-websearch could not start or communicate with the Codex execution subprocess";
23const EXEC_CANCELLED: &str =
24    "WebSearch failed: kcode-k1-codex-websearch Codex execution subprocess was cancelled";
25const EXEC_EXIT_ERROR: &str =
26    "WebSearch failed: kcode-k1-codex-websearch Codex execution subprocess exited unsuccessfully";
27const EXEC_ERROR: &str = "WebSearch failed: Codex execution failed";
28const BAD_JSON: &str = "WebSearch failed: invalid Codex JSONL output";
29const INCOMPLETE_ERROR: &str = "WebSearch failed: Codex response was incomplete";
30const MESSAGE_ERROR: &str = "WebSearch failed: Codex returned no agent message";
31
32const CONFIG: &str = r#"web_search="live"
33tools.web_search=true
34tools.view_image=false
35apps._default.enabled=false
36agents.enabled=false
37features.apps=false
38features.code_mode.enabled=false
39features.goals=false
40features.hooks=false
41features.memories=false
42features.multi_agent=false
43features.remote_plugin=false
44features.shell_snapshot=false
45features.shell_tool=false
46features.skill_mcp_dependency_install=false
47features.unified_exec=false
48memories.generate_memories=false
49memories.use_memories=false
50history.persistence="none"
51check_for_update_on_startup=false
52feedback.enabled=false
53analytics.enabled=false
54allow_login_shell=false
55skills.config=[]
56mcp_servers={}
57plugins={}
58marketplaces={}
59hooks={}"#;
60
61const EXEC_PREFIX: &str = r#"exec
62--search
63--ephemeral
64--ignore-user-config
65--ignore-rules
66--json
67--sandbox
68read-only
69--ask-for-approval
70never
71--skip-git-repo-check
72--no-daemon
73--strict-config
74--model"#;
75
76#[derive(Clone)]
77pub struct Runner {
78    executable: PathBuf,
79}
80
81pub struct Request {
82    pub query: String,
83    pub model: String,
84    pub reasoning_effort: String,
85    pub deadline: Instant,
86}
87
88impl Runner {
89    pub fn new(executable: PathBuf) -> Self {
90        Self { executable }
91    }
92
93    pub async fn run(&self, request: Request) -> Result<String, String> {
94        if request.model.contains('\0') || request.reasoning_effort.contains('\0') {
95            return Err(ARG_ERROR.to_owned());
96        }
97        let deadline = request.deadline;
98        if Instant::now() >= deadline {
99            return Err(deadline_error("before compatibility preflight"));
100        }
101
102        let mut help = self
103            .check_help(&["--help"], "codex --help", deadline)
104            .await?;
105        help.extend(
106            self.check_help(&["exec", "--help"], "codex exec --help", deadline)
107                .await?,
108        );
109        let help = String::from_utf8_lossy(&help);
110        let missing = required_flags()
111            .into_iter()
112            .filter(|flag| !help.contains(flag.as_str()))
113            .collect::<Vec<_>>();
114        if !missing.is_empty() {
115            return Err(format!(
116                "{PREFLIGHT}: combined required-flag validation omitted {}",
117                missing.join(", ")
118            ));
119        }
120        if Instant::now() >= deadline {
121            return Err(deadline_error("after compatibility preflight"));
122        }
123
124        let output = run_process(
125            self.executable.clone(),
126            exec_args(&request.model, &request.reasoning_effort),
127            request.query.into_bytes(),
128            deadline,
129        )
130        .await
131        .map_err(|error| match error {
132            ProcError::Timeout => deadline_error("during Codex execution"),
133            ProcError::Io => EXEC_IO_ERROR.to_owned(),
134            ProcError::Cancelled => EXEC_CANCELLED.to_owned(),
135        })?;
136        if !output.status.success() {
137            return Err(EXEC_EXIT_ERROR.to_owned());
138        }
139        parse_jsonl(&output.stdout)
140    }
141
142    async fn check_help(
143        &self,
144        args: &[&str],
145        phase: &str,
146        deadline: Instant,
147    ) -> Result<Vec<u8>, String> {
148        let output = run_process(
149            self.executable.clone(),
150            args.iter().map(OsString::from).collect(),
151            Vec::new(),
152            deadline,
153        )
154        .await
155        .map_err(|error| match error {
156            ProcError::Timeout => {
157                deadline_error(&format!("during compatibility preflight ({phase})"))
158            }
159            ProcError::Io => preflight_error(phase, "subprocess setup, spawn, or I/O failed"),
160            ProcError::Cancelled => preflight_error(phase, "subprocess was cancelled"),
161        })?;
162        if !output.status.success() {
163            return Err(preflight_error(phase, "subprocess exited unsuccessfully"));
164        }
165        let mut text = output.stdout;
166        text.extend(output.stderr);
167        Ok(text)
168    }
169}
170
171fn deadline_error(phase: &str) -> String {
172    format!("WebSearch failed: absolute operation deadline exceeded {phase}")
173}
174
175fn preflight_error(phase: &str, cause: &str) -> String {
176    format!("{PREFLIGHT}: {phase} {cause}")
177}
178
179fn exec_args(model: &str, effort: &str) -> Vec<OsString> {
180    let mut args = EXEC_PREFIX.lines().map(OsString::from).collect::<Vec<_>>();
181    args.push(model.into());
182    args.push("-c".into());
183    args.push(format!("model_reasoning_effort={}", toml_quote(effort)).into());
184    for config in CONFIG.lines() {
185        args.push("-c".into());
186        args.push(config.into());
187    }
188    args.push("-".into());
189    args
190}
191
192fn required_flags() -> Vec<String> {
193    let mut flags = exec_args("", "")
194        .into_iter()
195        .filter_map(|arg| arg.into_string().ok())
196        .filter(|arg| arg.starts_with("--"))
197        .collect::<Vec<_>>();
198    flags.push("--config".into());
199    flags
200}
201
202fn toml_quote(value: &str) -> String {
203    let mut output = String::from("\"");
204    for character in value.chars() {
205        match character {
206            '"' => output.push_str("\\\""),
207            '\\' => output.push_str("\\\\"),
208            '\u{8}' => output.push_str("\\b"),
209            '\t' => output.push_str("\\t"),
210            '\n' => output.push_str("\\n"),
211            '\u{c}' => output.push_str("\\f"),
212            '\r' => output.push_str("\\r"),
213            value if value.is_control() => output.push_str(&format!("\\u{:04X}", value as u32)),
214            value => output.push(value),
215        }
216    }
217    output.push('"');
218    output
219}
220
221fn parse_jsonl(bytes: &[u8]) -> Result<String, String> {
222    let mut turn_completed = false;
223    let mut turn_failed = false;
224    let mut message = None;
225    for raw in bytes.split(|byte| *byte == b'\n') {
226        let line = raw.trim_ascii();
227        if line.is_empty() {
228            continue;
229        }
230        let value: Value = serde_json::from_slice(line).map_err(|_| BAD_JSON.to_owned())?;
231        match value.get("type").and_then(Value::as_str) {
232            Some("turn.completed") => turn_completed = true,
233            Some("turn.failed") => turn_failed = true,
234            Some("item.completed")
235                if value.pointer("/item/type").and_then(Value::as_str) == Some("agent_message") =>
236            {
237                let text = value
238                    .pointer("/item/text")
239                    .and_then(Value::as_str)
240                    .ok_or_else(|| BAD_JSON.to_owned())?;
241                message = Some(text.to_owned());
242            }
243            _ => {}
244        }
245    }
246    if turn_failed {
247        return Err(EXEC_ERROR.to_owned());
248    }
249    if !turn_completed {
250        return Err(INCOMPLETE_ERROR.to_owned());
251    }
252    message.ok_or_else(|| MESSAGE_ERROR.to_owned())
253}
254
255enum ProcError {
256    Io,
257    Timeout,
258    Cancelled,
259}
260
261async fn drain<R: AsyncRead + Unpin>(mut reader: R) -> io::Result<Vec<u8>> {
262    let mut bytes = Vec::new();
263    reader.read_to_end(&mut bytes).await?;
264    Ok(bytes)
265}
266
267async fn run_process(
268    executable: PathBuf,
269    args: Vec<OsString>,
270    input: Vec<u8>,
271    until: Instant,
272) -> Result<std::process::Output, ProcError> {
273    if Instant::now() >= until {
274        return Err(ProcError::Timeout);
275    }
276    let directory = tempfile::tempdir().map_err(|_| ProcError::Io)?;
277    let mut command = Command::new(executable);
278    command
279        .args(args)
280        .current_dir(directory.path())
281        .stdin(std::process::Stdio::piped())
282        .stdout(std::process::Stdio::piped())
283        .stderr(std::process::Stdio::piped())
284        .kill_on_drop(true)
285        .process_group(0);
286    let mut child = command.spawn().map_err(|_| ProcError::Io)?;
287    let pid = Pid::from_raw(child.id().ok_or(ProcError::Io)? as i32);
288    let mut stdin = child.stdin.take().ok_or(ProcError::Io)?;
289    let stdout = child.stdout.take().ok_or(ProcError::Io)?;
290    let stderr = child.stderr.take().ok_or(ProcError::Io)?;
291    let (cancel_tx, mut cancel_rx) = oneshot::channel::<()>();
292
293    let worker = tokio::spawn(async move {
294        let _directory = directory;
295        let complete = async move {
296            let write = async move {
297                if input.is_empty() {
298                    return Ok(());
299                }
300                stdin.write_all(&input).await?;
301                stdin.shutdown().await
302            };
303            let (status, written, stdout, stderr) =
304                tokio::join!(child.wait(), write, drain(stdout), drain(stderr));
305            written.map_err(|_| ProcError::Io)?;
306            Ok::<_, ProcError>(std::process::Output {
307                status: status.map_err(|_| ProcError::Io)?,
308                stdout: stdout.map_err(|_| ProcError::Io)?,
309                stderr: stderr.map_err(|_| ProcError::Io)?,
310            })
311        };
312        tokio::pin!(complete);
313        let timer = tokio::time::sleep_until(tokio::time::Instant::from_std(until));
314        tokio::pin!(timer);
315        let error = tokio::select! {
316            result = &mut complete => return result,
317            _ = &mut timer => ProcError::Timeout,
318            _ = &mut cancel_rx => ProcError::Cancelled,
319        };
320        let _ = killpg(pid, Signal::SIGKILL);
321        let _ = (&mut complete).await;
322        Err(error)
323    });
324    let result = worker.await.map_err(|_| ProcError::Io)?;
325    drop(cancel_tx);
326    result
327}
328
329#[cfg(test)]
330mod tests {
331    use super::*;
332    use std::{fs, os::unix::fs::PermissionsExt, path::Path, time::Duration};
333
334    struct Fake(tempfile::TempDir);
335
336    impl Fake {
337        fn file(&self, name: &str) -> PathBuf {
338            self.0.path().join(name)
339        }
340
341        async fn run(&self, duration: Duration) -> Result<String, String> {
342            Runner::new(self.file("codex"))
343                .run(request("q", duration))
344                .await
345        }
346    }
347
348    fn fake(body: &str, exit: i32, help: Option<String>) -> Fake {
349        let directory = tempfile::tempdir().unwrap();
350        let executable = directory.path().join("codex");
351        let flags = required_flags().join(" ");
352        let help = help.unwrap_or_else(|| format!("printf '%s\\n' '{flags}'"));
353        let script = format!(
354            "#!/bin/sh\nbase=${{0%/*}}\nif [ \"$1\" = \"--help\" ] || {{ [ \"$1\" = exec ] && [ \"$2\" = \"--help\" ]; }}; then\n{}\nexit 0\nfi\nprintf '%s\\n' \"$@\" > \"$base/args\"\ncat > \"$base/input\"\n{}\nexit {}\n",
355            help, body, exit,
356        );
357        fs::write(&executable, script).unwrap();
358        fs::set_permissions(&executable, fs::Permissions::from_mode(0o755)).unwrap();
359        Fake(directory)
360    }
361
362    fn request(query: &str, duration: Duration) -> Request {
363        Request {
364            query: query.into(),
365            model: "model-x".into(),
366            reasoning_effort: "high".into(),
367            deadline: Instant::now() + duration,
368        }
369    }
370
371    const SUCCESS: &str = r#"printf '%s\n' \
372'{"type":"item.completed","item":{"type":"agent_message","text":"answer"}}' \
373'{"type":"turn.completed"}'"#;
374    const DESCENDANTS: &str = r#"sleep 30 & echo "$$ $!" > "$base/pids"; wait"#;
375
376    #[tokio::test]
377    async fn invocation_stdin_and_last_message_are_exact() {
378        let fake = fake(
379            r#"printf '%s\n' \
380'{"type":"future.event","x":1}' \
381'{"type":"item.completed","item":{"type":"agent_message","text":"old"}}' \
382'{"type":"item.completed","item":{"type":"agent_message","text":"最後\nline"}}' \
383'{"type":"turn.completed","extra":true}'"#,
384            0,
385            None,
386        );
387        let query = "héllo\n世界\0tail";
388        let answer = Runner::new(fake.file("codex"))
389            .run(request(query, Duration::from_secs(5)))
390            .await
391            .unwrap();
392        assert_eq!(answer, "最後\nline");
393        assert_eq!(fs::read(fake.file("input")).unwrap(), query.as_bytes());
394        let text = fs::read_to_string(fake.file("args")).unwrap();
395        let got = text.lines().collect::<Vec<_>>();
396        assert_eq!(&got[..3], ["exec", "--search", "--ephemeral"]);
397        assert_eq!(
398            &got[3..6],
399            ["--ignore-user-config", "--ignore-rules", "--json"]
400        );
401        assert_eq!(
402            &got[6..10],
403            ["--sandbox", "read-only", "--ask-for-approval", "never"]
404        );
405        assert_eq!(
406            &got[10..13],
407            ["--skip-git-repo-check", "--no-daemon", "--strict-config"]
408        );
409        assert_eq!(&got[13..15], ["--model", "model-x"]);
410        assert_eq!(&got[15..17], ["-c", r#"model_reasoning_effort="high""#]);
411        for (index, config) in CONFIG.lines().enumerate() {
412            assert_eq!(&got[17 + index * 2..19 + index * 2], ["-c", config]);
413        }
414        assert_eq!(got.last(), Some(&"-"));
415        assert_eq!(toml_quote("a\"\n\u{7f}"), "\"a\\\"\\n\\u007F\"");
416    }
417
418    #[tokio::test]
419    async fn compatibility_diagnostic_names_missing_flag_without_help() {
420        let advertised = required_flags().join(" ").replace("--no-daemon", "");
421        let help = format!("printf '%s\\n' '{advertised} conspicuous-private-help-text'");
422        let fake = fake("", 0, Some(help));
423        let error = fake.run(Duration::from_secs(5)).await.unwrap_err();
424        assert_eq!(
425            error,
426            preflight_error("combined required-flag validation", "omitted --no-daemon")
427        );
428        assert!(!error.contains("conspicuous-private-help-text"));
429    }
430
431    #[tokio::test]
432    async fn help_and_execution_use_only_the_absolute_deadline() {
433        let help = format!(
434            "if [ \"$1\" = \"--help\" ]; then sleep 3; fi\nprintf '%s\\n' '{}'",
435            required_flags().join(" ")
436        );
437        let slow_help = fake(SUCCESS, 0, Some(help));
438        assert_eq!(
439            slow_help.run(Duration::from_secs(7)).await.unwrap(),
440            "answer"
441        );
442        let body = format!("sleep 2\n{SUCCESS}");
443        let slow_execution = fake(&body, 0, None);
444        assert_eq!(
445            slow_execution
446                .run(Duration::from_millis(2800))
447                .await
448                .unwrap(),
449            "answer"
450        );
451    }
452
453    #[tokio::test]
454    async fn rejects_failures_and_nul_arguments() {
455        let cases = [
456            ("{", 0, BAD_JSON),
457            (r#"{"type":"turn.failed"}"#, 0, EXEC_ERROR),
458            (r#"{"type":"turn.completed"}"#, 7, EXEC_EXIT_ERROR),
459            (r#"{"type":"turn.completed"}"#, 0, MESSAGE_ERROR),
460            (
461                r#"{"type":"item.completed","item":{"type":"agent_message","text":"x"}}"#,
462                0,
463                INCOMPLETE_ERROR,
464            ),
465        ];
466        for (json, status, error) in cases {
467            let body = format!("printf '%s\\n' '{json}'");
468            let fake = fake(&body, status, None);
469            assert_eq!(fake.run(Duration::from_secs(5)).await.unwrap_err(), error);
470        }
471
472        let runner = Runner::new("/not/executed".into());
473        for (model, effort) in [("bad\0", "high"), ("model-x", "bad\0")] {
474            let mut bad = request("q", Duration::from_secs(2));
475            bad.model = model.into();
476            bad.reasoning_effort = effort.into();
477            assert_eq!(runner.run(bad).await.unwrap_err(), ARG_ERROR);
478        }
479    }
480
481    async fn wait_until(mut condition: impl FnMut() -> bool) {
482        for _ in 0..200 {
483            if condition() {
484                return;
485            }
486            tokio::time::sleep(Duration::from_millis(10)).await;
487        }
488        panic!("condition was not reached");
489    }
490
491    fn process_is_live(pid: i32) -> bool {
492        fs::read_to_string(format!("/proc/{pid}/stat")).is_ok_and(|stat| {
493            let state = stat
494                .rsplit_once(") ")
495                .and_then(|(_, tail)| tail.bytes().next());
496            !matches!(state, Some(b'Z' | b'X'))
497        })
498    }
499
500    async fn wait_for_processes_to_stop(path: &Path) {
501        let pids = fs::read_to_string(path)
502            .unwrap()
503            .split_whitespace()
504            .map(|value| value.parse::<i32>().unwrap())
505            .collect::<Vec<_>>();
506        assert!(pids.len() >= 2);
507        wait_until(|| pids.iter().all(|pid| !process_is_live(*pid))).await;
508    }
509
510    #[tokio::test]
511    async fn deadline_and_cancellation_reap_descendants() {
512        let deadline = fake(DESCENDANTS, 0, None);
513        let error = deadline.run(Duration::from_millis(1600)).await.unwrap_err();
514        assert_eq!(error, deadline_error("during Codex execution"));
515        wait_for_processes_to_stop(&deadline.file("pids")).await;
516
517        let cancelled = fake(DESCENDANTS, 0, None);
518        let runner = Runner::new(cancelled.file("codex"));
519        let task =
520            tokio::spawn(async move { runner.run(request("q", Duration::from_secs(10))).await });
521        wait_until(|| cancelled.file("pids").exists()).await;
522        task.abort();
523        let _ = task.await;
524        wait_for_processes_to_stop(&cancelled.file("pids")).await;
525    }
526}
kcode_k1_codex_websearch/lib.rs

kcode_k1_codex_websearch/
lib.rs
