use std::sync::{Arc, Mutex};
use axiom_engine::config::AxiomConfig;
use axiom_engine::inference::InferencePipeline;
use axiom_engine::openai_forwarder::OpenAiForwarder;
use axiom_engine::server::{create_router, AppState};
use axum::body::{to_bytes, Body};
use axum::extract::State;
use axum::http::{header, HeaderMap, Method, Request, Response, StatusCode};
use axum::routing::post;
use axum::{Json, Router};
use candle_core::Device;
use serde_json::{json, Value};
use tower::ServiceExt;
#[derive(Clone, Default)]
struct Capture {
requests: Arc<Mutex<Vec<(Option<String>, Value)>>>,
}
fn build_pipeline() -> InferencePipeline {
InferencePipeline::new(
AxiomConfig {
d_model: 16,
n_layers: 1,
vocab_size: 64,
lr_inner: 1e-3,
norm_eps: 1e-6,
},
Device::Cpu,
)
.expect("tiny pipeline must build")
}
async fn start_mock_responses_upstream(
streaming: bool,
) -> (String, Capture, tokio::task::JoinHandle<()>) {
async fn handler(
State((capture, streaming)): State<(Capture, bool)>,
headers: HeaderMap,
Json(body): Json<Value>,
) -> Response<Body> {
let auth = headers
.get(header::AUTHORIZATION)
.and_then(|value| value.to_str().ok())
.map(str::to_string);
capture.requests.lock().unwrap().push((auth, body));
if streaming {
let events = concat!(
"data: {\"type\":\"response.created\",\"sequence_number\":0,\"response\":{\"id\":\"resp_test\",\"object\":\"response\",\"status\":\"in_progress\"}}\n\n",
"data: {\"type\":\"response.output_text.delta\",\"sequence_number\":1,\"item_id\":\"msg_test\",\"output_index\":0,\"content_index\":0,\"delta\":\"hello\"}\n\n",
"data: {\"type\":\"response.completed\",\"sequence_number\":2,\"response\":{\"id\":\"resp_test\",\"object\":\"response\",\"status\":\"completed\"}}\n\n"
);
Response::builder()
.status(StatusCode::OK)
.header(header::CONTENT_TYPE, "text/event-stream")
.body(Body::from(events))
.unwrap()
} else {
Response::builder()
.status(StatusCode::CREATED)
.header(header::CONTENT_TYPE, "application/json")
.body(Body::from(
json!({
"id": "resp_test",
"object": "response",
"status": "completed",
"output": [{
"id": "msg_test",
"type": "message",
"role": "assistant",
"status": "completed",
"content": [{"type": "output_text", "text": "hello", "annotations": []}]
}]
})
.to_string(),
))
.unwrap()
}
}
let capture = Capture::default();
let app = Router::new()
.route("/v1/responses", post(handler))
.with_state((capture.clone(), streaming));
let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap();
let address = listener.local_addr().unwrap();
let task = tokio::spawn(async move { axum::serve(listener, app).await.unwrap() });
(format!("http://{address}"), capture, task)
}
async fn proxy_app(upstream: String) -> Router {
let pipeline = tokio::task::spawn_blocking(build_pipeline).await.unwrap();
let state = AppState::new(pipeline, "responses-proxy-test".to_string())
.with_openai_forwarder(Some(OpenAiForwarder::new(None, Some(upstream))));
create_router(state)
}
#[tokio::test]
async fn responses_json_relays_auth_request_and_upstream_response() {
let (upstream, capture, task) = start_mock_responses_upstream(false).await;
let app = proxy_app(upstream).await;
let request_body = json!({
"model": "gpt-5.5",
"instructions": "Be concise.",
"input": "hello",
"tools": [{"type": "function", "name": "shell", "parameters": {"type": "object"}}],
"stream": false,
"store": false
});
let response = app
.oneshot(
Request::builder()
.method(Method::POST)
.uri("/v1/responses")
.header(header::CONTENT_TYPE, "application/json")
.header(header::AUTHORIZATION, "Bearer client-secret")
.body(Body::from(request_body.to_string()))
.unwrap(),
)
.await
.unwrap();
assert_eq!(response.status(), StatusCode::CREATED);
assert_eq!(
response.headers().get(header::CONTENT_TYPE).unwrap(),
"application/json"
);
let response_body = to_bytes(response.into_body(), usize::MAX).await.unwrap();
let response_json: Value = serde_json::from_slice(&response_body).unwrap();
assert_eq!(response_json["object"], "response");
assert_eq!(response_json["output"][0]["content"][0]["text"], "hello");
let captured = capture.requests.lock().unwrap();
assert_eq!(captured.len(), 1);
assert_eq!(captured[0].0.as_deref(), Some("Bearer client-secret"));
assert_eq!(captured[0].1, request_body);
task.abort();
}
#[tokio::test]
async fn responses_stream_relays_semantic_sse_events_unchanged() {
let (upstream, capture, task) = start_mock_responses_upstream(true).await;
let app = proxy_app(upstream).await;
let response = app
.oneshot(
Request::builder()
.method(Method::POST)
.uri("/v1/responses")
.header(header::CONTENT_TYPE, "application/json")
.header(header::AUTHORIZATION, "Bearer stream-secret")
.body(Body::from(
json!({"model": "gpt-5.5", "input": "hello", "stream": true}).to_string(),
))
.unwrap(),
)
.await
.unwrap();
assert_eq!(response.status(), StatusCode::OK);
assert!(response
.headers()
.get(header::CONTENT_TYPE)
.unwrap()
.to_str()
.unwrap()
.starts_with("text/event-stream"));
let body = to_bytes(response.into_body(), usize::MAX).await.unwrap();
let text = String::from_utf8(body.to_vec()).unwrap();
assert!(text.contains("response.output_text.delta"));
assert!(text.contains("response.completed"));
assert!(!text.contains("chat.completion.chunk"));
assert_eq!(capture.requests.lock().unwrap().len(), 1);
task.abort();
}