use std::pin::Pin;
use std::sync::atomic::{AtomicBool, AtomicU64, Ordering};
use std::sync::Arc;
use std::task::{Context, Poll};
use axum::body::Body;
use axum::extract::{Request, State};
use axum::http::header::{
ACCEPT_ENCODING, CONNECTION, CONTENT_LENGTH, HOST, PROXY_AUTHENTICATE, PROXY_AUTHORIZATION, TE,
TRAILER, TRANSFER_ENCODING, UPGRADE,
};
use axum::http::{request::Parts, HeaderMap, HeaderValue, StatusCode};
use axum::response::Response;
use axum::Json;
use bytes::Bytes;
use futures_core::Stream;
use tokio::sync::mpsc::Sender;
use tokio::sync::OwnedSemaphorePermit;
use crate::cloud::CloudEvent;
use crate::privacy::PrivacyFilter;
use super::billing::detect_billing;
use super::capture::{self, CaptureGap, CostBasis, Usage, UsageAccumulator};
use super::emit::{self, Observation};
use super::preflight::PREFLIGHT_HEADER;
use super::session;
use super::tokenize::{classify_model, Estimator};
use super::wire_format::{self, AuthMode, WireFormat};
use super::{ModelRelayState, MAX_MATERIALIZE_BYTES};
use crate::generated::types::Verdict;
use crate::zone_eval::types::OptimizeParams;
use crate::zone_eval::{
AgentContext as ZoneAgentContext, Decision as ZoneDecision, Event as ZoneEvent,
};
struct I10Outcome {
bytes: Option<Bytes>,
applied: Option<super::transforms::interventions::AppliedIntervention>,
decision: ZoneDecision,
blocked: bool,
policy_state: Option<crate::zone_eval::SessionState>,
policy_layout: crate::zone_eval::StateLayout,
policy_session_id: String,
second_optimize_degraded: bool,
}
enum I10Attempt {
Absent,
Handled(Box<I10Outcome>),
Failed,
}
struct I10Delivery {
applied: Option<super::transforms::interventions::AppliedIntervention>,
policy_state: Option<crate::zone_eval::SessionState>,
policy_layout: crate::zone_eval::StateLayout,
policy_session_id: String,
}
const INSTALL_ID_HEADER: &str = "x-openlatch-install-id";
static PASS_THROUGH_FAILURES: AtomicU64 = AtomicU64::new(0);
static UPSTREAM_FAILURES: AtomicU64 = AtomicU64::new(0);
static INJECT_OBSERVE_PANIC: AtomicBool = AtomicBool::new(false);
#[cfg(test)]
static INJECT_SCAN_PANIC: AtomicBool = AtomicBool::new(false);
pub fn pass_through_failures() -> u64 {
PASS_THROUGH_FAILURES.load(Ordering::Relaxed)
}
pub fn upstream_failures() -> u64 {
UPSTREAM_FAILURES.load(Ordering::Relaxed)
}
pub fn set_inject_observe_panic(on: bool) {
INJECT_OBSERVE_PANIC.store(on, Ordering::Relaxed);
}
#[cfg(test)]
pub fn set_inject_scan_panic(on: bool) {
INJECT_SCAN_PANIC.store(on, Ordering::Relaxed);
}
fn record_pass_through_failure(reason: &'static str) {
PASS_THROUGH_FAILURES.fetch_add(1, Ordering::Relaxed);
tracing::warn!(
reason,
"model relay pass-through failure — forwarding unmodified"
);
}
fn observe_request(
st: &ModelRelayState,
wire: WireFormat,
headers: &HeaderMap,
path: &str,
bytes: &Bytes,
) -> (Observation, Option<serde_json::Value>) {
if INJECT_OBSERVE_PANIC.load(Ordering::Relaxed) {
panic!("injected observe panic (D-24 bench)");
}
let occurred_at = crate::envelope::current_timestamp();
let event_id = uuid::Uuid::now_v7().to_string();
let body: serde_json::Value = serde_json::from_slice(bytes).unwrap_or(serde_json::Value::Null);
let (model, model_known) = match capture::model_of(&body) {
Some(m) => {
let known = classify_model(&m).is_some();
(Some(m), known)
}
None => match capture::model_from_path(wire, path) {
Some(m) => (Some(m), true),
None => (None, false),
},
};
let billing = detect_billing(headers);
let pricing = capture::derive_pricing_inputs(&body, headers);
let install_id = request_install_id(st, headers);
let signals = request_signals(wire, &body);
let session =
session::resolve_session_scoped(&st.registry, &install_id, &signals, endpoint_agent(st));
let session_key = session.session_id.clone().unwrap_or_default();
let churn = st.churn.observe(&install_id, &session_key, bytes);
if let Some(f) = &churn {
retention_store(f, &occurred_at);
}
let resident = st.resident_request_rules();
let authored: &[crate::generated::types::PolicyRule] = resident
.as_ref()
.and_then(|guard| guard.as_ref().as_ref())
.map(|bundle| bundle.request_rules.as_slice())
.unwrap_or(&[]);
let transform = super::transforms::evaluate_would_have_with(&body, authored);
let carries_i10 = resident
.as_ref()
.and_then(|guard| guard.as_ref().as_ref())
.and_then(|bundle| bundle.zone.as_ref())
.is_some();
let carry = (st.transforms_act || carries_i10).then_some(body);
(
Observation {
measured: true,
event_id,
occurred_at,
model,
model_known,
billing,
install_id: install_id.clone(),
session,
pricing,
churn,
request_body_len: bytes.len(),
has_breakpoint: capture::has_cache_breakpoint(bytes),
transform,
optimize: None,
wire_format: wire,
attributable_agent: attributable_agent(st, wire),
},
carry,
)
}
fn apply_i10(
st: &ModelRelayState,
obs: &Observation,
body: &serde_json::Value,
) -> Option<I10Outcome> {
let resident = st.policy.as_ref()?.load_full();
let resident = resident.as_ref().as_ref()?;
let zone = resident.zone.as_ref()?;
let session_id = obs.session.session_id.as_deref().unwrap_or(&obs.install_id);
let now = std::time::Instant::now();
let now_ms = chrono::Utc::now().timestamp_millis();
let event_for = |input: serde_json::Value| ZoneEvent {
event_type: "pre_tool_use".into(),
session_id: session_id.to_string(),
tool_use_id: obs.event_id.clone(),
tool_name: "ModelRequest".into(),
tool_input: input,
tool_result: None,
agent: Some(ZoneAgentContext {
agent_id: obs
.session
.agent_id
.clone()
.or_else(|| Some(obs.install_id.clone())),
agent_type: Some("coding_assistant".into()),
..ZoneAgentContext::default()
}),
binding: serde_json::json!({"kind":"model_relay","mode":1}),
env: None,
session: None,
spend_delta: None,
};
let prior = st
.policy_sessions
.preview(session_id, &zone.state_layout, now);
let (outcome, _) = (|| {
let prior = prior.as_ref();
let original_event = event_for(body.clone());
let (initial, initial_state) =
crate::zone_eval::evaluate(zone, &original_event, prior, now_ms);
let mut selected = initial.optimize.clone();
let Some(actual) = initial
.optimize
.as_ref()
.and_then(|ctx| ctx.actual.as_ref())
else {
let commit_original = initial.verdict != Verdict::Block;
return (
I10Outcome {
bytes: None,
applied: None,
blocked: initial.verdict == Verdict::Block,
decision: initial,
policy_state: commit_original.then_some(initial_state.clone()),
policy_layout: zone.state_layout,
policy_session_id: session_id.to_string(),
second_optimize_degraded: false,
},
initial_state,
);
};
let directive = actual
.artifact_id
.as_deref()
.and_then(|id| {
zone.artifacts
.iter()
.find(|artifact| artifact.artifact_id() == Some(id))
})
.and_then(|artifact| artifact.optimize.as_ref());
let Some(directive) = directive else {
return (
I10Outcome {
bytes: None,
applied: None,
blocked: false,
decision: initial,
policy_state: None,
policy_layout: zone.state_layout,
policy_session_id: session_id.to_string(),
second_optimize_degraded: false,
},
initial_state,
);
};
let capable = matches!(
directive.params,
OptimizeParams::EffortClamp(_)
| OptimizeParams::ContextEdit(_)
| OptimizeParams::PrefixGuard(_)
);
if !capable {
if let Some(candidate) = selected.as_mut().and_then(|ctx| ctx.actual.as_mut()) {
candidate.reason = "incapable".into();
}
let mut decision = initial;
decision.optimize = selected;
return (
I10Outcome {
bytes: None,
applied: None,
blocked: false,
decision,
policy_state: None,
policy_layout: zone.state_layout,
policy_session_id: session_id.to_string(),
second_optimize_degraded: false,
},
initial_state,
);
}
if st.inject_mutate_panic {
panic!("injected I-10 mutation panic");
}
if !obs.model_known {
if let Some(candidate) = selected.as_mut().and_then(|ctx| ctx.actual.as_mut()) {
candidate.reason = "incapable".into();
}
let mut decision = initial;
decision.optimize = selected;
return (
I10Outcome {
bytes: None,
applied: None,
blocked: false,
decision,
policy_state: None,
policy_layout: zone.state_layout,
policy_session_id: session_id.to_string(),
second_optimize_degraded: false,
},
initial_state,
);
}
let session_key = format!("{}\u{1}{}", obs.install_id, session_id);
let applied = match st.intervention_sessions.apply(
&session_key,
body,
&directive.params,
obs.billing,
obs.churn
.as_ref()
.is_some_and(|finding| finding.churn_layer != super::churn::ChurnLayer::Messages),
now,
) {
Ok(applied) => applied,
Err(reason) => {
if let Some(candidate) = selected.as_mut().and_then(|ctx| ctx.actual.as_mut()) {
candidate.reason = reason.into();
}
let mut decision = initial;
decision.optimize = selected;
return (
I10Outcome {
bytes: None,
applied: None,
blocked: false,
decision,
policy_state: None,
policy_layout: zone.state_layout,
policy_session_id: session_id.to_string(),
second_optimize_degraded: false,
},
initial_state,
);
}
};
let serialized = match serde_json::to_vec(&applied.body) {
Ok(bytes) => Bytes::from(bytes),
Err(_) => {
return (
I10Outcome {
bytes: None,
applied: None,
blocked: false,
decision: initial,
policy_state: None,
policy_layout: zone.state_layout,
policy_session_id: session_id.to_string(),
second_optimize_degraded: false,
},
initial_state,
)
}
};
let final_event = event_for(applied.body.clone());
let (mut final_decision, final_state) =
crate::zone_eval::evaluate(zone, &final_event, prior, now_ms);
let same = final_decision
.optimize
.as_ref()
.and_then(|ctx| ctx.actual.as_ref())
.and_then(|candidate| candidate.artifact_id.as_deref())
== actual.artifact_id.as_deref();
let second_optimize_degraded = final_decision.verdict == Verdict::Optimize && !same;
if second_optimize_degraded {
final_decision.verdict = Verdict::Ask;
}
final_decision.optimize = selected;
let blocked = final_decision.verdict == Verdict::Block;
let deliver = !blocked && final_decision.verdict != Verdict::Ask;
let changed = applied.changed;
(
I10Outcome {
bytes: (deliver && changed).then_some(serialized),
applied: deliver.then_some(applied),
decision: final_decision,
blocked,
policy_state: deliver.then_some(final_state.clone()),
policy_layout: zone.state_layout,
policy_session_id: session_id.to_string(),
second_optimize_degraded,
},
final_state,
)
})();
Some(outcome)
}
#[allow(clippy::type_complexity)]
fn apply_transform(
st: &ModelRelayState,
obs: &Observation,
body: &serde_json::Value,
) -> Option<(Option<Bytes>, super::transforms::TransformDecision)> {
if st.inject_mutate_panic {
panic!("injected mutate panic (D-28)");
}
let resident = st.resident_request_rules();
let bundle = resident
.as_ref()
.and_then(|guard| guard.as_ref().as_ref())?;
let authored = bundle.request_rules.as_slice();
if authored.is_empty() {
return None;
}
let enforcing = bundle.enforcement_enabled;
let session_key = obs.session.session_id.clone().unwrap_or_default();
let shape = st.prefix_shape.observe(&obs.install_id, &session_key, body);
let evaluated = super::transforms::evaluate_acting(body, authored, &shape, enforcing)?;
let Some(rewritten) = evaluated.rewritten else {
return Some((None, evaluated.decision));
};
match serde_json::to_vec(&rewritten) {
Ok(encoded) => Some((Some(Bytes::from(encoded)), evaluated.decision)),
Err(_) => None,
}
}
fn retention_store(f: &super::churn::ChurnFinding, occurred_at: &str) {
super::retention::store(&super::retention::FindingRecord {
finding_id: f.finding_id.clone(),
captured_at: occurred_at.to_string(),
churn_layer: f.churn_layer.as_str().to_string(),
churn_class: f.churn_class.as_str().to_string(),
divergence_offset: f.divergence_offset,
churn_byte_len: f.churn_byte_len,
churn_block_index: f.churn_block_index,
block: f.block.clone(),
});
}
fn observe_request_metadata_only(
st: &ModelRelayState,
wire: WireFormat,
headers: &HeaderMap,
path: &str,
) -> Observation {
let install_id = request_install_id(st, headers);
let model = capture::model_from_path(wire, path);
Observation {
measured: true,
model_known: model.is_some(),
model,
event_id: uuid::Uuid::now_v7().to_string(),
occurred_at: crate::envelope::current_timestamp(),
billing: detect_billing(headers),
session: session::resolve_session_scoped(
&st.registry,
&install_id,
&session::RequestSignals::default(),
endpoint_agent(st),
),
install_id,
wire_format: wire,
attributable_agent: attributable_agent(st, wire),
..Observation::none()
}
}
const TEXT_SCAN_BYTES: usize = 4096;
const PROMPT_SCAN_TURNS: usize = 4;
fn request_signals(fmt: WireFormat, body: &serde_json::Value) -> session::RequestSignals {
let mut out = session::RequestSignals {
declared_session_id: wire_format::declared_session_id(fmt, body),
..Default::default()
};
let Some(view) = wire_format::conversation_view(fmt, body) else {
return out;
};
let tail_from = view.user_turns.len().saturating_sub(PROMPT_SCAN_TURNS);
for (i, msg) in view.user_turns.iter().enumerate() {
if i != 0 && i < tail_from {
continue;
}
for text in [
first_text_block(
msg,
view.content_field,
view.text_block_type,
TEXT_SCAN_BYTES,
),
message_text(
msg,
view.content_field,
view.text_block_type,
TEXT_SCAN_BYTES,
),
] {
if text.trim().is_empty() {
continue;
}
let h = session::text_hash(&text);
if !out.prompt_hashes.contains(&h) {
out.prompt_hashes.push(h);
}
}
}
out.tool_use_ids = view.tool_result_ids;
out
}
fn is_text_block(b: &serde_json::Value, block_type: Option<&str>) -> bool {
match block_type {
Some(t) => b.get("type").and_then(|x| x.as_str()) == Some(t),
None => b.get("text").is_some(),
}
}
fn first_text_block(
msg: &serde_json::Value,
content_field: &str,
block_type: Option<&str>,
limit: usize,
) -> String {
let mut out = String::new();
match msg.get(content_field) {
Some(serde_json::Value::String(s)) => push_bounded(&mut out, s, limit),
Some(serde_json::Value::Array(blocks)) => {
for b in blocks {
if !is_text_block(b, block_type) {
continue;
}
if let Some(t) = b.get("text").and_then(|v| v.as_str()) {
push_bounded(&mut out, t, limit);
}
break;
}
}
_ => {}
}
out
}
fn message_text(
msg: &serde_json::Value,
content_field: &str,
block_type: Option<&str>,
limit: usize,
) -> String {
let mut out = String::new();
match msg.get(content_field) {
Some(serde_json::Value::String(s)) => push_bounded(&mut out, s, limit),
Some(serde_json::Value::Array(blocks)) => {
for b in blocks {
if is_text_block(b, block_type) {
if let Some(t) = b.get("text").and_then(|v| v.as_str()) {
push_bounded(&mut out, t, limit);
}
}
if out.len() >= limit {
break;
}
}
}
_ => {}
}
out
}
fn push_bounded(out: &mut String, s: &str, limit: usize) {
if out.len() >= limit {
return;
}
let room = limit - out.len();
if s.len() <= room {
out.push_str(s);
} else {
let mut end = room;
while end > 0 && !s.is_char_boundary(end) {
end -= 1;
}
out.push_str(&s[..end]);
}
}
fn attributable_agent(st: &ModelRelayState, wire: WireFormat) -> Option<&'static str> {
endpoint_agent(st).or_else(|| st.wiring.sole_wired_agent_for(wire))
}
fn request_install_id(st: &ModelRelayState, headers: &HeaderMap) -> String {
headers
.get(INSTALL_ID_HEADER)
.and_then(|v| v.to_str().ok())
.filter(|v| !v.is_empty())
.map(str::to_string)
.or_else(|| st.install_id.clone())
.unwrap_or_default()
}
fn endpoint_agent(st: &ModelRelayState) -> Option<&'static str> {
st.endpoint.as_ref().map(|e| e.agent())
}
fn is_preflight(headers: &HeaderMap) -> bool {
headers.contains_key(PREFLIGHT_HEADER)
}
fn opaque_measure_ctx(
st: &Arc<ModelRelayState>,
wire: WireFormat,
headers: &HeaderMap,
path: &str,
) -> Option<MeasureCtx> {
if is_preflight(headers) {
return None;
}
st.cloud_tx.as_ref().map(|_| MeasureCtx {
obs: observe_request_metadata_only(st, wire, headers, path),
tokenizer: st.tokenizer,
cloud_tx: st.cloud_tx.clone(),
privacy: st.privacy.clone(),
wire_format_unknown: true,
})
}
pub async fn proxy_any(State(st): State<Arc<ModelRelayState>>, req: Request) -> Response {
let (parts, body) = req.into_parts();
let endpoint = st.endpoint.as_ref();
let wire = WireFormat::resolve_with_family(&parts, endpoint.and_then(|e| e.family()));
if let Some(endpoint) = endpoint {
if !is_preflight(&parts.headers) {
endpoint.note_request();
}
}
if !wire.is_captured() {
return stream_through_opaque(&st, parts, body, None).await;
}
let len = content_length(&parts.headers);
if len.is_none_or(|n| n > MAX_MATERIALIZE_BYTES) {
let ctx = opaque_measure_ctx(&st, wire, &parts.headers, parts.uri.path());
return stream_through_opaque(&st, parts, body, ctx).await;
}
let permit = match st.inflight.clone().try_acquire_owned() {
Ok(p) => p,
Err(_) => {
let ctx = opaque_measure_ctx(&st, wire, &parts.headers, parts.uri.path());
return stream_through_opaque(&st, parts, body, ctx).await;
}
};
let bytes = match axum::body::to_bytes(body, MAX_MATERIALIZE_BYTES).await {
Ok(b) => b,
Err(_) => {
record_pass_through_failure("body_read");
return synth_502(&st);
}
};
let (mut observation, parsed) = std::panic::catch_unwind(std::panic::AssertUnwindSafe(|| {
observe_request(&st, wire, &parts.headers, parts.uri.path(), &bytes)
}))
.unwrap_or_else(|_| {
record_pass_through_failure("observe_panic");
(Observation::none(), None)
});
if is_preflight(&parts.headers) {
observation.measured = false;
}
let original_bytes = bytes.clone();
let mut i10_delivery = None;
let bytes = if !wire.transforms_apply() {
bytes } else {
match parsed {
None => bytes,
Some(body) => {
let i10 = std::panic::catch_unwind(std::panic::AssertUnwindSafe(|| {
apply_i10(&st, &observation, &body)
}));
let attempt = match i10 {
Ok(None) => I10Attempt::Absent,
Ok(Some(outcome)) if outcome.blocked || outcome.decision.optimize.is_some() => {
I10Attempt::Handled(Box::new(outcome))
}
Ok(Some(_)) => I10Attempt::Absent,
Err(_) => {
record_pass_through_failure("i10_mutate_panic");
I10Attempt::Failed
}
};
match attempt {
I10Attempt::Handled(outcome) => {
let outcome = *outcome;
if outcome.blocked {
observation.optimize =
outcome
.decision
.optimize
.map(|context| emit::OptimizeEvidence {
context,
enforced: true,
result: "blocked",
second_optimize_degraded: false,
});
emit::build_and_emit(
&observation,
&Usage::default(),
CostBasis::ProviderReported,
None,
true,
&st.privacy,
st.cloud_tx.as_ref(),
);
return synth_blocked();
}
let enforced = outcome.bytes.is_some();
observation.optimize =
outcome
.decision
.optimize
.map(|context| emit::OptimizeEvidence {
context,
enforced,
result: if enforced { "rewritten" } else { "flagged" },
second_optimize_degraded: outcome.second_optimize_degraded,
});
i10_delivery = Some(I10Delivery {
applied: outcome.applied,
policy_state: outcome.policy_state,
policy_layout: outcome.policy_layout,
policy_session_id: outcome.policy_session_id,
});
outcome.bytes.unwrap_or(bytes)
}
I10Attempt::Failed => bytes,
I10Attempt::Absent => {
let mutated =
std::panic::catch_unwind(std::panic::AssertUnwindSafe(|| {
apply_transform(&st, &observation, &body)
}))
.unwrap_or_else(|_| {
record_pass_through_failure("mutate_panic");
None
});
match mutated {
Some((rewritten, decision)) => {
observation.transform = Some(decision);
rewritten.unwrap_or(bytes)
}
None => bytes,
}
}
}
}
}
};
forward_streaming(
&st,
parts,
bytes,
permit,
observation,
i10_delivery.map(|delivery| (original_bytes, delivery)),
)
.await
}
async fn stream_through_opaque(
st: &Arc<ModelRelayState>,
parts: Parts,
body: Body,
measure_ctx: Option<MeasureCtx>,
) -> Response {
let url = match upstream_url(st, &parts) {
Some(u) => u,
None => {
emit_terminal_error(measure_ctx);
return synth_502(st);
}
};
let Some(client) = st.client.current() else {
emit_terminal_error(measure_ctx);
return synth_502(st);
};
let req_body = reqwest::Body::wrap_stream(body.into_data_stream());
let send = client
.request(parts.method.clone(), url.clone())
.headers(forward_headers(&parts.headers))
.body(req_body)
.send();
let upstream = match tokio::time::timeout(st.header_timeout, send).await {
Ok(res) => res,
Err(_elapsed) => {
emit_terminal_error(measure_ctx);
return synth_502(st);
}
};
relay(st, &url, upstream, None, measure_ctx)
}
async fn forward_streaming(
st: &Arc<ModelRelayState>,
parts: Parts,
bytes: Bytes,
permit: OwnedSemaphorePermit,
observation: Observation,
i10_fallback: Option<(Bytes, I10Delivery)>,
) -> Response {
let undecoded = !observation.wire_format.has_decoder();
let mut measure_ctx = if observation.measured && st.cloud_tx.is_some() {
Some(MeasureCtx {
obs: observation,
tokenizer: st.tokenizer,
cloud_tx: st.cloud_tx.clone(),
privacy: st.privacy.clone(),
wire_format_unknown: undecoded,
})
} else {
None
};
let url = match upstream_url(st, &parts) {
Some(u) => u,
None => {
clear_i10_enforcement(&mut measure_ctx);
emit_terminal_error(measure_ctx);
return synth_502(st);
}
};
let Some(client) = st.client.current() else {
clear_i10_enforcement(&mut measure_ctx);
emit_terminal_error(measure_ctx);
return synth_502(st);
};
let send = client
.request(parts.method.clone(), url.clone())
.headers(forward_headers(&parts.headers))
.body(bytes) .send();
let mut upstream = match tokio::time::timeout(st.header_timeout, send).await {
Ok(res) => res,
Err(_elapsed) => {
clear_i10_enforcement(&mut measure_ctx);
emit_terminal_error(measure_ctx);
return synth_502(st);
}
};
if upstream
.as_ref()
.is_ok_and(|response| response.status() == StatusCode::BAD_REQUEST)
{
if let Some((original, delivery)) = i10_fallback {
if let Some(applied) = delivery.applied.as_ref() {
st.intervention_sessions.rejected(applied);
}
clear_i10_enforcement(&mut measure_ctx);
let retry = client
.request(parts.method.clone(), url.clone())
.headers(forward_headers(&parts.headers))
.body(original)
.send();
upstream = match tokio::time::timeout(st.header_timeout, retry).await {
Ok(result) => result,
Err(_) => {
clear_i10_enforcement(&mut measure_ctx);
emit_terminal_error(measure_ctx);
return synth_502(st);
}
};
}
} else if upstream
.as_ref()
.is_ok_and(|response| response.status().is_success())
{
if let Some((_, delivery)) = i10_fallback {
if let Some(applied) = delivery.applied.as_ref() {
st.intervention_sessions
.accepted(applied, std::time::Instant::now());
}
if let Some(policy_state) = delivery.policy_state {
st.policy_sessions.commit(
delivery.policy_session_id,
delivery.policy_layout,
policy_state,
std::time::Instant::now(),
);
}
}
} else {
clear_i10_enforcement(&mut measure_ctx);
}
relay(st, &url, upstream, Some(permit), measure_ctx)
}
fn clear_i10_enforcement(measure_ctx: &mut Option<MeasureCtx>) {
if let Some(evidence) = measure_ctx
.as_mut()
.and_then(|ctx| ctx.obs.optimize.as_mut())
{
evidence.enforced = false;
evidence.result = "flagged";
}
}
fn relay(
st: &Arc<ModelRelayState>,
base: &reqwest::Url,
upstream: Result<reqwest::Response, reqwest::Error>,
permit: Option<OwnedSemaphorePermit>,
measure_ctx: Option<MeasureCtx>,
) -> Response {
let resp = match upstream {
Ok(r) => {
st.egress.record_ok(base.as_str());
r
}
Err(e) => {
st.egress.record_failure(base.as_str(), &e);
emit_terminal_error(measure_ctx);
return synth_502(st);
}
};
let status = resp.status();
let mut headers = resp.headers().clone();
strip_hop_by_hop_headers(&mut headers);
let measure = measure_ctx.map(|ctx| ctx.into_measure(status.is_success()));
let guarded = GuardedBody {
inner: Box::pin(resp.bytes_stream()),
_permit: permit,
measure,
};
let mut out = Response::new(Body::from_stream(guarded));
*out.status_mut() = status;
*out.headers_mut() = headers;
out
}
fn emit_terminal_error(measure_ctx: Option<MeasureCtx>) {
if let Some(ctx) = measure_ctx {
let mut m = ctx.into_measure(false);
m.finalize();
}
}
fn upstream_url(st: &Arc<ModelRelayState>, parts: &Parts) -> Option<reqwest::Url> {
let pq = parts
.uri
.path_and_query()
.map(|x| x.as_str())
.unwrap_or("/");
let base = st.upstream_for_request(
WireFormat::resolve_upstream(parts),
AuthMode::resolve(&parts.headers),
);
join_upstream(&base, pq)
}
fn join_upstream(base: &reqwest::Url, pq: &str) -> Option<reqwest::Url> {
let prefix = base.path().trim_end_matches('/');
if prefix.is_empty() {
return base.join(pq).ok();
}
let last = prefix.rsplit('/').next().filter(|seg| !seg.is_empty());
let strip_segment = |segment: &str| {
pq.strip_prefix('/')
.and_then(|p| p.strip_prefix(segment))
.filter(|rest| rest.is_empty() || rest.starts_with('/') || rest.starts_with('?'))
};
let suffix = strip_segment("v1")
.or_else(|| last.and_then(strip_segment))
.unwrap_or(pq);
base.join(&format!("{prefix}{suffix}")).ok()
}
fn content_length(headers: &HeaderMap) -> Option<usize> {
headers.get(CONTENT_LENGTH)?.to_str().ok()?.parse().ok()
}
fn forward_headers(src: &HeaderMap) -> HeaderMap {
let mut h = src.clone();
h.remove(HOST); h.remove(ACCEPT_ENCODING);
h.remove(PREFLIGHT_HEADER);
strip_hop_by_hop_headers(&mut h);
h
}
fn strip_hop_by_hop_headers(h: &mut HeaderMap) {
let connection_named: Vec<String> = h
.get_all(CONNECTION)
.iter()
.filter_map(|v| v.to_str().ok())
.flat_map(|v| v.split(','))
.map(|t| t.trim().to_ascii_lowercase())
.filter(|t| !t.is_empty())
.collect();
for name in connection_named {
h.remove(name.as_str());
}
h.remove(CONNECTION);
h.remove("keep-alive");
h.remove(TRANSFER_ENCODING);
h.remove(TE);
h.remove(TRAILER);
h.remove(UPGRADE);
h.remove(PROXY_AUTHENTICATE);
h.remove(PROXY_AUTHORIZATION);
h.remove(CONTENT_LENGTH);
}
fn synth_502(st: &ModelRelayState) -> Response {
match &st.endpoint {
Some(endpoint) => endpoint.note_upstream_failure(),
None => {
UPSTREAM_FAILURES.fetch_add(1, Ordering::Relaxed);
}
}
let mut out = Response::new(Body::empty());
*out.status_mut() = StatusCode::BAD_GATEWAY;
out.headers_mut().insert(
"x-openlatch-upstream",
HeaderValue::from_static("unreachable"),
);
out
}
fn synth_blocked() -> Response {
let mut out = Response::new(Body::empty());
*out.status_mut() = StatusCode::FORBIDDEN;
out.headers_mut()
.insert("x-openlatch-verdict", HeaderValue::from_static("block"));
out
}
pub async fn model_relay_status(State(st): State<Arc<ModelRelayState>>) -> Json<serde_json::Value> {
let verdicts = st.wiring.verdicts();
Json(serde_json::json!({
"status": "up",
"port": st.port,
"upstream": WireFormat::ALL
.iter()
.map(|f| (f.as_str().to_string(), serde_json::json!(st.upstream_for(*f).as_str())))
.collect::<serde_json::Map<_, _>>(),
"upstream_chatgpt": st.chatgpt_upstream().map(reqwest::Url::as_str),
"inflight_available": st.inflight.available_permits(),
"uptime_secs": st.started_at.elapsed().as_secs(),
"pass_through_failures": pass_through_failures(),
"upstream_failures": upstream_failures(),
"selector_wins": session::selector_wins()
.into_iter()
.map(|(k, v)| (k.to_string(), serde_json::json!(v)))
.collect::<serde_json::Map<_, _>>(),
"wired": st.wiring
.wired_agents()
.into_iter()
.map(|(a, w)| (a.to_string(), serde_json::json!(w)))
.collect::<serde_json::Map<_, _>>(),
"preflight": verdicts
.iter()
.map(|(a, v)| (a.to_string(), serde_json::json!(v.label())))
.collect::<serde_json::Map<_, _>>(),
"preflight_error": verdicts
.iter()
.map(|(a, v)| (a.to_string(), serde_json::json!(v.error())))
.collect::<serde_json::Map<_, _>>(),
"endpoint_verdicts": st.wiring
.endpoint_verdicts()
.into_iter()
.map(|(k, v)| (k, serde_json::json!({"code": v.code, "detail": v.detail})))
.collect::<serde_json::Map<_, _>>(),
"endpoint": st.endpoint.as_ref().map(|e| e.to_json()),
"endpoints": match (&st.endpoint, &st.endpoints) {
(None, Some(listeners)) => serde_json::json!(listeners
.snapshot()
.iter()
.map(|e| e.to_json())
.collect::<Vec<_>>()),
(None, None) => serde_json::json!([]),
(Some(_), _) => serde_json::Value::Null,
},
}))
}
struct MeasureCtx {
obs: Observation,
tokenizer: Estimator,
cloud_tx: Option<Sender<CloudEvent>>,
privacy: PrivacyFilter,
wire_format_unknown: bool,
}
impl MeasureCtx {
fn into_measure(self, status_ok: bool) -> Measure {
Measure {
obs: self.obs,
acc: UsageAccumulator::default(),
status_ok,
emitted: false,
wire_format_unknown: self.wire_format_unknown,
tokenizer: self.tokenizer,
cloud_tx: self.cloud_tx,
privacy: self.privacy,
}
}
}
struct Measure {
obs: Observation,
acc: UsageAccumulator,
status_ok: bool,
emitted: bool,
wire_format_unknown: bool,
tokenizer: Estimator,
cloud_tx: Option<Sender<CloudEvent>>,
privacy: PrivacyFilter,
}
impl Measure {
fn finalize(&mut self) {
if self.emitted {
return;
}
self.emitted = true;
let model = self.obs.model.clone().unwrap_or_default();
let (usage, basis, gap) = if !self.status_ok {
(
Usage::default(),
CostBasis::ProviderReported,
Some(CaptureGap::ProviderError),
)
} else if self.acc.is_terminal() {
let gap = if self.acc.provider_arithmetic_bad() {
Some(CaptureGap::ProviderError)
} else {
base_gap(&self.obs)
};
(self.acc.usage(), CostBasis::ProviderReported, gap)
} else {
let est = self.tokenizer.estimate(&model, self.obs.request_body_len);
let usage = Usage {
input_tokens: est.input_tokens,
..Usage::default()
};
let gap = base_gap(&self.obs).or(Some(CaptureGap::StreamInterrupted));
(usage, CostBasis::TokenizerEstimated, gap)
};
let gap = if self.wire_format_unknown && self.status_ok {
Some(CaptureGap::UnknownWireFormat)
} else {
gap
};
let cache_preserved = capture::infer_cache_preserved(&usage);
emit::build_and_emit(
&self.obs,
&usage,
basis,
gap,
cache_preserved,
&self.privacy,
self.cloud_tx.as_ref(),
);
}
}
fn base_gap(obs: &Observation) -> Option<CaptureGap> {
if obs.model.is_some() && !obs.model_known {
Some(CaptureGap::UnknownModel)
} else {
None
}
}
struct GuardedBody {
inner: Pin<Box<dyn Stream<Item = reqwest::Result<Bytes>> + Send>>,
_permit: Option<OwnedSemaphorePermit>,
measure: Option<Measure>,
}
impl Stream for GuardedBody {
type Item = reqwest::Result<Bytes>;
fn poll_next(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
let this = self.get_mut();
let polled = this.inner.as_mut().poll_next(cx);
match &polled {
Poll::Ready(Some(Ok(chunk))) => {
let scan_panicked = if let Some(m) = this.measure.as_mut() {
let fmt = m.obs.wire_format;
let acc = &mut m.acc;
std::panic::catch_unwind(std::panic::AssertUnwindSafe(|| {
#[cfg(test)]
if INJECT_SCAN_PANIC.load(Ordering::Relaxed) {
panic!("injected usage-scan panic (FIX 3 test)");
}
acc.scan_chunk(fmt, chunk)
}))
.is_err()
} else {
false
};
if scan_panicked {
record_pass_through_failure("usage_scan_panic");
this.measure = None;
}
}
Poll::Ready(None) | Poll::Ready(Some(Err(_))) => {
if let Some(m) = this.measure.as_mut() {
m.finalize();
}
}
Poll::Pending => {}
}
polled
}
}
impl Drop for GuardedBody {
fn drop(&mut self) {
if let Some(m) = self.measure.as_mut() {
m.finalize();
}
}
}
#[cfg(test)]
mod tests {
use super::*;
use crate::model_relay::{mock, serve_ephemeral, ModelRelayState};
use futures_util::StreamExt;
use std::time::{Duration, Instant};
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, &[]))
}
fn joined(base: &str, pq: &str) -> String {
let url = reqwest::Url::parse(base).expect("base parses");
join_upstream(&url, pq).expect("join succeeds").to_string()
}
#[test]
fn a_bare_origin_joins_exactly_as_before() {
assert_eq!(
joined("https://api.anthropic.com", "/v1/messages"),
"https://api.anthropic.com/v1/messages"
);
assert_eq!(
joined("https://api.openai.com", "/v1/responses"),
"https://api.openai.com/v1/responses"
);
assert_eq!(
joined("https://api.anthropic.com", "/v1/models?beta=true"),
"https://api.anthropic.com/v1/models?beta=true"
);
assert_eq!(
joined("https://api.anthropic.com/", "/v1/messages"),
"https://api.anthropic.com/v1/messages"
);
}
#[test]
fn a_codex_install_forwards_each_route_to_the_right_origin() {
let st = Arc::new(
ModelRelayState::new(crate::model_relay::default_upstream(), 0, 8, &[])
.with_upstream_map(std::collections::BTreeMap::new()),
);
use axum::http::Method;
let sent = |method: Method, uri: &str, headers: &[(&str, &str)]| {
let mut req = axum::http::Request::builder().method(method).uri(uri);
for (k, v) in headers {
req = req.header(*k, *v);
}
let (parts, ()) = req.body(()).expect("build request").into_parts();
upstream_url(&st, &parts)
.expect("an upstream resolves")
.to_string()
};
const ACCOUNT: (&str, &str) = ("chatgpt-account-id", "ef1a0c98");
const ORIGINATOR: (&str, &str) = ("originator", "codex_exec");
const PLATFORM_KEY: (&str, &str) = ("authorization", "Bearer sk-proj-abc");
assert_eq!(
sent(Method::POST, "/v1/responses", &[ACCOUNT, ORIGINATOR]),
"https://chatgpt.com/backend-api/codex/responses"
);
assert_eq!(
sent(Method::POST, "/v1/responses", &[PLATFORM_KEY, ORIGINATOR]),
"https://api.openai.com/v1/responses"
);
assert_eq!(
sent(
Method::GET,
"/v1/models?client_version=0.150.1",
&[ACCOUNT, ORIGINATOR]
),
"https://chatgpt.com/backend-api/codex/models?client_version=0.150.1"
);
assert_eq!(
sent(
Method::GET,
"/v1/models?client_version=0.150.1",
&[PLATFORM_KEY, ORIGINATOR]
),
"https://api.openai.com/v1/models?client_version=0.150.1"
);
assert_eq!(
sent(Method::POST, "/v1/messages", &[("x-api-key", "sk-ant-abc")]),
"https://api.anthropic.com/v1/messages"
);
assert_eq!(
sent(
Method::GET,
"/v1/models",
&[("authorization", "Bearer sk-ant-oat01-abc")]
),
"https://api.anthropic.com/v1/models"
);
}
#[test]
fn a_based_upstream_keeps_its_prefix() {
assert_eq!(
joined(wire_format::CHATGPT_BASE, "/v1/responses"),
"https://chatgpt.com/backend-api/codex/responses"
);
assert_eq!(
joined(wire_format::CHATGPT_BASE, "/v1/models?client_version=0.1"),
"https://chatgpt.com/backend-api/codex/models?client_version=0.1"
);
let base = reqwest::Url::parse(wire_format::CHATGPT_BASE).unwrap();
assert_eq!(
base.join("/v1/responses").unwrap().to_string(),
"https://chatgpt.com/v1/responses",
"plain join drops the prefix — that is the bug"
);
}
#[test]
fn a_gateway_prefix_absorbs_the_version_segment() {
assert_eq!(
joined("https://gw.example/proxy/v1", "/v1/messages"),
"https://gw.example/proxy/v1/messages"
);
assert_eq!(
joined("https://gw.example/proxy/v1/", "/v1/messages"),
"https://gw.example/proxy/v1/messages"
);
}
#[test]
fn join_upstream_does_not_double_a_version_segment_the_base_ends_with() {
assert_eq!(
joined(
"https://gw.example/google/v1beta",
"/v1beta/models/gemini-2.5-pro:streamGenerateContent?alt=sse"
),
"https://gw.example/google/v1beta/models/gemini-2.5-pro:streamGenerateContent?alt=sse"
);
assert_eq!(
joined("https://gw.example/google/v1beta", "/v1beta2/models"),
"https://gw.example/google/v1beta/v1beta2/models"
);
assert_eq!(
joined("http://127.0.0.1:11434/", "/api/chat"),
"http://127.0.0.1:11434/api/chat"
);
assert_eq!(
joined(
"https://generativelanguage.googleapis.com/",
"/v1beta/models/m:streamGenerateContent?alt=sse"
),
"https://generativelanguage.googleapis.com/v1beta/models/m:streamGenerateContent?alt=sse"
);
}
#[test]
fn only_a_whole_version_segment_is_stripped() {
assert_eq!(
joined(wire_format::CHATGPT_BASE, "/v1beta/models"),
"https://chatgpt.com/backend-api/codex/v1beta/models"
);
assert_eq!(
joined(wire_format::CHATGPT_BASE, "/v2/responses"),
"https://chatgpt.com/backend-api/codex/v2/responses"
);
assert_eq!(
joined(wire_format::CHATGPT_BASE, "/v1"),
"https://chatgpt.com/backend-api/codex"
);
}
fn reporting_state() -> (Arc<ModelRelayState>, crate::egress::EgressState) {
let cfg = crate::egress::EgressConfig::direct();
let health = crate::egress::EgressState::new(&cfg);
let state = ModelRelayState::new(
reqwest::Url::parse("https://api.anthropic.test/").expect("url"),
0,
8,
&[],
)
.with_egress_reporter(crate::egress::EgressReporter::recording(
&cfg,
health.clone(),
));
(Arc::new(state), health)
}
#[tokio::test(flavor = "multi_thread")]
async fn a_forwarded_response_records_egress_ok() {
let (st, health) = reporting_state();
health.record_failure(crate::error::ERR_PROXY_UNREACHABLE, "an earlier outage");
health.record_failure(crate::error::ERR_PROXY_UNREACHABLE, "an earlier outage");
let upstream = mock::spawn_capture_200().await;
let response = crate::egress::client()
.get(format!("http://127.0.0.1:{}/", upstream.port))
.send()
.await
.expect("the mock answers 200");
let base = st.upstream_for(WireFormat::AnthropicMessages).clone();
let _ = relay(&st, &base, Ok(response), None, None);
assert_eq!(
health.consecutive_failures(),
0,
"a forwarded response is proof the route works"
);
assert!(health.last_ok_at().is_some());
}
#[tokio::test(flavor = "multi_thread")]
async fn the_connect_level_error_is_the_one_that_records_a_failure() {
let (st, health) = reporting_state();
let dead = mock::closed_port().await;
let err = crate::egress::client()
.get(format!("http://127.0.0.1:{dead}/"))
.send()
.await
.expect_err("a closed port cannot answer");
let base = st.upstream_for(WireFormat::AnthropicMessages).clone();
let _ = relay(&st, &base, Err(err), None, None);
assert_eq!(health.consecutive_failures(), 1);
let recorded = health.last_error().expect("a code and a message");
assert_eq!(
recorded.code,
crate::error::ERR_EGRESS_UNREACHABLE,
"a direct route names the platform side, not a proxy"
);
assert!(
!recorded.message.is_empty(),
"a code without a message is not a diagnosis"
);
}
fn measured_state(
upstream_port: u16,
) -> (
Arc<ModelRelayState>,
tokio::sync::mpsc::Receiver<CloudEvent>,
) {
let (tx, rx) = tokio::sync::mpsc::channel::<CloudEvent>(8);
let base = reqwest::Url::parse(&format!("http://127.0.0.1:{upstream_port}")).unwrap();
let registry = Arc::new(session::SessionRegistry::default());
registry.upsert("agt_1", "agt_1", "claude-code", "sess_a");
let st =
Arc::new(ModelRelayState::new(base, 0, 8, &[]).with_measurement(registry, Some(tx)));
(st, rx)
}
#[tokio::test(flavor = "multi_thread")]
async fn unknown_route_forwards_byte_identically_and_emits_nothing() {
let up = mock::spawn_capture_200().await;
let (st, mut rx) = measured_state(up.port);
let port = serve_ephemeral(st).await;
let body = br#"{"anything":"at all"}"#.to_vec();
let resp = crate::egress::client()
.post(format!("http://127.0.0.1:{port}/v1/unknown"))
.header("content-type", "application/json")
.header(INSTALL_ID_HEADER, "agt_1")
.body(body.clone())
.send()
.await
.expect("the forward answers");
assert!(resp.status().is_success());
let _ = resp.bytes().await.unwrap();
assert_eq!(
up.received_body.lock().unwrap().clone(),
Some(body),
"an uncaptured route forwards the bytes verbatim"
);
assert!(
tokio::time::timeout(Duration::from_millis(400), rx.recv())
.await
.ok()
.flatten()
.is_none(),
"an uncaptured route must emit NO economics event"
);
}
#[tokio::test(flavor = "multi_thread")]
async fn responses_decoded_turn_is_not_unknown_wire_format() {
let _serialised = SCAN_PANIC_LOCK.lock().await;
let up = mock::spawn_capture_responses_sse(14_342, 2_688, 0, 916).await;
let (st, mut rx) = measured_state(up.port);
let port = serve_ephemeral(st).await;
let body = br#"{"model":"gpt-5-codex","input":"hello"}"#.to_vec();
let resp = crate::egress::client()
.post(format!("http://127.0.0.1:{port}/v1/responses"))
.header("content-type", "application/json")
.header(INSTALL_ID_HEADER, "agt_1")
.body(body)
.send()
.await
.expect("the forward answers");
assert!(resp.status().is_success());
let _ = resp.bytes().await.unwrap();
let env = tokio::time::timeout(Duration::from_secs(3), rx.recv())
.await
.ok()
.flatten()
.expect("a captured route emits one economics event")
.envelope;
let data = &env["data"];
assert_ne!(
data["ai.openlatch.capture.gap"], "unknown_wire_format",
"the decoder ran — labelling a real measurement unmeasured is the D-14 defect"
);
assert!(
data["ai.openlatch.capture.gap"].is_null(),
"a fully provider-reported turn on a known model has no gap to report, got {}",
data["ai.openlatch.capture.gap"]
);
assert_eq!(data["ai.openlatch.cost.basis"], "provider_reported");
assert_eq!(data["gen_ai.provider.name"], "openai");
assert_eq!(data["gen_ai.usage.output_tokens"], 916);
assert_eq!(data["gen_ai.usage.cache_read.input_tokens"], 2_688);
assert_eq!(
data["gen_ai.usage.input_tokens"],
14_342 - 2_688,
"fresh input: the three-term subtraction, through the real proxy"
);
assert_eq!(data["ai.openlatch.cache.ephemeral_5m_input_tokens"], 0);
assert_eq!(data["ai.openlatch.cache.ephemeral_1h_input_tokens"], 0);
}
#[tokio::test(flavor = "multi_thread")]
async fn responses_clamp_sets_provider_error() {
let _serialised = SCAN_PANIC_LOCK.lock().await;
let up = mock::spawn_capture_responses_sse(10, 8, 5, 4).await;
let (st, mut rx) = measured_state(up.port);
let port = serve_ephemeral(st).await;
let body = br#"{"model":"gpt-5-codex","input":"hello"}"#.to_vec();
let resp = crate::egress::client()
.post(format!("http://127.0.0.1:{port}/v1/responses"))
.header("content-type", "application/json")
.header(INSTALL_ID_HEADER, "agt_1")
.body(body)
.send()
.await
.expect("the forward answers");
assert!(resp.status().is_success());
let _ = resp.bytes().await.unwrap();
let env = tokio::time::timeout(Duration::from_secs(3), rx.recv())
.await
.ok()
.flatten()
.expect("a captured route emits one economics event")
.envelope;
let data = &env["data"];
assert_eq!(
data["ai.openlatch.capture.gap"], "provider_error",
"a 2xx whose usage arithmetic does not reconcile is a provider error"
);
assert_eq!(
data["gen_ai.usage.output_tokens"], 4,
"the output count is KEPT — it is still the provider's own number"
);
assert_eq!(
data["gen_ai.usage.input_tokens"], 0,
"clamped to zero, never wrapped"
);
assert_eq!(
data["ai.openlatch.cost.basis"], "provider_reported",
"the terminal frame did arrive; it is the input SPLIT that is wrong"
);
}
#[tokio::test(flavor = "multi_thread")]
async fn ollama_native_decoded_turn_is_not_unknown_wire_format() {
let _serialised = SCAN_PANIC_LOCK.lock().await;
let up = mock::spawn_capture_ollama_native_stream(36, 35, 3).await;
let (st, mut rx) = measured_state(up.port);
let port = serve_ephemeral(st).await;
let body = br#"{"model":"qwen2.5-coder:7b","messages":[{"role":"user","content":"hello"}],"stream":true}"#.to_vec();
let resp = crate::egress::client()
.post(format!("http://127.0.0.1:{port}/api/chat"))
.header("content-type", "application/json")
.header(INSTALL_ID_HEADER, "agt_1")
.body(body)
.send()
.await
.expect("the forward answers");
assert!(resp.status().is_success());
let _ = resp.bytes().await.unwrap();
let env = tokio::time::timeout(Duration::from_secs(3), rx.recv())
.await
.ok()
.flatten()
.expect("a captured route emits one economics event")
.envelope;
let data = &env["data"];
assert_ne!(
data["ai.openlatch.capture.gap"], "unknown_wire_format",
"the decoder ran — labelling a real measurement unmeasured is the D-14 defect"
);
assert_eq!(data["ai.openlatch.cost.basis"], "provider_reported");
assert_eq!(data["gen_ai.provider.name"], "ollama");
assert_eq!(data["gen_ai.usage.input_tokens"], 36);
assert_eq!(data["gen_ai.usage.cache_read.input_tokens"], 35);
assert_eq!(data["gen_ai.usage.output_tokens"], 3);
}
#[tokio::test(flavor = "multi_thread")]
async fn ollama_native_model_listing_emits_no_economics_event() {
let up = mock::spawn_capture_200().await;
let (st, mut rx) = measured_state(up.port);
let port = serve_ephemeral(st).await;
let resp = crate::egress::client()
.get(format!("http://127.0.0.1:{port}/api/tags"))
.header(INSTALL_ID_HEADER, "agt_1")
.send()
.await
.expect("the forward answers");
assert!(
resp.status().is_success(),
"it still reaches the model server"
);
let _ = resp.bytes().await.unwrap();
assert!(
tokio::time::timeout(Duration::from_millis(500), rx.recv())
.await
.is_err(),
"an uncaptured route emits no economics event"
);
}
#[tokio::test(flavor = "multi_thread")]
async fn chat_completions_split_across_chunk_boundary() {
let _serialised = SCAN_PANIC_LOCK.lock().await;
let up = mock::spawn_capture_chat_completions_sse(34, 8, 2).await;
let (st, mut rx) = measured_state(up.port);
let port = serve_ephemeral(st).await;
let body = br#"{"model":"qwen2.5-coder:14b","messages":[{"role":"user","content":"hello"}],"stream":true,"stream_options":{"include_usage":true}}"#.to_vec();
let resp = crate::egress::client()
.post(format!("http://127.0.0.1:{port}/v1/chat/completions"))
.header("content-type", "application/json")
.header(INSTALL_ID_HEADER, "agt_1")
.body(body)
.send()
.await
.expect("the forward answers");
assert!(resp.status().is_success());
let _ = resp.bytes().await.unwrap();
let env = tokio::time::timeout(Duration::from_secs(3), rx.recv())
.await
.ok()
.flatten()
.expect("a captured route emits one economics event")
.envelope;
let data = &env["data"];
assert_ne!(
data["ai.openlatch.capture.gap"], "unknown_wire_format",
"the decoder ran — labelling a real measurement unmeasured is the D-14 defect"
);
assert_eq!(data["ai.openlatch.cost.basis"], "provider_reported");
assert_eq!(data["gen_ai.provider.name"], "openai-compatible");
assert_eq!(data["gen_ai.usage.output_tokens"], 2);
assert_eq!(data["gen_ai.usage.cache_read.input_tokens"], 8);
assert_eq!(
data["gen_ai.usage.input_tokens"],
34 - 8,
"fresh input: the two-term subtraction, through the real proxy"
);
assert_eq!(data["gen_ai.usage.cache_creation.input_tokens"], 0);
assert_eq!(data["ai.openlatch.cache.ephemeral_5m_input_tokens"], 0);
assert_eq!(data["ai.openlatch.cache.ephemeral_1h_input_tokens"], 0);
}
#[tokio::test(flavor = "multi_thread")]
async fn chat_completions_stream_without_usage_degrades() {
let _serialised = SCAN_PANIC_LOCK.lock().await;
let up = mock::spawn_capture_chat_completions_sse_without_usage().await;
let (st, mut rx) = measured_state(up.port);
let port = serve_ephemeral(st).await;
let body = br#"{"model":"qwen2.5-coder:14b","messages":[{"role":"user","content":"hello"}],"stream":true}"#.to_vec();
let resp = crate::egress::client()
.post(format!("http://127.0.0.1:{port}/v1/chat/completions"))
.header("content-type", "application/json")
.header(INSTALL_ID_HEADER, "agt_1")
.body(body.clone())
.send()
.await
.expect("the forward answers");
assert!(resp.status().is_success());
let _ = resp.bytes().await.unwrap();
assert_eq!(
up.received_body.lock().unwrap().clone(),
Some(body),
"the request reached the upstream byte-identically — no `stream_options` was injected"
);
let env = tokio::time::timeout(Duration::from_secs(3), rx.recv())
.await
.ok()
.flatten()
.expect("the turn still emits one economics event — degraded, not missing")
.envelope;
let data = &env["data"];
assert_eq!(
data["ai.openlatch.cost.basis"], "tokenizer_estimated",
"no usage object ever arrived, so the count is a local estimate and says so"
);
assert_ne!(
data["ai.openlatch.capture.gap"], "unknown_wire_format",
"the ROUTE was recognised and the decoder ran; it is the provider's \
numbers that are absent, which is a different gap"
);
assert_eq!(
data["gen_ai.usage.output_tokens"], 0,
"output is never estimated — an invented output count is worse than none"
);
assert!(
data["gen_ai.usage.input_tokens"].as_u64().unwrap_or(0) > 0,
"the estimate is non-zero for a non-empty body, got {}",
data["gen_ai.usage.input_tokens"]
);
assert_eq!(data["gen_ai.provider.name"], "openai-compatible");
}
#[tokio::test(flavor = "multi_thread")]
async fn google_split_across_chunk_boundary() {
let _serialised = SCAN_PANIC_LOCK.lock().await;
let up = mock::spawn_capture_google_sse(1024, 256, 64, 32).await;
let (st, mut rx) = measured_state(up.port);
let port = serve_ephemeral(st).await;
let body = br#"{"contents":[{"role":"user","parts":[{"text":"hello"}]}]}"#.to_vec();
let resp = crate::egress::client()
.post(format!(
"http://127.0.0.1:{port}/v1beta/models/gemini-2.5-pro:streamGenerateContent\
?alt=sse&key=AIzaSyFAKEKEY:v1"
))
.header("content-type", "application/json")
.header(INSTALL_ID_HEADER, "agt_1")
.body(body.clone())
.send()
.await
.expect("the forward answers");
assert!(resp.status().is_success());
let _ = resp.bytes().await.unwrap();
assert_eq!(
up.received_body.lock().unwrap().clone(),
Some(body),
"the request reached the upstream byte-identically — transforms do not apply here"
);
let env = tokio::time::timeout(Duration::from_secs(3), rx.recv())
.await
.ok()
.flatten()
.expect("a captured route emits one economics event")
.envelope;
let data = &env["data"];
assert_eq!(
data["gen_ai.request.model"], "gemini-2.5-pro",
"the model comes off the PATH — and off `path()`, not `path_and_query()`"
);
assert!(
!data["gen_ai.request.model"]
.as_str()
.unwrap_or_default()
.contains("AIzaSy"),
"a fragment of the caller's API key reached the economics record: {}",
data["gen_ai.request.model"]
);
assert!(
data["ai.openlatch.capture.gap"].is_null(),
"a fully provider-reported turn on a route-named model has no gap to \
report, got {}",
data["ai.openlatch.capture.gap"]
);
assert_eq!(data["ai.openlatch.cost.basis"], "provider_reported");
assert_eq!(data["gen_ai.provider.name"], "google");
assert_eq!(
data["gen_ai.usage.output_tokens"],
64 + 32,
"candidates + thoughts, taken from the LAST chunk — not summed with \
the lead chunk's 9 + 4"
);
assert_eq!(data["gen_ai.usage.cache_read.input_tokens"], 256);
assert_eq!(
data["gen_ai.usage.input_tokens"],
1024 - 256,
"fresh input, through the real proxy"
);
assert_eq!(data["gen_ai.usage.cache_creation.input_tokens"], 0);
assert_eq!(data["ai.openlatch.cache.ephemeral_5m_input_tokens"], 0);
assert_eq!(data["ai.openlatch.cache.ephemeral_1h_input_tokens"], 0);
}
#[test]
fn opaque_path_still_resolves_a_gemini_model_from_the_route() {
let base = reqwest::Url::parse("http://127.0.0.1:1").unwrap();
let st = ModelRelayState::new(base, 0, 8, &[]);
let headers = HeaderMap::new();
let obs = observe_request_metadata_only(
&st,
WireFormat::GoogleGenerateContent,
&headers,
"/v1beta/models/gemini-2.5-pro:streamGenerateContent",
);
assert_eq!(obs.model.as_deref(), Some("gemini-2.5-pro"));
assert!(
obs.model_known,
"the route named it, so `unknown_model` would be a false gap"
);
let anthropic = observe_request_metadata_only(
&st,
WireFormat::AnthropicMessages,
&headers,
"/v1/messages",
);
assert_eq!(anthropic.model, None);
assert!(!anthropic.model_known);
}
#[test]
fn observe_request_resolves_a_gemini_model_from_the_route() {
let base = reqwest::Url::parse("http://127.0.0.1:1").unwrap();
let st = ModelRelayState::new(base, 0, 8, &[]);
let mut headers = HeaderMap::new();
headers.insert(INSTALL_ID_HEADER, "agt_1".parse().unwrap());
let bytes = Bytes::from_static(
br#"{"contents":[{"role":"user","parts":[{"text":"build the parser"}]}]}"#,
);
let obs = observe_request(
&st,
WireFormat::GoogleGenerateContent,
&headers,
"/v1beta/models/gemini-2.5-pro:streamGenerateContent",
&bytes,
)
.0;
assert_eq!(obs.model.as_deref(), Some("gemini-2.5-pro"));
assert!(obs.model_known);
let anthropic = Bytes::from_static(br#"{"model":"claude-opus-4-8","messages":[]}"#);
let obs = observe_request(
&st,
WireFormat::AnthropicMessages,
&headers,
"/v1beta/models/gemini-2.5-pro:streamGenerateContent",
&anthropic,
)
.0;
assert_eq!(obs.model.as_deref(), Some("claude-opus-4-8"));
assert!(obs.model_known);
}
#[test]
fn request_signals_read_a_generate_content_body() {
let body = serde_json::json!({
"systemInstruction": {"parts": [{"text": "You are Cline."}]},
"contents": [
{"role": "user", "parts": [{"text": "build the parser"}]},
{"role": "model", "parts": [{"text": "on it"}]},
{"role": "user", "parts": [
{"functionResponse": {"id": "fc_a", "name": "read_file", "response": {}}},
{"text": "now fix it"}
]}
]
});
let s = request_signals(WireFormat::GoogleGenerateContent, &body);
assert!(
s.prompt_hashes
.contains(&crate::model_relay::session::text_hash("build the parser")),
"the opening turn's text must be a candidate — it is under `parts`, not `content`"
);
assert!(
s.prompt_hashes
.contains(&crate::model_relay::session::text_hash("now fix it")),
"and the turn being answered right now"
);
assert!(
!s.prompt_hashes
.contains(&crate::model_relay::session::text_hash("You are Cline.")),
"`systemInstruction` is the agent's own preamble, identical across every \
session it starts — hashing it offers a candidate that matches strangers"
);
assert_eq!(s.tool_use_ids, vec!["fc_a".to_string()]);
assert_eq!(s.declared_session_id, None);
}
#[test]
fn only_the_relay_records_an_egress_outcome() {
let production = include_str!("proxy.rs")
.split("mod tests {")
.next()
.expect("the file has a production half");
let sites = production.matches("st.egress.record_").count();
assert_eq!(
sites, 2,
"egress outcomes must be recorded ONLY at relay's two branches — a third call site means another synth_502 path started reporting"
);
}
fn conversation() -> serde_json::Value {
serde_json::json!({
"messages": [
{"role": "user", "content": [{"type": "text", "text": "build the parser"}]},
{"role": "assistant", "content": [
{"type": "tool_use", "id": "toolu_aaa", "name": "Bash", "input": {}}
]},
{"role": "user", "content": [
{"type": "tool_result", "tool_use_id": "toolu_aaa", "content": "ok"},
{"type": "text", "text": "now check /repo/alpha/src/main.rs please"}
]}
]
})
}
fn claude_code_metadata(session_id: &str) -> serde_json::Value {
serde_json::json!({
"metadata": {
"user_id": format!(
r#"{{"device_id":"600058d8","account_uuid":"212c36cb","session_id":"{session_id}"}}"#
)
}
})
}
#[test]
fn the_request_names_its_own_session() {
let body = claude_code_metadata("5c7d9833-5e3b-4f88-9ece-cf371fb48c6a");
assert_eq!(
request_signals(WireFormat::AnthropicMessages, &body)
.declared_session_id
.as_deref(),
Some("5c7d9833-5e3b-4f88-9ece-cf371fb48c6a")
);
}
#[test]
fn anything_that_is_not_that_shape_is_silently_ignored() {
let cases = [
serde_json::json!({}), serde_json::json!({"metadata": {}}), serde_json::json!({"metadata": {"user_id": "plain-id"}}), serde_json::json!({"metadata": {"user_id": 42}}), serde_json::json!({"metadata": {"user_id": "{\"account_uuid\":\"x\"}"}}), serde_json::json!({"metadata": {"user_id": "{\"session_id\":\"\"}"}}), serde_json::json!({"metadata": {"user_id": "{\"session_id\":\"a\\u0000b\"}"}}), ];
for body in cases {
assert_eq!(
request_signals(WireFormat::AnthropicMessages, &body).declared_session_id,
None,
"must not trust {body}"
);
}
}
#[test]
fn an_absurdly_long_declared_session_id_is_refused() {
let body = serde_json::json!({
"metadata": {"user_id": format!(r#"{{"session_id":"{}"}}"#, "x".repeat(500))}
});
assert_eq!(
request_signals(WireFormat::AnthropicMessages, &body).declared_session_id,
None
);
}
#[test]
fn declared_session_id_is_unchanged_for_anthropic() {
let body = claude_code_metadata("5c7d9833-5e3b-4f88-9ece-cf371fb48c6a");
assert_eq!(
wire_format::declared_session_id(WireFormat::AnthropicMessages, &body).as_deref(),
Some("5c7d9833-5e3b-4f88-9ece-cf371fb48c6a")
);
let refused = [
serde_json::json!({}), serde_json::json!({"metadata": {}}), serde_json::json!({"metadata": {"user_id": "plain-id"}}), serde_json::json!({"metadata": {"user_id": 42}}), serde_json::json!({"metadata": {"user_id": "{\"account_uuid\":\"x\"}"}}), serde_json::json!({"metadata": {"user_id": "{\"session_id\":\"\"}"}}), serde_json::json!({"metadata": {"user_id": "{\"session_id\":\"a\\u0000b\"}"}}), serde_json::json!({
"metadata": {"user_id": format!(r#"{{"session_id":"{}"}}"#, "x".repeat(500))}
}), ];
for body in refused {
assert_eq!(
wire_format::declared_session_id(WireFormat::AnthropicMessages, &body),
None,
"must not trust {body}"
);
}
}
#[test]
fn observe_request_believes_the_session_the_request_names() {
let base = reqwest::Url::parse("http://127.0.0.1:1").unwrap();
let st = ModelRelayState::new(base, 0, 8, &[]);
for zombie in ["6693e5f6", "a5a39553"] {
st.registry
.upsert("agt_test", "agt_test", "claude-code", zombie);
}
let mut body = conversation();
body["metadata"] = claude_code_metadata("5c7d9833")["metadata"].clone();
let mut headers = HeaderMap::new();
headers.insert(INSTALL_ID_HEADER, "agt_test".parse().unwrap());
let bytes = Bytes::from(serde_json::to_vec(&body).unwrap());
let obs = observe_request(
&st,
WireFormat::AnthropicMessages,
&headers,
"/v1/messages",
&bytes,
)
.0;
assert_eq!(obs.session.session_id.as_deref(), Some("5c7d9833"));
assert_eq!(
obs.session.assurance,
crate::model_relay::session::Assurance::Attested
);
}
#[test]
fn request_signals_reads_the_first_and_last_user_turns() {
let s = request_signals(WireFormat::AnthropicMessages, &conversation());
assert_eq!(s.tool_use_ids, vec!["toolu_aaa".to_string()]);
assert!(s
.prompt_hashes
.contains(&crate::model_relay::session::text_hash("build the parser")));
}
#[test]
fn request_signals_offers_recent_turns_not_only_the_opening_one() {
let body = serde_json::json!({"messages": [
{"role": "user", "content": "hello this a session without a cache"},
{"role": "assistant", "content": "hi"},
{"role": "user", "content": "what is session Id"}
]});
let s = request_signals(WireFormat::AnthropicMessages, &body);
for p in ["hello this a session without a cache", "what is session Id"] {
assert!(
s.prompt_hashes
.contains(&crate::model_relay::session::text_hash(p)),
"missing candidate for {p:?} — an entry born mid-conversation \
holds only the later turns"
);
}
}
#[test]
fn request_signals_handles_a_bare_string_content() {
let body = serde_json::json!({
"messages": [{"role": "user", "content": "build the parser"}]
});
let s = request_signals(WireFormat::AnthropicMessages, &body);
assert!(s
.prompt_hashes
.contains(&crate::model_relay::session::text_hash("build the parser")));
assert_eq!(s.prompt_hashes.len(), 1);
}
#[test]
fn first_prompt_survives_a_block_the_agent_appended() {
let body = serde_json::json!({
"messages": [{"role": "user", "content": [
{"type": "text", "text": "build the parser"},
{"type": "text", "text": "<system-reminder>appended by the agent</system-reminder>"}
]}]
});
let s = request_signals(WireFormat::AnthropicMessages, &body);
assert!(
s.prompt_hashes
.contains(&crate::model_relay::session::text_hash("build the parser")),
"the typed prompt must survive as a candidate, else selector 2 never fires"
);
}
#[test]
fn request_signals_is_empty_for_a_body_with_no_messages() {
for body in [serde_json::json!({}), serde_json::json!({"messages": []})] {
let s = request_signals(WireFormat::AnthropicMessages, &body);
assert!(s.tool_use_ids.is_empty());
assert!(s.prompt_hashes.is_empty());
}
}
fn codex_turn() -> serde_json::Value {
serde_json::json!({
"model": "gpt-5-codex",
"input": [
{"type": "message", "id": "msg_1", "role": "developer",
"content": [{"type": "input_text", "text": "<skills_instructions>…"}]},
{"type": "message", "id": "msg_2", "role": "user",
"content": [{"type": "input_text", "text": "<environment_context>\n <cwd>/repo</cwd>\n</environment_context>"}]},
{"type": "message", "id": "msg_3", "role": "user",
"content": [{"type": "input_text", "text": "build the parser"}]},
{"type": "function_call", "id": "fc_1", "call_id": "call_aaa",
"name": "shell", "arguments": "{}"},
{"type": "function_call_output", "id": "fco_1", "call_id": "call_aaa",
"output": "ok"}
]
})
}
#[test]
fn request_signals_reads_a_codex_responses_turn() {
let s = request_signals(WireFormat::OpenAiResponses, &codex_turn());
assert_eq!(s.tool_use_ids, vec!["call_aaa".to_string()]);
assert!(
s.prompt_hashes
.contains(&crate::model_relay::session::text_hash("build the parser")),
"the typed prompt must survive as a candidate, else selector 2 never fires"
);
assert!(
!s.prompt_hashes
.contains(&crate::model_relay::session::text_hash(
"<skills_instructions>…"
)),
"the agent's own preamble is not a signal about WHICH session this is"
);
}
#[test]
fn a_codex_body_read_as_the_wrong_format_yields_nothing() {
let s = request_signals(WireFormat::AnthropicMessages, &codex_turn());
assert!(s.tool_use_ids.is_empty());
assert!(s.prompt_hashes.is_empty());
let s = request_signals(WireFormat::OpenAiResponses, &conversation());
assert!(s.tool_use_ids.is_empty());
assert!(s.prompt_hashes.is_empty());
}
#[test]
fn message_text_truncates_on_a_char_boundary() {
let body = serde_json::json!({
"messages": [{"role": "user", "content": "é".repeat(4096)}]
});
let msg = &body["messages"][0];
let text = message_text(msg, "content", Some("text"), 101);
assert!(text.len() <= 101);
assert!(text.chars().all(|c| c == 'é'));
}
fn observed_session() -> Observation {
let base = reqwest::Url::parse("http://127.0.0.1:1").unwrap();
let st = ModelRelayState::new(base, 0, 8, &[]);
st.registry.upsert_signals(
"agt_test",
"agt_test",
"claude-code",
"sess_a",
&session::SessionSignals {
tool_use_id: Some("toolu_aaa"),
prompt: Some("build the parser"),
},
);
std::thread::sleep(Duration::from_millis(5));
st.registry
.upsert("agt_test", "agt_test", "claude-code", "sess_b");
let mut headers = HeaderMap::new();
headers.insert(INSTALL_ID_HEADER, "agt_test".parse().unwrap());
let bytes = Bytes::from(serde_json::to_vec(&conversation()).unwrap());
observe_request(
&st,
WireFormat::AnthropicMessages,
&headers,
"/v1/messages",
&bytes,
)
.0
}
#[test]
fn observe_request_attributes_to_the_session_the_content_identifies() {
let obs = observed_session();
assert_eq!(obs.session.session_id.as_deref(), Some("sess_a"));
assert_eq!(
obs.session.assurance,
crate::model_relay::session::Assurance::Attested,
"an exact tool-result join identifies the session; it does not guess"
);
}
#[test]
fn the_opaque_path_still_falls_back_and_says_it_guessed() {
let base = reqwest::Url::parse("http://127.0.0.1:1").unwrap();
let st = ModelRelayState::new(base, 0, 8, &[]);
st.registry.upsert_signals(
"agt_test",
"agt_test",
"claude-code",
"sess_a",
&session::SessionSignals {
tool_use_id: Some("toolu_aaa"),
prompt: Some("build the parser"),
},
);
std::thread::sleep(Duration::from_millis(5));
st.registry
.upsert("agt_test", "agt_test", "claude-code", "sess_b");
let mut headers = HeaderMap::new();
headers.insert(INSTALL_ID_HEADER, "agt_test".parse().unwrap());
let obs = observe_request_metadata_only(
&st,
WireFormat::AnthropicMessages,
&headers,
"/v1/messages",
);
assert_eq!(obs.session.session_id.as_deref(), Some("sess_b"));
assert_eq!(
obs.session.assurance,
crate::model_relay::session::Assurance::Inferred
);
}
#[test]
fn observe_request_attributes_a_codex_turn_to_the_session_that_ran_the_tool() {
let base = reqwest::Url::parse("http://127.0.0.1:1").unwrap();
let st = ModelRelayState::new(base, 0, 8, &[]);
st.registry.upsert_signals(
"agt_test",
"agt_test",
"codex-cli",
"sess_a",
&session::SessionSignals {
tool_use_id: Some("call_aaa"),
prompt: Some("build the parser"),
},
);
std::thread::sleep(Duration::from_millis(5));
st.registry
.upsert("agt_test", "agt_test", "codex-cli", "sess_b");
let mut headers = HeaderMap::new();
headers.insert(INSTALL_ID_HEADER, "agt_test".parse().unwrap());
let bytes = Bytes::from(serde_json::to_vec(&codex_turn()).unwrap());
let obs = observe_request(
&st,
WireFormat::OpenAiResponses,
&headers,
"/v1/responses",
&bytes,
)
.0;
assert_eq!(obs.session.session_id.as_deref(), Some("sess_a"));
assert_eq!(
obs.session.assurance,
crate::model_relay::session::Assurance::Attested,
"an exact call-id join identifies the session; it does not guess"
);
}
#[test]
fn a_codex_opening_turn_is_attributed_by_its_prompt() {
let base = reqwest::Url::parse("http://127.0.0.1:1").unwrap();
let st = ModelRelayState::new(base, 0, 8, &[]);
st.registry.upsert_signals(
"agt_test",
"agt_test",
"codex-cli",
"sess_a",
&session::SessionSignals {
tool_use_id: None,
prompt: Some("build the parser"),
},
);
std::thread::sleep(Duration::from_millis(5));
st.registry
.upsert("agt_test", "agt_test", "codex-cli", "sess_b");
let body = serde_json::json!({"input": [
{"type": "message", "role": "user",
"content": [{"type": "input_text", "text": "build the parser"}]}
]});
let mut headers = HeaderMap::new();
headers.insert(INSTALL_ID_HEADER, "agt_test".parse().unwrap());
let bytes = Bytes::from(serde_json::to_vec(&body).unwrap());
let obs = observe_request(
&st,
WireFormat::OpenAiResponses,
&headers,
"/v1/responses",
&bytes,
)
.0;
assert_eq!(obs.session.session_id.as_deref(), Some("sess_a"));
assert_eq!(
obs.session.assurance,
crate::model_relay::session::Assurance::Attested
);
}
#[test]
fn a_codex_request_that_names_its_session_is_believed_over_the_joins() {
let base = reqwest::Url::parse("http://127.0.0.1:1").unwrap();
let st = ModelRelayState::new(base, 0, 8, &[]);
st.registry
.upsert("agt_test", "agt_test", "codex-cli", "sess_a");
std::thread::sleep(Duration::from_millis(5));
st.registry
.upsert("agt_test", "agt_test", "codex-cli", "sess_b");
let mut body = codex_turn();
body["client_metadata"] = serde_json::json!({
"session_id": "01a07677-988f-7901-af48-225c7386aabc",
"thread_id": "01a07677-988f-7901-af48-225c7386aabc",
"turn_id": "01a07677-7334-74a2-be87-bdd0bbe8fb77",
"x-codex-installation-id": "5532dc79-4512-4bae-ab73-0dd3e814d329"
});
let mut headers = HeaderMap::new();
headers.insert(INSTALL_ID_HEADER, "agt_test".parse().unwrap());
let bytes = Bytes::from(serde_json::to_vec(&body).unwrap());
let obs = observe_request(
&st,
WireFormat::OpenAiResponses,
&headers,
"/v1/responses",
&bytes,
)
.0;
assert_eq!(
obs.session.session_id.as_deref(),
Some("01a07677-988f-7901-af48-225c7386aabc"),
"the request named its own session; nothing else gets a vote"
);
assert_eq!(
obs.session.assurance,
crate::model_relay::session::Assurance::Attested
);
}
#[test]
fn content_length_parses_and_defaults() {
let mut h = HeaderMap::new();
assert_eq!(content_length(&h), None); h.insert(CONTENT_LENGTH, HeaderValue::from_static("42"));
assert_eq!(content_length(&h), Some(42));
}
#[test]
fn forward_headers_strips_framing_keeps_credential() {
let mut h = HeaderMap::new();
h.insert(HOST, HeaderValue::from_static("127.0.0.1:7600"));
h.insert(CONTENT_LENGTH, HeaderValue::from_static("10"));
h.insert("x-api-key", HeaderValue::from_static("sk-ant-xyz"));
h.insert("anthropic-version", HeaderValue::from_static("2023-06-01"));
let out = forward_headers(&h);
assert!(
out.get(HOST).is_none(),
"host must be stripped (reqwest sets it)"
);
assert!(
out.get(CONTENT_LENGTH).is_none(),
"content-length must be stripped"
);
assert_eq!(out.get("x-api-key").unwrap(), "sk-ant-xyz");
assert_eq!(out.get("anthropic-version").unwrap(), "2023-06-01");
}
#[test]
fn forward_headers_strips_accept_encoding_for_readable_usage() {
let mut h = HeaderMap::new();
h.insert(ACCEPT_ENCODING, HeaderValue::from_static("gzip, br, zstd"));
h.insert("x-api-key", HeaderValue::from_static("sk-ant-xyz"));
let out = forward_headers(&h);
assert!(
out.get(ACCEPT_ENCODING).is_none(),
"accept-encoding must be stripped so the response body is scannable"
);
assert_eq!(out.get("x-api-key").unwrap(), "sk-ant-xyz");
}
#[test]
fn forward_headers_strips_connection_listed_and_hop_by_hop() {
let mut h = HeaderMap::new();
h.insert(CONNECTION, HeaderValue::from_static("x-internal-foo"));
h.insert("x-internal-foo", HeaderValue::from_static("secret"));
h.insert("keep-alive", HeaderValue::from_static("timeout=5"));
h.insert("x-api-key", HeaderValue::from_static("sk-ant-xyz"));
h.insert("authorization", HeaderValue::from_static("Bearer tok"));
let out = forward_headers(&h);
assert!(
out.get("x-internal-foo").is_none(),
"a Connection-listed header must be stripped"
);
assert!(
out.get("keep-alive").is_none(),
"Keep-Alive is hop-by-hop and must be stripped"
);
assert!(
out.get(CONNECTION).is_none(),
"Connection itself is stripped"
);
assert_eq!(out.get("x-api-key").unwrap(), "sk-ant-xyz");
assert_eq!(out.get("authorization").unwrap(), "Bearer tok");
}
#[tokio::test(flavor = "multi_thread")]
async fn header_wait_times_out_to_synth_502() {
let hang_port = mock::spawn_hang_after_accept().await;
let base = reqwest::Url::parse(&format!("http://127.0.0.1:{hang_port}")).unwrap();
let state = Arc::new(
ModelRelayState::new(base, 0, 8, &[]).with_header_timeout(Duration::from_millis(200)),
);
let port = serve_ephemeral(state).await;
let resp = tokio::time::timeout(
Duration::from_secs(5),
crate::egress::client()
.post(format!("http://127.0.0.1:{port}/v1/messages"))
.header("content-type", "application/json")
.body(br#"{"model":"x","messages":[]}"#.to_vec())
.send(),
)
.await
.expect("request must return within 5s, not hang")
.unwrap();
assert_eq!(resp.status(), StatusCode::BAD_GATEWAY);
assert_eq!(
resp.headers().get("x-openlatch-upstream").unwrap(),
"unreachable"
);
}
#[tokio::test(flavor = "multi_thread")]
async fn admin_status_endpoint_reports_the_listener() {
let port = serve_ephemeral(state_for(0)).await;
let resp = crate::egress::client()
.get(format!("http://127.0.0.1:{port}/admin/model-relay/status"))
.send()
.await
.unwrap();
assert!(resp.status().is_success());
let v: serde_json::Value = resp.json().await.unwrap();
assert_eq!(v["status"], "up");
assert!(v["port"].is_number());
assert!(v["pass_through_failures"].is_number());
}
#[tokio::test]
async fn synth_502_shape() {
let r = synth_502(&state_for(1));
assert_eq!(r.status(), StatusCode::BAD_GATEWAY);
assert_eq!(
r.headers().get("x-openlatch-upstream").unwrap(),
"unreachable"
);
}
#[tokio::test(flavor = "multi_thread")]
async fn opaque_get_models_forwarded() {
let up = mock::spawn_capture_200().await;
let port = serve_ephemeral(state_for(up.port)).await;
let resp = crate::egress::client()
.get(format!("http://127.0.0.1:{port}/v1/models"))
.send()
.await
.unwrap();
assert!(resp.status().is_success());
assert_eq!(resp.bytes().await.unwrap().as_ref(), b"ok");
let line = up.received_request_line.lock().unwrap().clone().unwrap();
assert!(
line.starts_with("GET /v1/models"),
"opaque path forwards verbatim: {line}"
);
}
#[tokio::test(flavor = "multi_thread")]
async fn materialized_messages_forwards_body_and_credential() {
let up = mock::spawn_capture_200().await;
let port = serve_ephemeral(state_for(up.port)).await;
let body = br#"{"model":"claude-opus-4-8","messages":[{"role":"user","content":"hi"}]}"#;
let resp = crate::egress::client()
.post(format!("http://127.0.0.1:{port}/v1/messages"))
.header("content-type", "application/json")
.header("x-api-key", "sk-ant-test123")
.body(body.to_vec())
.send()
.await
.unwrap();
assert!(resp.status().is_success());
assert_eq!(
up.received_body.lock().unwrap().clone().unwrap(),
body.to_vec()
);
assert_eq!(up.header("x-api-key").as_deref(), Some("sk-ant-test123"));
assert_eq!(
up.header("host").as_deref(),
Some(format!("127.0.0.1:{}", up.port).as_str())
);
}
#[tokio::test(flavor = "multi_thread")]
async fn breakpoint_body_round_trips_byte_identical() {
let up = mock::spawn_capture_200().await;
let port = serve_ephemeral(state_for(up.port)).await;
let body = br#"{"model":"claude-opus-4-8","system":[{"type":"text","text":"x","cache_control":{"type":"ephemeral"}}],"messages":[]}"#;
let resp = crate::egress::client()
.post(format!("http://127.0.0.1:{port}/v1/messages"))
.header("content-type", "application/json")
.body(body.to_vec())
.send()
.await
.unwrap();
assert!(resp.status().is_success());
let received = up.received_body.lock().unwrap().clone().unwrap();
assert_eq!(
received,
body.to_vec(),
"cache_control breakpoint must round-trip byte-identically"
);
}
#[tokio::test(flavor = "multi_thread")]
async fn count_tokens_forwarded_opaque_with_body() {
let up = mock::spawn_capture_200().await;
let port = serve_ephemeral(state_for(up.port)).await;
let body = br#"{"model":"claude-opus-4-8","messages":[]}"#;
let resp = crate::egress::client()
.post(format!("http://127.0.0.1:{port}/v1/messages/count_tokens"))
.header("content-type", "application/json")
.body(body.to_vec())
.send()
.await
.unwrap();
assert!(resp.status().is_success());
assert_eq!(
up.received_body.lock().unwrap().clone().unwrap(),
body.to_vec()
);
}
#[tokio::test(flavor = "multi_thread")]
async fn unreachable_upstream_yields_synth_502() {
let dead = mock::closed_port().await;
let port = serve_ephemeral(state_for(dead)).await;
let resp = crate::egress::client()
.post(format!("http://127.0.0.1:{port}/v1/messages"))
.header("content-type", "application/json")
.body(br#"{"model":"x","messages":[]}"#.to_vec())
.send()
.await
.unwrap();
assert_eq!(resp.status(), StatusCode::BAD_GATEWAY);
assert_eq!(
resp.headers().get("x-openlatch-upstream").unwrap(),
"unreachable"
);
}
#[tokio::test(flavor = "multi_thread")]
async fn an_unreachable_upstream_is_counted_as_an_upstream_failure() {
let _serialised = SCAN_PANIC_LOCK.lock().await;
let dead = mock::closed_port().await;
let port = serve_ephemeral(state_for(dead)).await;
let upstream_before = upstream_failures();
let pass_through_before = pass_through_failures();
let resp = crate::egress::client()
.post(format!("http://127.0.0.1:{port}/v1/messages"))
.header("content-type", "application/json")
.body(br#"{"model":"x","messages":[]}"#.to_vec())
.send()
.await
.unwrap();
assert_eq!(resp.status(), StatusCode::BAD_GATEWAY);
assert!(
upstream_failures() > upstream_before,
"a request the agent got no answer to must be counted"
);
assert_eq!(
pass_through_failures(),
pass_through_before,
"an unreachable upstream is not a degraded-to-pass-through step"
);
}
#[tokio::test(flavor = "multi_thread")]
async fn preflight_probe_forwards_without_measuring_or_leaking_its_marker() {
use crate::model_relay::preflight::PREFLIGHT_HEADER;
use crate::model_relay::session::SessionRegistry;
let up = mock::spawn_capture_200().await;
let base = reqwest::Url::parse(&format!("http://127.0.0.1:{}", up.port)).unwrap();
let (tx, mut rx) = tokio::sync::mpsc::channel(8);
let reg = Arc::new(SessionRegistry::default());
reg.upsert("agt_1", "agt_1", "claude-code", "sess_a");
let state = Arc::new(ModelRelayState::new(base, 0, 8, &[]).with_measurement(reg, Some(tx)));
let port = serve_ephemeral(state).await;
let resp = crate::egress::client()
.post(format!("http://127.0.0.1:{port}/v1/messages"))
.header("content-type", "application/json")
.header("x-openlatch-install-id", "agt_1")
.header(PREFLIGHT_HEADER, "1")
.body(
br#"{"model":"claude-opus-4-8","max_tokens":1,"messages":[{"role":"user","content":"ping"}]}"#
.to_vec(),
)
.send()
.await
.unwrap();
assert!(
resp.status().is_success(),
"the probe must be forwarded like any other request"
);
let _ = resp.bytes().await;
for _ in 0..50 {
if up.received_headers.lock().unwrap().is_some() {
break;
}
tokio::time::sleep(Duration::from_millis(20)).await;
}
let sent = up.received_headers.lock().unwrap().clone().unwrap();
assert!(
!sent.to_ascii_lowercase().contains(PREFLIGHT_HEADER),
"the preflight marker must be stripped before the request leaves for the provider, \
got headers: {sent}"
);
tokio::time::sleep(Duration::from_millis(50)).await;
assert!(
rx.try_recv().is_err(),
"a preflight probe must emit ZERO economics events"
);
}
static SCAN_PANIC_LOCK: tokio::sync::Mutex<()> = tokio::sync::Mutex::const_new(());
struct ScanPanicInjection;
impl ScanPanicInjection {
fn arm() -> Self {
set_inject_scan_panic(true);
Self
}
}
impl Drop for ScanPanicInjection {
fn drop(&mut self) {
set_inject_scan_panic(false);
}
}
#[tokio::test(flavor = "multi_thread")]
async fn an_unmarked_request_on_the_same_path_is_measured() {
let _serialised = SCAN_PANIC_LOCK.lock().await;
use crate::model_relay::session::SessionRegistry;
let up = mock::spawn_capture_200().await;
let base = reqwest::Url::parse(&format!("http://127.0.0.1:{}", up.port)).unwrap();
let (tx, mut rx) = tokio::sync::mpsc::channel(8);
let reg = Arc::new(SessionRegistry::default());
reg.upsert("agt_1", "agt_1", "claude-code", "sess_a");
let state = Arc::new(ModelRelayState::new(base, 0, 8, &[]).with_measurement(reg, Some(tx)));
let port = serve_ephemeral(state).await;
let resp = crate::egress::client()
.post(format!("http://127.0.0.1:{port}/v1/messages"))
.header("content-type", "application/json")
.header("x-openlatch-install-id", "agt_1")
.body(
br#"{"model":"claude-opus-4-8","max_tokens":1,"messages":[{"role":"user","content":"ping"}]}"#
.to_vec(),
)
.send()
.await
.unwrap();
let _ = resp.bytes().await;
tokio::time::sleep(Duration::from_millis(50)).await;
assert!(
rx.try_recv().is_ok(),
"an ordinary /v1/messages request must still emit its economics event"
);
}
#[tokio::test(flavor = "multi_thread")]
async fn usage_scan_panic_degrades_to_unmeasured() {
let _serialised = SCAN_PANIC_LOCK.lock().await;
use crate::model_relay::session::SessionRegistry;
let up = mock::spawn_capture_usage_sse(10, 0, 0, 0, 0, 20).await;
let base = reqwest::Url::parse(&format!("http://127.0.0.1:{}", up.port)).unwrap();
let (tx, mut rx) = tokio::sync::mpsc::channel(8);
let reg = Arc::new(SessionRegistry::default());
reg.upsert("agt_1", "agt_1", "claude-code", "sess_a");
let state = Arc::new(ModelRelayState::new(base, 0, 8, &[]).with_measurement(reg, Some(tx)));
let port = serve_ephemeral(state).await;
let failures_before = pass_through_failures();
let injection = ScanPanicInjection::arm();
let body = br#"{"model":"claude-opus-4-8","stream":true,"messages":[{"role":"user","content":"hi"}]}"#.to_vec();
let resp = crate::egress::client()
.post(format!("http://127.0.0.1:{port}/v1/messages"))
.header("content-type", "application/json")
.header("x-openlatch-install-id", "agt_1")
.body(body.clone())
.send()
.await
.unwrap();
assert!(
resp.status().is_success(),
"forward must complete despite the scan panic"
);
let received = resp.bytes().await.expect("response body drains cleanly");
assert!(!received.is_empty(), "response body still flows through");
drop(injection);
for _ in 0..50 {
if up.received_body.lock().unwrap().is_some() {
break;
}
tokio::time::sleep(Duration::from_millis(20)).await;
}
assert_eq!(
up.received_body.lock().unwrap().clone().unwrap(),
body,
"request body must be byte-identical despite the scan panic"
);
assert!(
pass_through_failures() > failures_before,
"a usage-scan panic must be recorded as a pass-through failure"
);
assert!(
tokio::time::timeout(Duration::from_millis(500), rx.recv())
.await
.ok()
.flatten()
.is_none(),
"a usage-scan panic must emit NO event (unmeasured, not corrupt)"
);
}
#[tokio::test(flavor = "multi_thread")]
async fn streaming_is_zero_buffer() {
let gap = Duration::from_millis(60);
let trickle = mock::spawn_trickle_sse(3, gap).await;
let port = serve_ephemeral(state_for(trickle.port)).await;
let resp = crate::egress::client()
.post(format!("http://127.0.0.1:{port}/v1/messages"))
.header("content-type", "application/json")
.body(br#"{"model":"x","messages":[],"stream":true}"#.to_vec())
.send()
.await
.unwrap();
assert!(resp.status().is_success());
let mut stream = resp.bytes_stream();
let first = stream.next().await;
let first_byte_at = Instant::now();
assert!(first.is_some(), "expected at least one streamed chunk");
assert!(first.unwrap().is_ok());
while stream.next().await.is_some() {}
let final_written_at = trickle
.final_written_at
.lock()
.unwrap()
.expect("mock must have finished writing");
assert!(
first_byte_at < final_written_at,
"first client byte must arrive BEFORE the upstream stream completes (zero-buffer)"
);
}
}