use std::io::Write as _;
use std::process;
use serde::Serialize;
use crate::cli::ModelsTestArgs;
use crate::dispatch;
use crate::env_guard::ScopedEnvVar;
const TEST_RESULT_ENV: &str = "HARN_MODELS_TEST_RESULT_JSON";
static DISPATCH_TEST_LOCK: tokio::sync::Mutex<()> = tokio::sync::Mutex::const_new(());
#[derive(Debug, Serialize)]
struct TestEnvelope<'a> {
ok: bool,
#[serde(skip_serializing_if = "Option::is_none")]
result: Option<&'a harn_vm::llm::ModelSmokeTestResult>,
#[serde(skip_serializing_if = "Option::is_none")]
error: Option<&'a str>,
}
pub(crate) async fn run(args: &ModelsTestArgs) {
if std::env::var("HARN_CLI_IMPL").as_deref() == Ok("rust") {
run_legacy(args).await;
return;
}
let exit_code = run_dispatch(args).await;
if exit_code != 0 {
process::exit(exit_code);
}
}
async fn run_dispatch(args: &ModelsTestArgs) -> i32 {
let result = harn_vm::llm::run_model_smoke_test(harn_vm::llm::ModelSmokeTestOptions {
model: args.model.clone(),
provider: args.provider.clone(),
prompt: args.prompt.clone(),
})
.await;
let envelope = match &result {
Ok(value) => TestEnvelope {
ok: true,
result: Some(value),
error: None,
},
Err(error) => TestEnvelope {
ok: false,
result: None,
error: Some(error.as_str()),
},
};
let envelope_json = match serde_json::to_string(&envelope) {
Ok(json) => json,
Err(error) => {
eprintln!("error: failed to serialise models-test envelope: {error}");
return 1;
}
};
let _guard = DISPATCH_TEST_LOCK.lock().await;
let _payload = ScopedEnvVar::set(TEST_RESULT_ENV, &envelope_json);
let outcome = dispatch::run_embedded_script("models/test", Vec::new(), args.json).await;
if !outcome.stderr.is_empty() {
let _ = std::io::stderr().write_all(outcome.stderr.as_bytes());
}
if !outcome.stdout.is_empty() {
let _ = std::io::stdout().write_all(outcome.stdout.as_bytes());
}
outcome.exit_code
}
async fn run_legacy(args: &ModelsTestArgs) {
let result = harn_vm::llm::run_model_smoke_test(harn_vm::llm::ModelSmokeTestOptions {
model: args.model.clone(),
provider: args.provider.clone(),
prompt: args.prompt.clone(),
})
.await;
match result {
Ok(result) if args.json => match serde_json::to_string_pretty(&result) {
Ok(payload) => println!("{payload}"),
Err(error) => {
crate::command_error(&format!("failed to serialize model test result: {error}"))
}
},
Ok(result) => {
let first_token = result
.first_token_ms
.map(|value| value.to_string())
.unwrap_or_else(|| "-".to_string());
println!(
"model_id={} provider={} latency_ms={} first_token_ms={} input_tokens={} output_tokens={} estimated_cost_usd={:.6}",
result.model_id,
result.provider,
result.latency_ms,
first_token,
result.input_tokens,
result.output_tokens,
result.estimated_cost_usd
);
}
Err(error) if args.json => {
println!("{}", serde_json::json!({ "ok": false, "error": error }));
process::exit(1);
}
Err(error) => {
eprintln!("{error}");
process::exit(1);
}
}
}