use std::sync::Arc;
use std::time::Duration;
use axum_test::TestServer;
use bitrouter_sdk::config;
use bitrouter_sdk::server::{AppState, build_router};
use wiremock::matchers::{method, path as wm_path};
use wiremock::{Mock, MockServer, ResponseTemplate};
use bitrouter::auth::{NewApiKey, db as auth_db, generate};
use bitrouter::daemon::ObserveStatusProvider;
use bitrouter::metering::entities::requests;
use sea_orm::{ColumnTrait, DatabaseConnection, EntityTrait, PaginatorTrait, QueryFilter};
struct FullStack {
server: TestServer,
db: DatabaseConnection,
brvk_secret: String,
upstream: MockServer,
otlp_collector: MockServer,
observe: Arc<dyn ObserveStatusProvider>,
_policy_dir: PolicyDir,
}
impl FullStack {
async fn teardown(self) {
self.observe.shutdown().await;
}
}
struct PolicyDir(std::path::PathBuf);
impl Drop for PolicyDir {
fn drop(&mut self) {
let _ = std::fs::remove_dir_all(&self.0);
}
}
async fn assemble_full_stack() -> FullStack {
assemble_full_stack_inner(true).await
}
async fn assemble_full_stack_traces_only() -> FullStack {
assemble_full_stack_inner(false).await
}
async fn assemble_full_stack_inner(metrics_enabled: bool) -> FullStack {
let policy_dir = std::env::temp_dir().join(format!(
"bitrouter-fullstack-pol-{}",
std::time::SystemTime::now()
.duration_since(std::time::UNIX_EPOCH)
.unwrap()
.as_nanos()
));
tokio::fs::create_dir_all(&policy_dir).await.unwrap();
tokio::fs::write(
policy_dir.join("pol_main.yaml"),
"id: pol_main\n\
allowed_models: [test-model]\n\
allowed_tools: [search]\n\
max_spend_micro_usd: 100000000\n",
)
.await
.unwrap();
let upstream = MockServer::start().await;
let sse_body = build_sse_stream_with_ssn();
Mock::given(method("POST"))
.and(wm_path("/chat/completions"))
.respond_with(
ResponseTemplate::new(200)
.insert_header("content-type", "text/event-stream")
.set_body_raw(sse_body, "text/event-stream"),
)
.mount(&upstream)
.await;
let otlp_collector = MockServer::start().await;
Mock::given(method("POST"))
.and(wm_path("/v1/traces"))
.respond_with(ResponseTemplate::new(200))
.mount(&otlp_collector)
.await;
let observe_block = if metrics_enabled {
format!(" otlp_endpoint: \"{otlp}\"", otlp = otlp_collector.uri())
} else {
format!(
" otel:\n endpoint: \"{otlp}\"\n metrics:\n enabled: false",
otlp = otlp_collector.uri()
)
};
let yaml = format!(
r#"
server:
listen: "127.0.0.1:0"
skip_auth: false
database:
url: "sqlite::memory:"
providers:
mock:
api_base: {upstream}
api_key: test-key
api_protocol:
- "*": chat_completions
models:
- id: test-model
pricing:
input_micro_usd_per_token: 2.0
output_micro_usd_per_token: 10.0
plugins:
bitrouter-policy:
# Single-quoted scalar — double quotes would interpret the backslashes
# in a Windows temp path (`C:\Users\…`) as escape sequences and trip
# the parser on `\U`. Single quotes treat the value literally.
policy_dir: '{policy_path}'
bitrouter-guardrails:
custom_patterns:
- {{ name: ssn, pattern: '\d{{3}}-\d{{2}}-\d{{4}}', action: redact }}
- {{ name: forbidden, pattern: '(?i)forbidden', action: block }}
bitrouter-observe:
{observe_block}
"#,
upstream = upstream.uri(),
policy_path = policy_dir.display(),
);
let cfg: config::Config = config::parse_with(&yaml, |_| None).expect("config parses");
let assembled = bitrouter::build_app(&cfg).await.expect("app assembles");
let observe = assembled.observe.clone();
let user = "fullstack-user";
auth_db::upsert_user(&assembled.db, user).await.unwrap();
let key = generate();
let key_id = "fullstack-key".to_string();
auth_db::insert_api_key(
&assembled.db,
&NewApiKey {
id: key_id,
key_hash: key.hash.clone(),
user_id: user.to_string(),
spend_limit_micro_usd: None,
rpm_limit: None,
policy_id: Some("pol_main".to_string()),
},
)
.await
.unwrap();
let state = AppState {
language_model: assembled.app.language_model().unwrap().clone(),
mcp: assembled.app.mcp().cloned(),
skip_auth: assembled.app.skip_auth(),
metrics_renderer: assembled.app.metrics_renderer().cloned(),
};
let router = build_router(state);
let server = TestServer::new(router);
FullStack {
server,
db: assembled.db,
brvk_secret: key.secret,
upstream,
otlp_collector,
observe,
_policy_dir: PolicyDir(policy_dir),
}
}
fn build_sse_stream_with_ssn() -> String {
let mut out = String::new();
let chunk = |delta: serde_json::Value, finish: Option<&str>| -> String {
let body = serde_json::json!({
"id": "chatcmpl-mock",
"object": "chat.completion.chunk",
"model": "test-model",
"choices": [{"index": 0, "delta": delta, "finish_reason": finish}],
});
format!("data: {}\n\n", body)
};
out.push_str(&chunk(
serde_json::json!({"role": "assistant", "content": "your "}),
None,
));
out.push_str(&chunk(
serde_json::json!({"content": "ssn is 123-45-6789 ok"}),
None,
));
let finish = serde_json::json!({
"id": "chatcmpl-mock",
"object": "chat.completion.chunk",
"model": "test-model",
"choices": [{"index": 0, "delta": {}, "finish_reason": "stop"}],
"usage": {"prompt_tokens": 11, "completion_tokens": 7, "total_tokens": 18},
});
out.push_str(&format!("data: {}\n\n", finish));
out.push_str("data: [DONE]\n\n");
out
}
fn clean_body(stream: bool) -> serde_json::Value {
serde_json::json!({
"model": "test-model",
"messages": [{"role": "user", "content": "hello"}],
"tools": [{"function": {"name": "search"}}],
"stream": stream,
})
}
#[tokio::test]
async fn e2e_full_stack_streaming_redacts_and_meters() {
let fs = assemble_full_stack().await;
let resp = fs
.server
.post("/v1/chat/completions")
.add_header("authorization", format!("Bearer {}", fs.brvk_secret))
.add_header("accept", "text/event-stream")
.json(&clean_body(true))
.await;
resp.assert_status_ok();
assert_eq!(
resp.header("content-type").to_str().unwrap(),
"text/event-stream",
"streaming response must advertise text/event-stream",
);
let body = resp.text();
assert!(
body.contains("[REDACTED]"),
"GuardrailStreamHook should have inserted [REDACTED] into the SSE body; got:\n{body}"
);
assert!(
!body.contains("123-45-6789"),
"GuardrailStreamHook must have stripped the SSN-shaped span from the SSE body; got:\n{body}"
);
wait_for_metering_row(&fs.db, "fullstack-key").await;
let row = requests::Entity::find()
.filter(requests::Column::ApiKeyId.eq("fullstack-key"))
.one(&fs.db)
.await
.expect("metering query runs")
.expect("metering wrote a row for the streamed request");
assert_eq!(
row.model_id, "test-model",
"metering row records the routed model"
);
assert_eq!(
row.streamed, 1,
"metering row marks the request as streamed"
);
assert_eq!(
row.estimated_charge_micro_usd, 92,
"estimated_charge_micro_usd should be 2*11 + 10*7 = 92µ$",
);
wait_for_otlp(&fs.otlp_collector).await;
fs.teardown().await;
}
#[tokio::test]
async fn e2e_full_stack_guardrail_blocks_at_pre_request() {
let fs = assemble_full_stack().await;
let resp = fs
.server
.post("/v1/chat/completions")
.add_header("authorization", format!("Bearer {}", fs.brvk_secret))
.expect_failure()
.json(&serde_json::json!({
"model": "test-model",
"messages": [{"role": "user", "content": "please do the forbidden thing"}],
"tools": [{"function": {"name": "search"}}],
}))
.await;
resp.assert_status_bad_request();
let upstream_requests = fs.upstream.received_requests().await.unwrap_or_default();
assert!(
upstream_requests.is_empty(),
"blocked request must not reach the upstream; saw {} request(s)",
upstream_requests.len(),
);
let row_count = requests::Entity::find()
.filter(requests::Column::ApiKeyId.eq("fullstack-key"))
.count(&fs.db)
.await
.expect("count query runs");
assert_eq!(
row_count, 0,
"blocked request must not produce a metering row",
);
fs.teardown().await;
}
#[tokio::test]
async fn e2e_full_stack_policy_denies_disallowed_tool_before_guardrails() {
let fs = assemble_full_stack().await;
let resp = fs
.server
.post("/v1/chat/completions")
.add_header("authorization", format!("Bearer {}", fs.brvk_secret))
.expect_failure()
.json(&serde_json::json!({
"model": "test-model",
"messages": [{"role": "user", "content": "clean prompt"}],
"tools": [{"function": {"name": "filesystem"}}],
}))
.await;
resp.assert_status_forbidden();
let err: serde_json::Value = resp.json();
let msg = err["error"]["message"].as_str().unwrap_or("");
assert!(
msg.contains("tool") && msg.contains("filesystem"),
"deny reason must come from PolicyHook (tool restriction), not guardrails; got: {msg}",
);
let upstream_requests = fs.upstream.received_requests().await.unwrap_or_default();
assert!(upstream_requests.is_empty());
let row_count = requests::Entity::find()
.filter(requests::Column::ApiKeyId.eq("fullstack-key"))
.count(&fs.db)
.await
.expect("count query runs");
assert_eq!(row_count, 0);
fs.teardown().await;
}
const ASYNC_WAIT_BUDGET_MS: u64 = 5_000;
const POLL_INTERVAL_MS: u64 = 50;
async fn wait_for_metering_row(db: &DatabaseConnection, api_key_id: &str) {
for _ in 0..(ASYNC_WAIT_BUDGET_MS / POLL_INTERVAL_MS) {
let count = requests::Entity::find()
.filter(requests::Column::ApiKeyId.eq(api_key_id))
.count(db)
.await
.expect("count query runs");
if count >= 1 {
return;
}
tokio::time::sleep(Duration::from_millis(POLL_INTERVAL_MS)).await;
}
panic!(
"metering row never appeared within {}ms of stream completion",
ASYNC_WAIT_BUDGET_MS
);
}
async fn wait_for_otlp(collector: &MockServer) {
for _ in 0..(ASYNC_WAIT_BUDGET_MS / POLL_INTERVAL_MS) {
let n = collector
.received_requests()
.await
.map(|r| r.len())
.unwrap_or(0);
if n >= 1 {
return;
}
tokio::time::sleep(Duration::from_millis(POLL_INTERVAL_MS)).await;
}
panic!(
"OTLP collector never received an export within {}ms",
ASYNC_WAIT_BUDGET_MS
);
}
#[tokio::test]
async fn e2e_full_stack_outbound_traceparent_propagation() {
let fs = assemble_full_stack().await;
let inbound_trace_id = "aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa";
let inbound_traceparent = format!("00-{inbound_trace_id}-bbbbbbbbbbbbbbbb-01");
let resp = fs
.server
.post("/v1/chat/completions")
.add_header("authorization", format!("Bearer {}", fs.brvk_secret))
.add_header("accept", "text/event-stream")
.add_header("traceparent", inbound_traceparent.clone())
.json(&clean_body(true))
.await;
resp.assert_status_ok();
let _ = resp.text();
wait_for_metering_row(&fs.db, "fullstack-key").await;
let upstream_reqs = fs
.upstream
.received_requests()
.await
.expect("upstream received at least one request");
let upstream = upstream_reqs
.first()
.expect("upstream got the chat completion call");
let tp = upstream
.headers
.get("traceparent")
.expect("`HttpExecutor` must merge an outbound `traceparent` set by the observe hook")
.to_str()
.expect("`traceparent` is ASCII");
assert!(
tp.starts_with(&format!("00-{inbound_trace_id}-")),
"outbound traceparent must continue the inbound trace_id; got {tp}"
);
assert_eq!(
tp.len(),
55,
"outbound traceparent must be 55 chars; got {tp}"
);
fs.teardown().await;
}
async fn collect_exported_trace_spans(
collector: &MockServer,
) -> Vec<opentelemetry_proto::tonic::trace::v1::Span> {
use opentelemetry_proto::tonic::collector::trace::v1::ExportTraceServiceRequest;
use prost::Message;
let requests = collector.received_requests().await.unwrap_or_default();
let mut spans = Vec::new();
for req in &requests {
let parsed = ExportTraceServiceRequest::decode(req.body.as_slice())
.expect("every collector body is a trace export (metric export is off)");
for resource_spans in parsed.resource_spans {
for scope_spans in resource_spans.scope_spans {
spans.extend(scope_spans.spans);
}
}
}
spans
}
fn span_str_attr<'a>(
span: &'a opentelemetry_proto::tonic::trace::v1::Span,
key: &str,
) -> Option<&'a str> {
use opentelemetry_proto::tonic::common::v1::any_value::Value as AnyValue;
let kv = span.attributes.iter().find(|kv| kv.key == key)?;
let value = kv.value.as_ref()?;
match value.value.as_ref()? {
AnyValue::StringValue(s) => Some(s.as_str()),
_ => None,
}
}
#[tokio::test]
async fn e2e_full_stack_otel_span_hierarchy() {
let fs = assemble_full_stack_traces_only().await;
let inbound_trace_id = "aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa";
let inbound_traceparent = format!("00-{inbound_trace_id}-bbbbbbbbbbbbbbbb-01");
let resp = fs
.server
.post("/v1/chat/completions")
.add_header("authorization", format!("Bearer {}", fs.brvk_secret))
.add_header("accept", "text/event-stream")
.add_header("traceparent", inbound_traceparent)
.json(&clean_body(true))
.await;
resp.assert_status_ok();
let _ = resp.text();
wait_for_metering_row(&fs.db, "fullstack-key").await;
wait_for_otlp(&fs.otlp_collector).await;
fs.observe.shutdown().await;
let spans = collect_exported_trace_spans(&fs.otlp_collector).await;
assert!(
!spans.is_empty(),
"collector must have received at least one trace span"
);
use opentelemetry_proto::tonic::trace::v1::span::SpanKind as ProtoKind;
let root_chat = spans
.iter()
.find(|s| s.name == "chat test-model" && s.kind == ProtoKind::Internal as i32)
.expect("root chat INTERNAL span present");
let hop_chat = spans
.iter()
.find(|s| s.name == "chat test-model" && s.kind == ProtoKind::Client as i32)
.expect("per-hop chat CLIENT span present");
let route = spans
.iter()
.find(|s| s.name == "route")
.expect("route INTERNAL span present");
let settle = spans
.iter()
.find(|s| s.name == "settle")
.expect("settle INTERNAL span present");
assert_eq!(
hex::encode(&root_chat.trace_id),
inbound_trace_id,
"root chat must continue the inbound trace_id"
);
assert_eq!(route.parent_span_id, root_chat.span_id);
assert_eq!(hop_chat.parent_span_id, root_chat.span_id);
assert_eq!(settle.parent_span_id, root_chat.span_id);
assert_eq!(route.trace_id, root_chat.trace_id);
assert_eq!(hop_chat.trace_id, root_chat.trace_id);
assert_eq!(settle.trace_id, root_chat.trace_id);
assert_eq!(
span_str_attr(hop_chat, "gen_ai.operation.name"),
Some("chat")
);
assert_eq!(
span_str_attr(hop_chat, "gen_ai.provider.name"),
Some("mock"),
"the hop carries the provider id under the current spec name"
);
assert_eq!(
span_str_attr(hop_chat, "gen_ai.request.model"),
Some("test-model")
);
assert!(
span_str_attr(hop_chat, "server.address").is_some(),
"hop must carry `server.address` parsed from the target's api_base"
);
fs.teardown().await;
}