use std::collections::BTreeMap;
use std::sync::Mutex;
use std::time::Duration;
use super::wire_format::WireFormat;
pub const PREFLIGHT_HEADER: &str = "x-openlatch-preflight";
const UPSTREAM_UNREACHABLE_HEADER: &str = "x-openlatch-upstream";
pub const PREFLIGHT_TIMEOUT: Duration = Duration::from_secs(5);
const PREFLIGHT_BODY: &str =
r#"{"model":"claude-sonnet-4-5","max_tokens":1,"messages":[{"role":"user","content":"ping"}]}"#;
const PREFLIGHT_BODY_RESPONSES: &str =
r#"{"model":"gpt-5-codex","input":"preflight","max_output_tokens":16}"#;
const PREFLIGHT_BODY_CHAT_COMPLETIONS: &str = r#"{"model":"openlatch-preflight","max_tokens":1,"messages":[{"role":"user","content":"ping"}]}"#;
const PREFLIGHT_BODY_GENERATE_CONTENT: &str =
r#"{"contents":[{"role":"user","parts":[{"text":"ping"}]}]}"#;
const PREFLIGHT_BODY_OLLAMA_NATIVE: &str =
r#"{"model":"openlatch-preflight","messages":[],"stream":false}"#;
const PREFLIGHT_GOOGLE_MODEL: &str = "openlatch-preflight";
const PREFLIGHT_API_KEY: &str = "ol-preflight-not-a-key";
#[derive(Debug, Clone, Default, PartialEq, Eq)]
pub enum Verdict {
#[default]
Pending,
Ok,
Failed(String),
}
impl Verdict {
pub fn label(&self) -> &'static str {
match self {
Self::Pending => "pending",
Self::Ok => "ok",
Self::Failed(_) => "failed",
}
}
pub fn error(&self) -> Option<&str> {
match self {
Self::Failed(e) => Some(e.as_str()),
_ => None,
}
}
pub fn is_ok(&self) -> bool {
matches!(self, Self::Ok)
}
}
#[derive(Debug, Default)]
pub struct WiringState {
wired: Mutex<BTreeMap<&'static str, bool>>,
verdict: Mutex<BTreeMap<&'static str, Verdict>>,
format: Mutex<BTreeMap<&'static str, WireFormat>>,
endpoints: std::sync::Mutex<BTreeMap<String, EndpointVerdict>>,
}
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct EndpointVerdict {
pub code: &'static str,
pub detail: String,
}
impl WiringState {
fn read<T: Clone>(m: &Mutex<BTreeMap<&'static str, T>>) -> BTreeMap<&'static str, T> {
match m.lock() {
Ok(v) => v.clone(),
Err(poisoned) => poisoned.into_inner().clone(),
}
}
fn write<T>(
m: &Mutex<BTreeMap<&'static str, T>>,
f: impl FnOnce(&mut BTreeMap<&'static str, T>),
) {
match m.lock() {
Ok(mut v) => f(&mut v),
Err(poisoned) => f(&mut poisoned.into_inner()),
}
}
pub fn is_wired(&self, agent: &str) -> bool {
match self.wired.lock() {
Ok(v) => v.get(agent).copied().unwrap_or(false),
Err(poisoned) => poisoned.into_inner().get(agent).copied().unwrap_or(false),
}
}
pub fn set_wired(&self, agent: &'static str, wired: bool) {
Self::write(&self.wired, |m| {
m.insert(agent, wired);
});
}
pub fn set_wired_format(&self, agent: &'static str, format: WireFormat) {
Self::write(&self.format, |m| {
m.insert(agent, format);
});
}
pub fn wired_format(&self, agent: &str) -> Option<WireFormat> {
match self.format.lock() {
Ok(v) => v.get(agent).copied(),
Err(poisoned) => poisoned.into_inner().get(agent).copied(),
}
}
pub fn sole_wired_agent_for(&self, format: WireFormat) -> Option<&'static str> {
if !format.is_captured() {
return None;
}
let wired = Self::read(&self.wired);
let mut hit = None;
for (agent, agent_format) in Self::read(&self.format) {
if agent_format != format || !wired.get(agent).copied().unwrap_or(false) {
continue;
}
if hit.is_some() {
return None;
}
hit = Some(agent);
}
hit
}
pub fn set_endpoint_verdict(&self, key: &str, verdict: Option<EndpointVerdict>) {
let mut map = match self.endpoints.lock() {
Ok(m) => m,
Err(poisoned) => poisoned.into_inner(),
};
match verdict {
Some(v) => {
map.insert(key.to_string(), v);
}
None => {
map.remove(key);
}
}
}
pub fn endpoint_verdicts(&self) -> BTreeMap<String, EndpointVerdict> {
match self.endpoints.lock() {
Ok(m) => m.clone(),
Err(poisoned) => poisoned.into_inner().clone(),
}
}
pub fn verdict(&self, agent: &str) -> Verdict {
match self.verdict.lock() {
Ok(v) => v.get(agent).cloned().unwrap_or_default(),
Err(poisoned) => poisoned
.into_inner()
.get(agent)
.cloned()
.unwrap_or_default(),
}
}
pub fn set_verdict(&self, agent: &'static str, verdict: Verdict) {
Self::write(&self.verdict, |m| {
m.insert(agent, verdict);
});
}
pub fn verdicts(&self) -> BTreeMap<&'static str, Verdict> {
Self::read(&self.verdict)
}
pub fn wired_agents(&self) -> BTreeMap<&'static str, bool> {
Self::read(&self.wired)
}
pub fn seed(&self, agent: &'static str) {
Self::write(&self.wired, |m| {
m.entry(agent).or_insert(false);
});
Self::write(&self.verdict, |m| {
m.entry(agent).or_default();
});
}
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
enum ProbeShape {
AnthropicMessages,
OpenAiResponses,
OpenAiChatCompletions,
GoogleGenerateContent,
OllamaNative,
Reachability,
}
impl ProbeShape {
fn route(self) -> String {
match self {
Self::AnthropicMessages => "/v1/messages".to_string(),
Self::OpenAiResponses => "/v1/responses".to_string(),
Self::OpenAiChatCompletions => "/v1/chat/completions".to_string(),
Self::GoogleGenerateContent => {
format!("/v1beta/models/{PREFLIGHT_GOOGLE_MODEL}:generateContent")
}
Self::OllamaNative => "/api/chat".to_string(),
Self::Reachability => "/".to_string(),
}
}
fn method(self) -> reqwest::Method {
match self {
Self::Reachability => reqwest::Method::GET,
_ => reqwest::Method::POST,
}
}
fn body(self) -> &'static str {
match self {
Self::AnthropicMessages => PREFLIGHT_BODY,
Self::OpenAiResponses => PREFLIGHT_BODY_RESPONSES,
Self::OpenAiChatCompletions => PREFLIGHT_BODY_CHAT_COMPLETIONS,
Self::GoogleGenerateContent => PREFLIGHT_BODY_GENERATE_CONTENT,
Self::OllamaNative => PREFLIGHT_BODY_OLLAMA_NATIVE,
Self::Reachability => "",
}
}
fn auth_headers(self) -> &'static [(&'static str, &'static str)] {
match self {
Self::AnthropicMessages => &[
("anthropic-version", "2023-06-01"),
("x-api-key", PREFLIGHT_API_KEY),
],
Self::OpenAiResponses | Self::OpenAiChatCompletions => &[(
"authorization",
concat!("Bearer ", "ol-preflight-not-a-key"),
)],
Self::GoogleGenerateContent => &[("x-goog-api-key", PREFLIGHT_API_KEY)],
Self::OllamaNative | Self::Reachability => &[],
}
}
fn of(fmt: WireFormat) -> Self {
match fmt {
WireFormat::AnthropicMessages => Self::AnthropicMessages,
WireFormat::OpenAiResponses => Self::OpenAiResponses,
WireFormat::OpenAiChatCompletions => Self::OpenAiChatCompletions,
WireFormat::GoogleGenerateContent => Self::GoogleGenerateContent,
WireFormat::OllamaNative => Self::OllamaNative,
WireFormat::Unknown => Self::Reachability,
}
}
}
pub async fn probe(
port: u16,
fmt: WireFormat,
upstream: &str,
timeout: Duration,
) -> Result<(), String> {
let client = match crate::egress::client_builder().timeout(timeout).build() {
Ok(c) => c,
Err(e) => return Err(format!("could not build the preflight client: {e}")),
};
let shape = ProbeShape::of(fmt);
let url = format!("http://127.0.0.1:{port}{}", shape.route());
let mut req = client
.request(shape.method(), &url)
.header(PREFLIGHT_HEADER, "1");
if !shape.body().is_empty() {
req = req
.header("content-type", "application/json")
.body(shape.body());
}
for (name, value) in shape.auth_headers() {
req = req.header(*name, *value);
}
let sent = req.send().await;
let resp = match sent {
Ok(r) => r,
Err(e) if e.is_timeout() => {
return Err(format!(
"no response from the model relay on 127.0.0.1:{port} within {}s",
timeout.as_secs()
))
}
Err(e) => {
return Err(format!(
"could not reach the model relay on 127.0.0.1:{port}: {}",
e.without_url()
))
}
};
if resp.headers().contains_key(UPSTREAM_UNREACHABLE_HEADER) {
return Err(format!(
"the model relay is listening but could not reach {upstream} — model calls would fail"
));
}
Ok(())
}
#[cfg(test)]
mod tests {
use super::*;
use crate::model_relay::{mock, serve_ephemeral, ModelRelayState};
use std::sync::Arc;
fn state_for(upstream_port: u16) -> Arc<ModelRelayState> {
let base = reqwest::Url::parse(&format!("http://127.0.0.1:{upstream_port}")).unwrap();
Arc::new(ModelRelayState::new(base, 0, 8, &[]))
}
#[test]
fn probe_paths_resolve_to_their_own_format() {
for fmt in [
WireFormat::AnthropicMessages,
WireFormat::OpenAiResponses,
WireFormat::OpenAiChatCompletions,
WireFormat::GoogleGenerateContent,
WireFormat::OllamaNative,
] {
let route = ProbeShape::of(fmt).route();
let (parts, _) = axum::http::Request::builder()
.method(axum::http::Method::POST)
.uri(format!("http://127.0.0.1:1{route}"))
.body(())
.expect("a request")
.into_parts();
assert_eq!(
WireFormat::resolve(&parts),
fmt,
"{fmt:?}'s probe route {route} must resolve to {fmt:?}, or the probe \
is forwarded to another format's upstream"
);
}
assert!(
ProbeShape::of(WireFormat::GoogleGenerateContent)
.route()
.ends_with(":generateContent"),
"the colon suffix is what the Google arm matches on — a bare \
/v1beta/models resolves Unknown and goes to Anthropic"
);
}
#[tokio::test(flavor = "multi_thread")]
async fn preflight_probes_each_format_in_its_own_shape() {
let cases = [
(
WireFormat::OpenAiChatCompletions,
"/v1/chat/completions",
"\"messages\"",
"\"input\"",
),
(
WireFormat::GoogleGenerateContent,
":generateContent",
"\"contents\"",
"\"messages\"",
),
(
WireFormat::OllamaNative,
"POST /api/chat",
"\"stream\":false",
"max_tokens",
),
];
for (fmt, route_marker, expected, forbidden) in cases {
let upstream = mock::spawn_capture_200().await;
let port = serve_ephemeral(state_for(upstream.port)).await;
probe(port, fmt, "http://127.0.0.1", PREFLIGHT_TIMEOUT)
.await
.unwrap_or_else(|e| panic!("{fmt:?} probe: {e}"));
let line = upstream
.received_request_line
.lock()
.expect("the mock recorded a request line")
.clone()
.expect("a request arrived");
assert!(
line.contains(route_marker),
"{fmt:?} must probe its own route, got: {line}"
);
let body = String::from_utf8(
upstream
.received_body
.lock()
.expect("the mock recorded a body")
.clone()
.expect("a body arrived"),
)
.expect("utf8");
assert!(
body.contains(expected),
"{fmt:?} must probe in its own body shape, got: {body}"
);
assert!(
!body.contains(forbidden),
"{fmt:?} must not carry another format's body, got: {body}"
);
}
}
#[tokio::test(flavor = "multi_thread")]
async fn ollama_native_probe_posts_api_chat_and_unknown_probes_reachability() {
let upstream = mock::spawn_capture_200().await;
let port = serve_ephemeral(state_for(upstream.port)).await;
probe(
port,
WireFormat::Unknown,
"http://127.0.0.1",
PREFLIGHT_TIMEOUT,
)
.await
.expect("reachability probe");
let line = upstream
.received_request_line
.lock()
.expect("lock")
.clone()
.expect("a request arrived");
assert!(line.starts_with("GET / "), "{line}");
assert_eq!(upstream.header("x-api-key"), None);
assert_eq!(upstream.header("authorization"), None);
}
#[tokio::test(flavor = "multi_thread")]
async fn probe_passes_when_upstream_answers() {
let upstream = mock::spawn_capture_200().await;
let port = serve_ephemeral(state_for(upstream.port)).await;
assert_eq!(
probe(
port,
WireFormat::AnthropicMessages,
crate::model_relay::ANTHROPIC_BASE,
PREFLIGHT_TIMEOUT
)
.await,
Ok(())
);
}
#[tokio::test(flavor = "multi_thread")]
async fn probe_fails_when_upstream_is_unreachable() {
let dead = mock::closed_port().await;
let port = serve_ephemeral(state_for(dead)).await;
let err = probe(
port,
WireFormat::AnthropicMessages,
crate::model_relay::ANTHROPIC_BASE,
PREFLIGHT_TIMEOUT,
)
.await
.unwrap_err();
assert!(
err.contains("could not reach"),
"an unreachable upstream must be named as such, got: {err}"
);
}
#[tokio::test(flavor = "multi_thread")]
async fn probe_fails_fast_on_a_silent_upstream() {
let hung = mock::spawn_hang_after_accept().await;
let upstream = reqwest::Url::parse(&format!("http://127.0.0.1:{hung}")).unwrap();
let state = Arc::new(
ModelRelayState::new(upstream, 0, 8, &[]).with_header_timeout(Duration::from_secs(60)),
);
let port = serve_ephemeral(state).await;
let started = std::time::Instant::now();
let err = probe(
port,
WireFormat::AnthropicMessages,
crate::model_relay::ANTHROPIC_BASE,
Duration::from_millis(300),
)
.await
.unwrap_err();
assert!(
err.contains("no response"),
"a silent upstream must read as no response, got: {err}"
);
assert!(
started.elapsed() < Duration::from_secs(5),
"the probe must return on its own budget, not the forward path's"
);
}
#[tokio::test(flavor = "multi_thread")]
async fn probe_fails_when_nothing_is_listening() {
let port = mock::closed_port().await;
assert!(probe(
port,
WireFormat::AnthropicMessages,
crate::model_relay::ANTHROPIC_BASE,
Duration::from_millis(500)
)
.await
.is_err());
}
#[tokio::test(flavor = "multi_thread")]
async fn probe_sends_the_format_it_was_given() {
let upstream = mock::spawn_capture_200().await;
let port = serve_ephemeral(state_for(upstream.port)).await;
assert_eq!(
probe(
port,
WireFormat::OpenAiResponses,
crate::model_relay::wire_format::OPENAI_BASE,
PREFLIGHT_TIMEOUT
)
.await,
Ok(())
);
let line = upstream
.received_request_line
.lock()
.unwrap()
.clone()
.expect("the mock recorded the request line");
assert!(
line.starts_with("POST /v1/responses"),
"a Responses probe must speak the Responses route, got: {line}"
);
assert_eq!(
upstream.header("anthropic-version"),
None,
"an Anthropic protocol header on an OpenAI request is the exact defect this gate exists to catch"
);
}
#[tokio::test(flavor = "multi_thread")]
async fn probe_carries_no_real_credential() {
let anthropic_up = mock::spawn_capture_200().await;
let port = serve_ephemeral(state_for(anthropic_up.port)).await;
assert_eq!(
probe(
port,
WireFormat::AnthropicMessages,
crate::model_relay::ANTHROPIC_BASE,
PREFLIGHT_TIMEOUT
)
.await,
Ok(())
);
assert_eq!(
anthropic_up.header("x-api-key").as_deref(),
Some(PREFLIGHT_API_KEY),
"the Anthropic probe must send the not-a-key literal"
);
let responses_up = mock::spawn_capture_200().await;
let port = serve_ephemeral(state_for(responses_up.port)).await;
assert_eq!(
probe(
port,
WireFormat::OpenAiResponses,
crate::model_relay::wire_format::OPENAI_BASE,
PREFLIGHT_TIMEOUT
)
.await,
Ok(())
);
assert_eq!(
responses_up.header("authorization").as_deref(),
Some(format!("Bearer {PREFLIGHT_API_KEY}").as_str()),
"the Responses probe must send the same not-a-key literal as a Bearer"
);
}
#[test]
fn verdict_labels_are_stable() {
assert_eq!(Verdict::default(), Verdict::Pending);
assert_eq!(Verdict::Pending.label(), "pending");
assert_eq!(Verdict::Ok.label(), "ok");
assert_eq!(Verdict::Failed("boom".into()).label(), "failed");
assert_eq!(Verdict::Failed("boom".into()).error(), Some("boom"));
assert_eq!(Verdict::Ok.error(), None);
assert!(Verdict::Ok.is_ok());
assert!(!Verdict::Pending.is_ok());
}
#[test]
fn wiring_state_round_trips() {
let st = WiringState::default();
assert!(!st.is_wired("claude-code"));
assert_eq!(st.verdict("claude-code"), Verdict::Pending);
st.set_wired("claude-code", true);
st.set_verdict("claude-code", Verdict::Ok);
assert!(st.is_wired("claude-code"));
assert_eq!(st.verdict("claude-code"), Verdict::Ok);
}
#[test]
fn wiring_state_is_keyed_per_agent() {
let st = WiringState::default();
st.set_wired("claude-code", true);
assert!(st.is_wired("claude-code"));
assert!(
!st.is_wired("codex-cli"),
"one agent's wiring says nothing about another's"
);
st.set_verdict("claude-code", Verdict::Ok);
assert_eq!(
st.verdict("codex-cli"),
Verdict::Pending,
"an agent with no entry is Pending — not the other agent's verdict"
);
st.set_verdict("codex-cli", Verdict::Failed("no round trip".into()));
let verdicts = st.verdicts();
assert_eq!(verdicts.get("claude-code"), Some(&Verdict::Ok));
assert_eq!(
verdicts.get("codex-cli"),
Some(&Verdict::Failed("no round trip".into())),
"the snapshot the admin surface renders carries every agent"
);
let wired = st.wired_agents();
assert_eq!(wired.get("claude-code"), Some(&true));
assert_eq!(
wired.get("codex-cli"),
None,
"a verdict is not a write: codex-cli was judged, never wired"
);
st.seed("cursor");
st.seed("claude-code");
assert_eq!(st.verdict("cursor"), Verdict::Pending);
assert_eq!(st.wired_agents().get("cursor"), Some(&false));
assert_eq!(
st.verdict("claude-code"),
Verdict::Ok,
"re-seeding must never overwrite a verdict the supervisor recorded"
);
assert_eq!(st.wired_agents().get("claude-code"), Some(&true));
}
#[test]
fn a_lone_speaker_is_named() {
let st = WiringState::default();
st.set_wired("claude-code", true);
st.set_wired_format("claude-code", WireFormat::AnthropicMessages);
st.set_wired("codex-cli", true);
st.set_wired_format("codex-cli", WireFormat::OpenAiResponses);
assert_eq!(
st.sole_wired_agent_for(WireFormat::AnthropicMessages),
Some("claude-code")
);
assert_eq!(
st.sole_wired_agent_for(WireFormat::OpenAiResponses),
Some("codex-cli")
);
}
#[test]
fn two_speakers_of_one_format_name_nobody() {
let st = WiringState::default();
st.set_wired("claude-code", true);
st.set_wired_format("claude-code", WireFormat::AnthropicMessages);
st.set_wired("cline", true);
st.set_wired_format("cline", WireFormat::AnthropicMessages);
assert_eq!(st.sole_wired_agent_for(WireFormat::AnthropicMessages), None);
}
#[test]
fn an_unwired_agent_does_not_speak() {
let st = WiringState::default();
st.set_wired_format("claude-code", WireFormat::AnthropicMessages);
st.set_wired("claude-code", false);
assert_eq!(st.sole_wired_agent_for(WireFormat::AnthropicMessages), None);
st.set_wired("cline", true);
st.set_wired_format("cline", WireFormat::AnthropicMessages);
assert_eq!(
st.sole_wired_agent_for(WireFormat::AnthropicMessages),
Some("cline"),
"the unwired agent must not block the one that is actually wired"
);
}
#[test]
fn no_speaker_names_nobody() {
let st = WiringState::default();
st.set_wired("claude-code", true);
st.set_wired_format("claude-code", WireFormat::AnthropicMessages);
assert_eq!(st.sole_wired_agent_for(WireFormat::OpenAiResponses), None);
}
#[test]
fn the_uncaptured_format_is_never_attributable() {
let st = WiringState::default();
st.set_wired("claude-code", true);
st.set_wired_format("claude-code", WireFormat::Unknown);
assert_eq!(st.sole_wired_agent_for(WireFormat::Unknown), None);
}
}