Skip to main content

kcode_k1_codex_testkit/
lib.rs

1//! Deterministic fake app-server conformance for the K1 Codex adapter family.
2
3use kcode_k1_codex_conversations::State;
4use kcode_k1_codex_events::{Inbound, RpcId, decode, rejection_payload};
5use kcode_k1_codex_runtime::{
6    Adapter, Config, DynamicTool, ErrorKind, Event, ToolCall, ToolResult, Turn,
7};
8use kcode_k1_codex_shim::{
9    ASYNC_TOOL_ACKNOWLEDGEMENT, BoxCodec, Shim, ShimItem, ToolCallLauncher, ToolLaunchFuture,
10};
11use serde_json::{Value, json};
12use std::{
13    fmt::Display,
14    fs::{self, OpenOptions},
15    io::Write as _,
16    os::unix::fs::PermissionsExt,
17    path::{Path, PathBuf},
18    sync::{
19        Arc, Mutex,
20        atomic::{AtomicU64, AtomicUsize, Ordering},
21    },
22    time::Duration,
23};
24
25const SCRIPT: &str = r#"#!/bin/sh
26set -eu
27: > transcript.log; : > launch.log
28fail() { printf '%s\n' "$1" > fake.fail; exit 1; }
29printf '%s\n' "$@" > args.log
30[ "$1" = app-server ] || fail subcommand
31history=0; effort=0
32for arg in "$@"; do case "$arg" in 'history.persistence="none"') history=1 ;; 'model_reasoning_effort="medium"') effort=1 ;; model=*) fail global-model ;; esac; done
33[ "$history/$effort" = 1/1 ] || fail global-config
34printf 'diagnostic-begin\n' >&2; i=0
35while [ "$i" -lt 2048 ]; do printf '0123456789abcdef' >&2; i=$((i + 1)); done
36printf '\ndiagnostic-end\n' >&2
37init=0; thread=0; turn=0; unsub=0; intr=0; nr=0; sr=0; ur=0; finished=0
38out() { printf '%s\n' "$1"; }
39rpc_id() { printf '%s\n' "$1" | sed -n 's/.*"id":\([0-9][0-9]*\).*/\1/p'; }
40field() { printf '%s\n' "$1" | sed -n "s/.*\"$2\":\"\([^\"]*\)\".*/\1/p"; }
41complete_main() {
42  if [ "$nr$sr$ur$finished" = 1110 ]; then
43    out '{"method":"turn/completed","params":{"threadId":"th1","turn":{"id":"tu1","status":"completed"}}}'
44    out '{"method":"turn/completed","params":{"threadId":"th2","turn":{"id":"tu2","status":"completed"}}}'
45    finished=1
46  fi
47}
48while IFS= read -r line; do
49  printf '%s\n' "$line" >> transcript.log
50  case "$line" in
51    *'"method":"initialize"'*)
52      init=$((init + 1)); [ "$init" -eq 1 ] || fail initialization-count; out '{"id":0,"result":{}}' ;;
53    *'"method":"initialized"'*) ;;
54    *'"method":"thread/start"'*)
55      case "$line" in *'"model":"requested-model"'*) ;; *) fail thread-model ;; esac
56      case "$line" in *'"dynamicTools":'*'"name":"launch"'*) ;; *) fail tool-registry ;; esac
57      case "$line" in *ephemeral*) fail ephemeral-thread ;; esac
58      thread=$((thread + 1)); id=$(rpc_id "$line"); out "{\"id\":$id,\"result\":{\"thread\":{\"id\":\"th$thread\"}}}" ;;
59    *'"method":"turn/start"'*)
60      turn=$((turn + 1)); id=$(rpc_id "$line"); th=$(field "$line" threadId); tu="tu$turn"
61      case "$turn:$th" in 1:th1|2:th2|3:th1|4:th3|5:th4|6:th5) ;; *) fail thread-routing ;; esac
62      if [ "$turn" -eq 6 ]; then case "$line" in *'H1\nH2\nprompt'*) ;; *) fail shim-history ;; esac; fi
63      out "{\"id\":$id,\"result\":{\"turn\":{\"id\":\"$tu\"}}}"
64      case "$turn" in
65        2)
66          out '{"method":"future/request","id":"unsupported","params":{}}'; out '{"method":"item/agentMessage/delta","params":{"threadId":"th1","turnId":"tu1","delta":"A"}}'; out '{"method":"item/agentMessage/delta","params":{"threadId":"th2","turnId":"tu2","delta":"B"}}'; out '{"method":"item/tool/call","id":41,"params":{"threadId":"th1","turnId":"tu1","callId":"num","tool":"launch","arguments":{"kind":"number"}}}'; out '{"method":"item/tool/call","id":"41","params":{"threadId":"th2","turnId":"tu2","callId":"str","tool":"launch","arguments":{"kind":"string"}}}' ;;
67        3)
68          out "{\"method\":\"item/agentMessage/delta\",\"params\":{\"threadId\":\"$th\",\"turnId\":\"$tu\",\"delta\":\"again\"}}"; out "{\"method\":\"item/tool/call\",\"id\":\"repeat-gate\",\"params\":{\"threadId\":\"$th\",\"turnId\":\"$tu\",\"callId\":\"gate\",\"tool\":\"launch\",\"arguments\":{}}}" ;;
69        4) out "{\"method\":\"turn/completed\",\"params\":{\"threadId\":\"$th\",\"turn\":{\"id\":\"$tu\",\"status\":\"completed\"}}}" ;;
70        5) out "{\"method\":\"item/agentMessage/delta\",\"params\":{\"threadId\":\"$th\",\"turnId\":\"$tu\",\"delta\":\"drop-ready\"}}" ;;
71        6)
72          out "{\"method\":\"item/agentMessage/delta\",\"params\":{\"threadId\":\"$th\",\"turnId\":\"$tu\",\"delta\":\"L\"}}"; out "{\"method\":\"item/agentMessage/delta\",\"params\":{\"threadId\":\"$th\",\"turnId\":\"$tu\",\"delta\":\"1\"}}"; out "{\"method\":\"item/tool/call\",\"id\":\"wave-1\",\"params\":{\"threadId\":\"$th\",\"turnId\":\"$tu\",\"callId\":\"c1\",\"tool\":\"launch\",\"arguments\":{\"wave\":1}}}" ;;
73      esac ;;
74    *'"method":"thread/unsubscribe"'*)
75      unsub=$((unsub + 1)); id=$(rpc_id "$line"); out "{\"id\":$id,\"result\":{}}" ;;
76    *'"method":"turn/interrupt"'*)
77      intr=$((intr + 1)); id=$(rpc_id "$line"); th=$(field "$line" threadId); tu=$(field "$line" turnId)
78      printf '%s/%s\n' "$th" "$tu" > interrupt.log; out "{\"id\":$id,\"result\":{}}"; out "{\"method\":\"turn/completed\",\"params\":{\"threadId\":\"$th\",\"turn\":{\"id\":\"$tu\",\"status\":\"interrupted\"}}}" ;;
79    *)
80      case "$line" in
81        *'"id":"unsupported"'*)
82          printf '%s\n' "$line" | grep -q '"code":-32601' || fail unsupported; ur=1 ;;
83        *'"id":41'*)
84          printf '%s\n' "$line" | grep -q '"text":"numeric"' || fail numeric-result; nr=1 ;;
85        *'"id":"41"'*)
86          printf '%s\n' "$line" | grep -q '"text":"string"' || fail string-result; sr=1 ;;
87        *'"id":"repeat-gate"'*) out '{"method":"turn/completed","params":{"threadId":"th1","turn":{"id":"tu3","status":"completed"}}}' ;;
88        *'"id":"wave-1"'*)
89          grep -qx 'c1:1' launch.log || fail launch-before-first-ack; out '{"method":"item/agentMessage/delta","params":{"threadId":"th5","turnId":"tu6","delta":"M"}}'; out '{"method":"item/agentMessage/delta","params":{"threadId":"th5","turnId":"tu6","delta":"2"}}'; out '{"method":"item/tool/call","id":72,"params":{"threadId":"th5","turnId":"tu6","callId":"c2","tool":"launch","arguments":{"wave":2}}}' ;;
90        *'"id":72'*)
91          grep -qx 'c2:2' launch.log || fail launch-before-second-ack; out '{"method":"item/agentMessage/delta","params":{"threadId":"th5","turnId":"tu6","delta":"R"}}'; out '{"method":"item/agentMessage/delta","params":{"threadId":"th5","turnId":"tu6","delta":"3"}}'; out '{"method":"turn/completed","params":{"threadId":"th5","turn":{"id":"tu6","status":"completed"}}}' ;;
92      esac
93      complete_main ;;
94  esac
95done
96[ "$init/$thread/$turn/$unsub/$intr" = 1/5/6/2/1 ] || fail final-counts
97printf stopped > stopped.log
98"#;
99
100trait Text<T> {
101    fn text(self) -> Result<T, String>;
102}
103
104impl<T, E: Display> Text<T> for Result<T, E> {
105    fn text(self) -> Result<T, String> {
106        self.map_err(|error| error.to_string())
107    }
108}
109
110macro_rules! check {
111    ($condition:expr, $message:literal) => {
112        if !$condition {
113            return Err($message.to_owned());
114        }
115    };
116}
117
118struct App {
119    directory: PathBuf,
120    executable: PathBuf,
121}
122
123impl App {
124    fn new() -> Result<Self, String> {
125        static NEXT: AtomicU64 = AtomicU64::new(0);
126        let suffix = NEXT.fetch_add(1, Ordering::Relaxed);
127        let directory =
128            std::env::temp_dir().join(format!("k1-codex-testkit-{}-{suffix}", std::process::id()));
129        fs::create_dir(&directory).text()?;
130        let executable = directory.join("codex");
131        fs::write(&executable, SCRIPT).text()?;
132        let mut permissions = fs::metadata(&executable).text()?.permissions();
133        permissions.set_mode(0o700);
134        fs::set_permissions(&executable, permissions).text()?;
135        Ok(Self {
136            directory,
137            executable,
138        })
139    }
140
141    fn path(&self, name: &str) -> PathBuf {
142        self.directory.join(name)
143    }
144}
145
146impl Drop for App {
147    fn drop(&mut self) {
148        let _ = fs::remove_dir_all(&self.directory);
149    }
150}
151
152struct Codec(Arc<AtomicUsize>);
153
154impl BoxCodec for Codec {
155    type Box = String;
156
157    fn tool_call_box(&mut self, call: &ToolCall) -> String {
158        self.0.fetch_add(1, Ordering::Relaxed);
159        format!("{}:{}", call.call_id, call.arguments["wave"])
160    }
161
162    fn box_text<'a>(&self, box_: &'a String) -> &'a str {
163        box_
164    }
165}
166
167struct Launcher {
168    seen: Arc<Mutex<Vec<String>>>,
169    log: PathBuf,
170}
171
172impl ToolCallLauncher<String> for Launcher {
173    fn launch<'a>(&'a mut self, box_: &'a String) -> ToolLaunchFuture<'a> {
174        Box::pin(async move {
175            self.seen
176                .lock()
177                .map_err(|_| "launch record lock poisoned".to_owned())?
178                .push(box_.clone());
179            let mut log = OpenOptions::new().append(true).open(&self.log).text()?;
180            writeln!(log, "{box_}").text()
181        })
182    }
183}
184
185fn config(app: &App) -> Config {
186    Config {
187        executable: app.executable.clone(),
188        working_directory: app.directory.to_string_lossy().into_owned(),
189        model: "requested-model".to_owned(),
190        reasoning_effort: Some("medium".to_owned()),
191        base_instructions: "dynamic tools only".to_owned(),
192        tools: vec![DynamicTool {
193            name: "launch".to_owned(),
194            description: "launch a deterministic test call".to_owned(),
195            input_schema: json!({"type": "object"}),
196        }],
197    }
198}
199
200async fn next(turn: &mut Turn) -> Result<Event, String> {
201    turn.next_event()
202        .await
203        .ok_or_else(|| "turn event stream ended early".to_owned())
204}
205
206async fn text(turn: &mut Turn, expected: &str) -> Result<(), String> {
207    check!(
208        next(turn).await? == Event::TextDelta(expected.to_owned()),
209        "text event was wrong"
210    );
211    Ok(())
212}
213
214async fn done(turn: &mut Turn) -> Result<(), String> {
215    check!(next(turn).await? == Event::Done, "turn did not complete");
216    Ok(())
217}
218
219fn tool(event: Event, id: &str) -> Result<ToolCall, String> {
220    match event {
221        Event::ToolCall(call) if call.call_id == id => Ok(call),
222        other => Err(format!("expected tool call {id}, got {other:?}")),
223    }
224}
225
226async fn wait_for(mut ready: impl FnMut() -> bool, label: &str) -> Result<(), String> {
227    tokio::time::timeout(Duration::from_secs(5), async {
228        while !ready() {
229            tokio::task::yield_now().await;
230        }
231    })
232    .await
233    .map_err(|_| format!("timed out waiting for {label}"))
234}
235
236fn lower_layer_checks() -> Result<(), String> {
237    let Inbound::UnsupportedServerRequest(unsupported) =
238        decode(json!({"method": "future/request", "id": "u"}))
239    else {
240        return Err("unsupported request was misclassified".to_owned());
241    };
242    check!(
243        rejection_payload(&unsupported.id, -32601, "unsupported")["id"] == "u",
244        "unsupported rejection lost its ID"
245    );
246    let malformed = json!({
247        "method": "item/tool/call",
248        "id": 9,
249        "params": {"threadId": "t", "turnId": "u", "callId": "c", "tool": "launch"}
250    });
251    check!(
252        matches!(decode(malformed.clone()), Inbound::Malformed(_)),
253        "malformed call did not fail closed"
254    );
255    let malformed_id = RpcId::try_from(&malformed["id"]).text()?;
256    check!(
257        rejection_payload(&malformed_id, -32602, "malformed")["id"] == 9,
258        "malformed rejection lost its ID"
259    );
260    let Inbound::ModelReroute(reroute) = decode(json!({"method": "model/rerouted", "id": "r"}))
261    else {
262        return Err("model reroute was misclassified".to_owned());
263    };
264    let reroute_id = reroute.id.ok_or_else(|| "reroute lost its ID".to_owned())?;
265    check!(
266        rejection_payload(&reroute_id, -32602, "reroute")["id"] == "r",
267        "reroute rejection lost its ID"
268    );
269    let mut state = State::<(), ()>::default();
270    let serial = state.begin_turn("k", ()).text()?;
271    state.track_tool("k", serial, "call", &json!(7)).text()?;
272    let duplicate = state
273        .track_tool("k", serial, "call", &json!(8))
274        .expect_err("duplicate call must fail");
275    check!(
276        duplicate.kind == ErrorKind::Protocol
277            && duplicate.message == "duplicate dynamic tool call id",
278        "duplicate call did not fail closed"
279    );
280    let duplicate = state
281        .track_tool("k", serial, "other", &json!(7))
282        .expect_err("duplicate RPC ID must fail");
283    check!(
284        duplicate.kind == ErrorKind::Protocol
285            && duplicate.message == "duplicate app-server request id",
286        "duplicate RPC ID did not fail closed"
287    );
288    Ok(())
289}
290
291fn messages(path: &Path) -> Result<Vec<Value>, String> {
292    fs::read_to_string(path)
293        .text()?
294        .lines()
295        .map(|line| serde_json::from_str(line).text())
296        .collect()
297}
298
299fn response<'a>(messages: &'a [Value], id: &Value) -> Option<&'a Value> {
300    messages
301        .iter()
302        .find(|message| message.get("id") == Some(id) && message.get("result").is_some())
303}
304
305fn inspect(app: &App) -> Result<(), String> {
306    check!(!app.path("fake.fail").exists(), "fake app-server failed");
307    check!(
308        fs::read_to_string(app.path("stopped.log")).text()? == "stopped",
309        "app-server did not stop cleanly"
310    );
311    let messages = messages(&app.path("transcript.log"))?;
312    for (id, output) in [(json!(41), "numeric"), (json!("41"), "string")] {
313        let response =
314            response(&messages, &id).ok_or_else(|| "typed tool response missing".to_owned())?;
315        check!(
316            response
317                .pointer("/result/contentItems/0/text")
318                .and_then(Value::as_str)
319                == Some(output),
320            "typed tool response crossed IDs"
321        );
322    }
323    for id in [json!("wave-1"), json!(72)] {
324        let response =
325            response(&messages, &id).ok_or_else(|| "shim acknowledgement missing".to_owned())?;
326        check!(
327            response.pointer("/result/success").and_then(Value::as_bool) == Some(true)
328                && response
329                    .pointer("/result/contentItems/0/text")
330                    .and_then(Value::as_str)
331                    == Some(ASYNC_TOOL_ACKNOWLEDGEMENT),
332            "shim acknowledgement was not exact"
333        );
334    }
335    check!(
336        fs::read_to_string(app.path("launch.log")).text()? == "c1:1\nc2:2\n",
337        "launch order was wrong"
338    );
339    Ok(())
340}
341
342/// Runs the complete local fake-app-server conformance scenario.
343pub async fn run_conformance() -> Result<(), String> {
344    lower_layer_checks()?;
345    let app = App::new()?;
346    let adapter = Adapter::open(config(&app)).await.text()?;
347    wait_for(
348        || {
349            adapter
350                .diagnostics()
351                .windows(b"diagnostic-end".len())
352                .any(|window| window == b"diagnostic-end")
353        },
354        "stderr capture",
355    )
356    .await?;
357    let mut alpha = adapter.start_turn("alpha", "first").await.text()?;
358    let busy = adapter
359        .start_turn("alpha", "overlap")
360        .await
361        .expect_err("same-key overlap must fail");
362    check!(
363        busy.kind == ErrorKind::Busy,
364        "same-key overlap was not Busy"
365    );
366    check!(
367        busy.diagnostics.len() >= 32_768 && !busy.to_string().contains("diagnostic-begin"),
368        "diagnostics were truncated or leaked into Display"
369    );
370    let mut beta = adapter.start_turn("beta", "second").await.text()?;
371    text(&mut alpha, "A").await?;
372    text(&mut beta, "B").await?;
373    let numeric = tool(next(&mut alpha).await?, "num")?;
374    let string = tool(next(&mut beta).await?, "str")?;
375    alpha
376        .respond(
377            numeric.call_id,
378            ToolResult {
379                success: true,
380                output: "numeric".to_owned(),
381            },
382        )
383        .await
384        .text()?;
385    beta.respond(
386        string.call_id,
387        ToolResult {
388            success: true,
389            output: "string".to_owned(),
390        },
391    )
392    .await
393    .text()?;
394    done(&mut alpha).await?;
395    done(&mut beta).await?;
396
397    let mut repeat = adapter.start_turn("alpha", "again").await.text()?;
398    text(&mut repeat, "again").await?;
399    let gate = tool(next(&mut repeat).await?, "gate")?;
400    let busy = adapter
401        .close_conversation("alpha")
402        .await
403        .expect_err("active close must fail");
404    check!(busy.kind == ErrorKind::Busy, "active close was not Busy");
405    repeat
406        .respond(
407            gate.call_id,
408            ToolResult {
409                success: true,
410                output: "gate".to_owned(),
411            },
412        )
413        .await
414        .text()?;
415    done(&mut repeat).await?;
416    adapter.close_conversation("alpha").await.text()?;
417    adapter.close_conversation("alpha").await.text()?;
418    let mut fresh = adapter.start_turn("alpha", "fresh").await.text()?;
419    done(&mut fresh).await?;
420
421    let mut dropped = adapter.start_turn("drop", "drop").await.text()?;
422    text(&mut dropped, "drop-ready").await?;
423    drop(dropped);
424    wait_for(|| app.path("interrupt.log").exists(), "turn interrupt").await?;
425
426    let conversions = Arc::new(AtomicUsize::new(0));
427    let launches = Arc::new(Mutex::new(Vec::new()));
428    let launcher = Launcher {
429        seen: launches.clone(),
430        log: app.path("launch.log"),
431    };
432    let mut shim = Shim::new(
433        adapter.clone(),
434        "shim",
435        Codec(conversions.clone()),
436        Box::new(launcher),
437    );
438    shim.record_boxes(["H1".to_owned(), "H2".to_owned()]);
439    let output = shim.infer("prompt").await.text()?;
440    check!(
441        output.items
442            == vec![
443                ShimItem::Text("L1".to_owned()),
444                ShimItem::Box("c1:1".to_owned()),
445                ShimItem::Text("M2".to_owned()),
446                ShimItem::Box("c2:2".to_owned()),
447                ShimItem::Text("R3".to_owned()),
448            ],
449        "shim event ordering or text coalescing was wrong"
450    );
451    check!(
452        shim.pending_box_count() == 0 && conversions.load(Ordering::Relaxed) == 2,
453        "shim conversion or history count was wrong"
454    );
455    {
456        let launches = launches
457            .lock()
458            .map_err(|_| "launch record lock poisoned".to_owned())?;
459        check!(
460            launches.as_slice() == ["c1:1", "c2:2"],
461            "shim launch order was wrong"
462        );
463    }
464    shim.close_conversation().await.text()?;
465    drop(shim);
466    drop(adapter);
467    wait_for(|| app.path("stopped.log").exists(), "bounded shutdown").await?;
468    inspect(&app)
469}
470
471#[cfg(test)]
472mod tests {
473    use super::*;
474
475    #[tokio::test(flavor = "current_thread")]
476    async fn adapter_family_conforms() {
477        let total = [
478            include_str!("../Cargo.toml"),
479            include_str!("../Documentation.md"),
480            include_str!("lib.rs"),
481        ]
482        .iter()
483        .flat_map(|source| source.lines())
484        .filter(|line| !line.trim().is_empty())
485        .count();
486        assert!(total < 500, "package has {total} nonblank lines");
487        run_conformance().await.expect("adapter family conformance");
488    }
489}