use std::io::Write as _;
use std::process;
use harn_vm::llm::readiness::{probe_provider_readiness, ProviderReadiness};
use harn_vm::llm_config;
use serde::Serialize;
use crate::cli::{ProviderProbeArgs, ProviderToolProbeArgs, ProviderToolProbeModeArg};
use crate::commands::local::runtime::{fetch_ollama_ps, LoadedModel, LOCAL_PROVIDERS};
use crate::dispatch;
use crate::env_guard::ScopedEnvVar;
const PROBE_PAYLOAD_ENV: &str = "HARN_PROVIDER_PROBE_PAYLOAD_JSON";
const PROBE_PAYLOAD_PRETTY_ENV: &str = "HARN_PROVIDER_PROBE_PAYLOAD_PRETTY";
const TOOL_PROBE_PAYLOAD_ENV: &str = "HARN_PROVIDER_TOOL_PROBE_PAYLOAD_JSON";
const TOOL_PROBE_PAYLOAD_PRETTY_ENV: &str = "HARN_PROVIDER_TOOL_PROBE_PAYLOAD_PRETTY";
static DISPATCH_PROBE_LOCK: tokio::sync::Mutex<()> = tokio::sync::Mutex::const_new(());
static DISPATCH_TOOL_PROBE_LOCK: tokio::sync::Mutex<()> = tokio::sync::Mutex::const_new(());
#[derive(Debug, Serialize)]
struct ProviderProbe {
provider: String,
base_url: Option<String>,
readiness: ProviderReadiness,
#[serde(skip_serializing_if = "Option::is_none")]
runtime_profile: Option<harn_vm::llm::local_profiles::LocalRuntimeProfileReport>,
#[serde(skip_serializing_if = "Vec::is_empty")]
loaded_models: Vec<LoadedModel>,
}
pub(crate) async fn run_provider_probe(args: ProviderProbeArgs) {
let exit_code = dispatch_provider_probe(args).await;
if exit_code != 0 {
process::exit(exit_code);
}
}
async fn dispatch_provider_probe(args: ProviderProbeArgs) -> i32 {
let probe = aggregate_provider_probe(&args).await;
let exit_code = i32::from(!probe.readiness.ok);
let payload_json = match serde_json::to_string(&probe) {
Ok(json) => json,
Err(error) => {
eprintln!("error: failed to serialise provider-probe payload: {error}");
return 1;
}
};
let payload_pretty = match serde_json::to_string_pretty(&probe) {
Ok(json) => json,
Err(error) => {
eprintln!("error: failed to render provider-probe payload: {error}");
return 1;
}
};
let _guard = DISPATCH_PROBE_LOCK.lock().await;
let _payload_guard = ScopedEnvVar::set(PROBE_PAYLOAD_ENV, &payload_json);
let _pretty_guard = ScopedEnvVar::set(PROBE_PAYLOAD_PRETTY_ENV, &payload_pretty);
let outcome = dispatch::run_embedded_script("providers/probe", 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());
}
if outcome.exit_code != 0 {
outcome.exit_code
} else {
exit_code
}
}
async fn aggregate_provider_probe(args: &ProviderProbeArgs) -> ProviderProbe {
let readiness = probe_provider_readiness(
&args.provider,
args.model.as_deref(),
args.base_url.as_deref(),
)
.await;
let base_url = readiness.base_url.clone().or_else(|| {
llm_config::provider_config(&args.provider).map(|def| llm_config::resolve_base_url(&def))
});
let loaded_models = if args.provider == "ollama" {
let base = base_url
.clone()
.unwrap_or_else(|| "http://localhost:11434".to_string());
match fetch_ollama_ps(&base).await {
Ok(entries) => entries,
Err(error) => {
eprintln!("warning: /api/ps unavailable: {error}");
Vec::new()
}
}
} else {
Vec::new()
};
ProviderProbe {
provider: args.provider.clone(),
base_url,
readiness,
runtime_profile: if LOCAL_PROVIDERS.contains(&args.provider.as_str()) {
args.model.as_deref().map(|model| {
harn_vm::llm::local_profiles::local_runtime_profile_report(
model,
Some(&args.provider),
)
})
} else {
None
},
loaded_models,
}
}
pub(crate) async fn run_provider_tool_probe(args: ProviderToolProbeArgs) {
let exit_code = dispatch_provider_tool_probe(args).await;
if exit_code != 0 {
process::exit(exit_code);
}
}
async fn dispatch_provider_tool_probe(args: ProviderToolProbeArgs) -> i32 {
let report = match aggregate_tool_conformance_report(&args).await {
Ok(report) => report,
Err(error) => {
eprintln!("{error}");
return 1;
}
};
let payload_json = match serde_json::to_string(&report) {
Ok(json) => json,
Err(error) => {
eprintln!("error: failed to serialise tool-probe payload: {error}");
return 1;
}
};
let payload_pretty = match serde_json::to_string_pretty(&report) {
Ok(json) => json,
Err(error) => {
eprintln!("error: failed to render tool-probe payload: {error}");
return 1;
}
};
let fallback_disabled = report.tool_calling.fallback_mode
== harn_vm::llm::tool_conformance::ToolProbeFallbackMode::Disabled;
let _guard = DISPATCH_TOOL_PROBE_LOCK.lock().await;
let _payload_guard = ScopedEnvVar::set(TOOL_PROBE_PAYLOAD_ENV, &payload_json);
let _pretty_guard = ScopedEnvVar::set(TOOL_PROBE_PAYLOAD_PRETTY_ENV, &payload_pretty);
let outcome =
dispatch::run_embedded_script("providers/tool_probe", 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());
}
if outcome.exit_code != 0 {
return outcome.exit_code;
}
i32::from(fallback_disabled)
}
async fn aggregate_tool_conformance_report(
args: &ProviderToolProbeArgs,
) -> Result<harn_vm::llm::tool_conformance::ToolConformanceReport, String> {
if let Some(path) = args.response_fixture.as_ref() {
let raw = std::fs::read_to_string(path)
.map_err(|error| format!("error: failed to read {}: {error}", path.display()))?;
Ok(
harn_vm::llm::tool_conformance::classify_tool_conformance_fixture(
args.provider.clone(),
args.model.clone(),
modes_for_arg(args.mode)
.into_iter()
.next()
.unwrap_or(harn_vm::llm::tool_conformance::ToolProbeMode::NonStreaming),
args.marker.clone(),
&raw,
),
)
} else {
let mut options = harn_vm::llm::tool_conformance::ToolConformanceProbeOptions::new(
args.provider.clone(),
args.model.clone(),
);
options.base_url = args.base_url.clone();
options.modes = modes_for_arg(args.mode);
options.marker = args.marker.clone();
options.timeout_secs = args.timeout_secs;
Ok(harn_vm::llm::tool_conformance::run_tool_conformance_probe(options).await)
}
}
fn modes_for_arg(
mode: ProviderToolProbeModeArg,
) -> Vec<harn_vm::llm::tool_conformance::ToolProbeMode> {
use harn_vm::llm::tool_conformance::ToolProbeMode;
match mode {
ProviderToolProbeModeArg::Both => {
vec![ToolProbeMode::NonStreaming, ToolProbeMode::Streaming]
}
ProviderToolProbeModeArg::NonStreaming => vec![ToolProbeMode::NonStreaming],
ProviderToolProbeModeArg::Streaming => vec![ToolProbeMode::Streaming],
}
}