1//! `--send`: every gate written as a replay fixture. Layout parity and the
2//! unsure band are sent as `main.rs` builds them, through a
3//! `jev_mock::forward::Forward` directly; the label-tree and ranking runs
4//! are driven through the built extension in a throwaway Postgres whose
5//! endpoint is a forwarding jev-mock.
6//!
7//! The Postgres harness is `tests/support`, which finds the extension and
8//! the server through the environment `buck2 test` sets. This binary is
9//! built with the same values (`BUCK`), and re-executes itself with them
10//! when they are missing, since changing its own environment is unsafe.
11
12use std::path::Path;
13
14use jev_mock::MockJev;
15use jev_mock::forward::Forward;
16use support::postgres::TestPostgres;
17
18use support::gates::{LABEL_TREE_SQL, RANKING_SQL, setup_people};
19
20use crate::{Args, QUESTION, Run, TREE_PATHS};
21
22const REAL_ENDPOINT: &str = "https://api.typesafe.ai/v1/systemone";
23
24const LAYOUT_PARITY: &str = "layout_parity";
25const UNSURE_BAND: &str = "unsure_band";
26const LABEL_TREE: &str = "label_tree";
27const RANKING: &str = "ranking";
28
29/// Every fixture `record` writes, as `<out-dir>/<gate>.json`.
30const GATES: [&str; 4] = [LAYOUT_PARITY, UNSURE_BAND, LABEL_TREE, RANKING];
31
32fn fixture_path(args: &Args, gate: &str) -> std::path::PathBuf {
33    Path::new(&args.out_dir).join(format!("{gate}.json"))
34}
35
36/// Refuses a recording that would overwrite a gate already recorded,
37/// unless `--rerecord` asks for it: each recording spends live quota.
38pub fn refuse_recorded(args: &Args) -> Result<(), String> {
39    if args.rerecord {
40        return Ok(());
41    }
42    let recorded: Vec<String> =
43        GATES.iter().map(|g| fixture_path(args, g)).filter(|p| p.exists()).map(|p| p.display().to_string()).collect();
44    if recorded.is_empty() {
45        return Ok(());
46    }
47    Err(format!(
48        "already recorded, nothing sent: {}. The gates replay from these fixtures (tests/gate_*.rs); \
49         pass --rerecord to spend live quota again and overwrite them",
50        recorded.join(", ")
51    ))
52}
53
54/// Absent under cargo, which has no location of the built extension.
55const BUILT_ENV: [(&str, Option<&str>); 3] = [
56    ("POSTJEVSQL_EXT", option_env!("POSTJEVSQL_EXT")),
57    ("POSTJEVSQL_PG_MAJOR", option_env!("POSTJEVSQL_PG_MAJOR")),
58    ("POSTJEVSQL_POSTGRES_BIN", option_env!("POSTJEVSQL_POSTGRES_BIN")),
59];
60
61pub fn record(args: &Args, key: &str) -> Result<(), String> {
62    if BUILT_ENV.iter().any(|(name, _)| std::env::var_os(name).is_none()) {
63        let mut built = Vec::new();
64        for (name, value) in BUILT_ENV {
65            built.push((name, value.ok_or("--send needs the extension built by buck: buck2 run //tools/record-gates:record-gates")?));
66        }
67        let status = std::process::Command::new(std::env::current_exe().map_err(|e| e.to_string())?)
68            .args(std::env::args_os().skip(1))
69            .envs(built)
70            .status()
71            .map_err(|e| e.to_string())?;
72        return if status.success() { Ok(()) } else { Err("the recording run failed".into()) };
73    }
74    let runtime = tokio::runtime::Builder::new_multi_thread().enable_all().build().map_err(|e| e.to_string())?;
75    runtime.block_on(async {
76        direct(args, key, LAYOUT_PARITY, crate::layout_parity(&args.model)).await?;
77        direct(args, key, UNSURE_BAND, vec![crate::unsure_band(&args.model, 2)]).await?;
78        let paths: Vec<&str> = TREE_PATHS.to_vec();
79        let tree = gate(
80            args,
81            key,
82            LABEL_TREE,
83            LABEL_TREE_SQL,
84            Some(paths),
85        );
86        tree.await?;
87        gate(
88            args,
89            key,
90            RANKING,
91            RANKING_SQL,
92            None,
93        )
94        .await
95    })
96}
97
98/// Sends each distinct request of `runs` once and writes the fixture.
99async fn direct(args: &Args, key: &str, name: &str, runs: Vec<Run>) -> Result<(), String> {
100    let forward = Forward::new(REAL_ENDPOINT, key, args.model.as_str())?;
101    let mut sent: Vec<&[u8]> = Vec::new();
102    for request in runs.iter().flat_map(|run| &run.requests) {
103        if sent.contains(&request.as_slice()) {
104            continue;
105        }
106        sent.push(request);
107        let response = forward.exchange(request).await.map_err(|e| format!("{name}: {e}"))?;
108        if response.status != 200 {
109            return Err(format!("{name}: the endpoint answered {}: {}", response.status, response.body));
110        }
111    }
112    write(args, name, &forward, None)
113}
114
115fn write(args: &Args, name: &str, forward: &Forward, answer: Option<String>) -> Result<(), String> {
116    let fixture = forward.fixture();
117    let out = fixture_path(args, name);
118    std::fs::create_dir_all(&args.out_dir).map_err(|e| e.to_string())?;
119    // Without --rerecord, never replace a fixture another run wrote since
120    // `refuse_recorded` looked.
121    let mut open = std::fs::OpenOptions::new();
122    open.write(true);
123    if args.rerecord {
124        open.create(true).truncate(true);
125    } else {
126        open.create_new(true);
127    }
128    let mut file = open.open(&out).map_err(|e| format!("{}: {e}", out.display()))?;
129    std::io::Write::write_all(&mut file, (fixture.to_json() + "\n").as_bytes()).map_err(|e| format!("{}: {e}", out.display()))?;
130    println!("{name}: {} request(s), answer {answer:?}, wrote {}", fixture.exchanges.len(), out.display());
131    Ok(())
132}
133
134/// Runs `sql` once through a recording mock and writes the fixture.
135async fn gate(args: &Args, key: &str, name: &str, sql: &str, paths: Option<Vec<&str>>) -> Result<(), String> {
136    let forward = Forward::new(REAL_ENDPOINT, key, args.model.as_str())?;
137    let mock = MockJev::start(forward.responder()).await;
138    let pg = TestPostgres::new()
139        .with_extension(Path::new(&std::env::var("POSTJEVSQL_EXT").unwrap()))
140        .setting("jev.endpoint", &mock.endpoint())
141        .setting("jev.model", args.model.as_str())
142        .setting("jev.ca_file", mock.ca_file().to_str().unwrap())
143        .setting("jev.price_per_mtok", &args.price_per_mtok.to_string())
144        .setting("jev.max_cost", &args.max_cost.to_string())
145        // The mock adds the real key; the server's never leaves loopback.
146        .env_var("TYPESAFE_API_KEY", "record-gates")
147        .start()
148        .await;
149    let (client, conn) = tokio_postgres::connect(&pg.conn_str(), tokio_postgres::NoTls).await.map_err(|e| e.to_string())?;
150    tokio::spawn(conn);
151    setup_people(&client).await.map_err(|e| e.to_string())?;
152
153    let rows = match &paths {
154        Some(paths) => client.query(sql, &[&QUESTION, paths]).await,
155        None => client.query(sql, &[&QUESTION]).await,
156    }
157    .map_err(|e| format!("{name}: {}", e.as_db_error().map_or(e.to_string(), |d| d.message().to_owned())))?;
158    let answer: Option<String> = rows.first().and_then(|r| r.get(0));
159    write(args, name, &forward, answer)
160}