use axum::{
extract::{Path, State},
http::{HeaderMap, StatusCode},
response::sse::{Event, Sse},
routing::{get, post},
Json, Router,
};
use futures::{stream::BoxStream, StreamExt};
use parking_lot::Mutex;
use serde::{Deserialize, Serialize};
use std::{collections::HashMap, convert::Infallible, net::SocketAddr, sync::Arc};
use tokio::{net::TcpListener, sync::Notify};
#[derive(Clone, Debug, Deserialize, Serialize)]
#[serde(tag = "type", rename_all = "snake_case")]
pub enum MockResponse {
Text {
content: String,
},
ToolCall {
name: String,
input: serde_json::Value,
},
ToolCalls {
calls: Vec<(String, serde_json::Value)>,
},
Error {
status: u16,
message: String,
},
TextStream {
chunks: Vec<String>,
},
ToolCallStream {
name: String,
id: String,
input: serde_json::Value,
},
Thinking {
text: String,
signature: String,
},
Truncated {
content: String,
},
Reasoning {
reasoning: String,
content: String,
},
CutStream {
chunks: Vec<String>,
after: usize,
},
CutToolCallStream {
name: String,
id: String,
partial_input_json: String,
},
}
pub(crate) struct QueueEntry {
pub(crate) response: MockResponse,
pub(crate) reached: Option<Arc<Notify>>,
pub(crate) gate: Option<Arc<Notify>>,
pub(crate) delay: Option<std::time::Duration>,
}
impl QueueEntry {
fn immediate(response: MockResponse) -> Self {
Self {
response,
reached: None,
gate: None,
delay: None,
}
}
}
pub struct BlockHandle {
gate: Arc<Notify>,
reached: Arc<Notify>,
}
impl BlockHandle {
pub async fn wait_until_received(&self) {
self.reached.notified().await;
}
pub fn release(&self) {
self.gate.notify_one();
}
}
#[derive(Clone, Debug, Deserialize, Serialize)]
pub struct Scenario {
#[serde(default)]
pub description: String,
pub responses: Vec<MockResponse>,
}
struct ScenarioState {
responses: Vec<MockResponse>,
cursor: usize,
}
impl ScenarioState {
fn from_scenario(s: &Scenario) -> Self {
Self {
responses: s.responses.clone(),
cursor: 0,
}
}
fn next_response(&mut self) -> Option<MockResponse> {
let resp = self.responses.get(self.cursor)?.clone();
self.cursor += 1;
Some(resp)
}
}
#[derive(Clone, Debug, Deserialize, Serialize)]
pub struct ScenarioConfig {
pub scenarios: HashMap<String, Scenario>,
}
#[derive(Default)]
pub(crate) struct MockState {
queue: Mutex<Vec<QueueEntry>>,
scenarios: Mutex<HashMap<String, Scenario>>,
session_bindings: Mutex<HashMap<String, String>>,
session_states: Mutex<HashMap<String, ScenarioState>>,
captured: Mutex<Vec<serde_json::Value>>,
hold_next: Mutex<Option<std::time::Duration>>,
}
impl MockState {
pub(crate) fn dequeue_entry(&self) -> Option<QueueEntry> {
let mut entry = {
let mut q = self.queue.lock();
(!q.is_empty()).then(|| q.remove(0))
}?;
if let Some(hold) = self.hold_next.lock().take() {
if entry.delay.is_none() {
entry.delay = Some(hold);
}
}
Some(entry)
}
pub(crate) fn capture(&self, body: serde_json::Value) {
self.captured.lock().push(body);
}
}
#[derive(Default)]
pub struct MockLlmServerBuilder {
responses: Vec<MockResponse>,
scenarios: HashMap<String, Scenario>,
bind_all: bool,
port: Option<u16>,
}
impl MockLlmServerBuilder {
#[must_use]
pub fn response(mut self, text: impl Into<String>) -> Self {
self.responses.push(MockResponse::Text {
content: text.into(),
});
self
}
#[must_use]
pub fn tool_call(mut self, name: impl Into<String>, input: serde_json::Value) -> Self {
self.responses.push(MockResponse::ToolCall {
name: name.into(),
input,
});
self
}
#[must_use]
pub fn error(mut self, status: u16, message: impl Into<String>) -> Self {
self.responses.push(MockResponse::Error {
status,
message: message.into(),
});
self
}
#[must_use]
pub fn response_stream(mut self, chunks: impl IntoIterator<Item = impl Into<String>>) -> Self {
self.responses.push(MockResponse::TextStream {
chunks: chunks.into_iter().map(Into::into).collect(),
});
self
}
#[must_use]
pub fn tool_call_stream(mut self, name: impl Into<String>, input: serde_json::Value) -> Self {
self.responses.push(MockResponse::ToolCallStream {
name: name.into(),
id: format!("toolu_{}", uuid::Uuid::new_v4()),
input,
});
self
}
#[must_use]
pub fn thinking(mut self, text: impl Into<String>, signature: impl Into<String>) -> Self {
self.responses.push(MockResponse::Thinking {
text: text.into(),
signature: signature.into(),
});
self
}
#[must_use]
pub fn with_scenarios(mut self, config: ScenarioConfig) -> Self {
self.scenarios = config.scenarios;
self
}
#[must_use]
pub fn bind_all_interfaces(mut self) -> Self {
self.bind_all = true;
self
}
#[must_use]
pub fn port(mut self, port: u16) -> Self {
self.port = Some(port);
self
}
pub async fn build(self) -> MockLlmServer {
let queue = self
.responses
.into_iter()
.map(QueueEntry::immediate)
.collect();
let state = Arc::new(MockState {
queue: Mutex::new(queue),
scenarios: Mutex::new(self.scenarios),
session_bindings: Mutex::new(HashMap::new()),
session_states: Mutex::new(HashMap::new()),
captured: Mutex::new(Vec::new()),
hold_next: Mutex::new(None),
});
let app = Router::new()
.route(
"/v1/messages",
post(crate::mock::anthropic::handle_messages),
)
.route(
"/v1/chat/completions",
post(crate::mock::openai::handle_chat_completions),
)
.route("/responses", post(crate::mock::responses::handle_responses))
.route(
"/v1/responses",
post(crate::mock::responses::handle_responses),
)
.route("/queue", post(handle_queue))
.route("/received", get(handle_received))
.route("/reset", post(handle_reset))
.route("/scenarios/load", post(handle_load_scenarios))
.route("/scenarios", get(handle_list_scenarios))
.route(
"/scenarios/{name}/register/{session_id}",
post(handle_register_session),
)
.with_state(state.clone());
let port = self.port.unwrap_or(0);
let bind = if self.bind_all {
format!("0.0.0.0:{port}")
} else {
format!("127.0.0.1:{port}")
};
let listener = TcpListener::bind(&bind).await.unwrap();
let addr = listener.local_addr().unwrap();
let handle = tokio::spawn(async move { axum::serve(listener, app).await.unwrap() });
MockLlmServer {
addr,
_handle: handle,
state,
}
}
}
pub struct MockLlmServer {
addr: SocketAddr,
_handle: tokio::task::JoinHandle<()>,
state: Arc<MockState>,
}
impl MockLlmServer {
#[must_use]
pub fn builder() -> MockLlmServerBuilder {
MockLlmServerBuilder::default()
}
#[must_use]
pub fn url(&self) -> String {
format!("http://{}", self.addr)
}
#[must_use]
pub fn port(&self) -> u16 {
self.addr.port()
}
pub fn queued_count(&self) -> usize {
self.state.queue.lock().len()
}
pub fn queue_response(&self, text: impl Into<String>) {
self.state
.queue
.lock()
.push(QueueEntry::immediate(MockResponse::Text {
content: text.into(),
}));
}
pub fn queue_tool_calls(&self, calls: Vec<(String, serde_json::Value)>) {
self.state
.queue
.lock()
.push(QueueEntry::immediate(MockResponse::ToolCalls { calls }));
}
pub fn queue_tool_call(&self, name: impl Into<String>, input: serde_json::Value) {
self.state
.queue
.lock()
.push(QueueEntry::immediate(MockResponse::ToolCall {
name: name.into(),
input,
}));
}
pub fn queue_error(&self, status: u16, message: impl Into<String>) {
self.state
.queue
.lock()
.push(QueueEntry::immediate(MockResponse::Error {
status,
message: message.into(),
}));
}
pub fn queue_truncated(&self, content: impl Into<String>) {
self.state
.queue
.lock()
.push(QueueEntry::immediate(MockResponse::Truncated {
content: content.into(),
}));
}
pub fn queue_delayed(&self, text: impl Into<String>, delay: std::time::Duration) {
self.state.queue.lock().push(QueueEntry {
response: MockResponse::Text {
content: text.into(),
},
reached: None,
gate: None,
delay: Some(delay),
});
}
pub fn queue_cut_stream(
&self,
chunks: impl IntoIterator<Item = impl Into<String>>,
after: usize,
) {
self.state
.queue
.lock()
.push(QueueEntry::immediate(MockResponse::CutStream {
chunks: chunks.into_iter().map(Into::into).collect(),
after,
}));
}
pub fn queue_cut_tool_call(
&self,
name: impl Into<String>,
id: impl Into<String>,
partial_input_json: impl Into<String>,
) {
self.state
.queue
.lock()
.push(QueueEntry::immediate(MockResponse::CutToolCallStream {
name: name.into(),
id: id.into(),
partial_input_json: partial_input_json.into(),
}));
}
pub fn queue_reasoning(&self, reasoning: impl Into<String>, content: impl Into<String>) {
self.state
.queue
.lock()
.push(QueueEntry::immediate(MockResponse::Reasoning {
reasoning: reasoning.into(),
content: content.into(),
}));
}
pub fn blocking_response(&self, text: impl Into<String>) -> BlockHandle {
let gate = Arc::new(Notify::new());
let reached = Arc::new(Notify::new());
self.state.queue.lock().push(QueueEntry {
response: MockResponse::Text {
content: text.into(),
},
reached: Some(Arc::clone(&reached)),
gate: Some(Arc::clone(&gate)),
delay: None,
});
BlockHandle { gate, reached }
}
pub fn load_scenarios(&self, config: ScenarioConfig) {
self.state.scenarios.lock().extend(config.scenarios);
}
pub fn register_session(
&self,
session_id: impl Into<String>,
scenario_name: impl Into<String>,
) {
let session_id = session_id.into();
let scenario_name = scenario_name.into();
let scenarios = self.state.scenarios.lock();
if let Some(scenario) = scenarios.get(&scenario_name) {
let state = ScenarioState::from_scenario(scenario);
drop(scenarios);
self.state
.session_states
.lock()
.insert(session_id.clone(), state);
self.state
.session_bindings
.lock()
.insert(session_id, scenario_name);
}
}
}
#[derive(Serialize)]
struct StatusResponse {
status: String,
#[serde(skip_serializing_if = "Option::is_none")]
message: Option<String>,
}
#[derive(Serialize)]
struct ScenariosListResponse {
scenarios: Vec<String>,
}
pub(crate) enum ResponseKind {
Json(axum::Json<serde_json::Value>),
Sse(Sse<BoxStream<'static, Result<Event, Infallible>>>),
HttpError(StatusCode, axum::Json<serde_json::Value>),
}
impl axum::response::IntoResponse for ResponseKind {
fn into_response(self) -> axum::response::Response {
match self {
ResponseKind::Json(j) => j.into_response(),
ResponseKind::Sse(s) => s.into_response(),
ResponseKind::HttpError(status, body) => (status, body).into_response(),
}
}
}
pub(crate) async fn handle_messages(
State(state): State<Arc<MockState>>,
headers: HeaderMap,
Json(req): Json<serde_json::Value>,
) -> ResponseKind {
state.capture(req.clone());
let entry = if let Some(sid) = headers.get("X-Session-Id").and_then(|v| v.to_str().ok()) {
let mut ss = state.session_states.lock();
if let Some(scenario) = ss.get_mut(sid) {
scenario.next_response().map(QueueEntry::immediate)
} else {
state.dequeue_entry()
}
} else {
state.dequeue_entry()
};
if let Some(e) = &entry {
if let Some(r) = &e.reached {
r.notify_one();
}
if let Some(g) = &e.gate {
g.notified().await;
}
if let Some(d) = e.delay {
tokio::time::sleep(d).await;
}
}
let response = entry.map(|e| e.response);
let is_stream = req
.get("stream")
.and_then(serde_json::Value::as_bool)
.unwrap_or(false);
match response {
Some(MockResponse::TextStream { chunks }) => sse_from_pairs(text_stream_sse(&chunks)),
Some(MockResponse::ToolCallStream { name, id, input }) => {
sse_from_pairs(tool_call_stream_sse(&name, &id, &input))
}
Some(MockResponse::CutStream { chunks, after }) => {
sse_from_pairs(cut_text_stream_sse(&chunks, after))
}
Some(MockResponse::CutToolCallStream {
name,
id,
partial_input_json,
}) => sse_from_pairs(cut_tool_call_stream_sse(&name, &id, &partial_input_json)),
other => {
let resp = other;
if is_stream {
let msg_id = format!("msg_{}", uuid::Uuid::new_v4());
let tool_id = format!("toolu_{}", uuid::Uuid::new_v4());
let pairs = match resp {
Some(MockResponse::Text { content }) => text_sse(&msg_id, &content),
Some(MockResponse::ToolCall { name, input }) => {
tool_sse(&msg_id, &tool_id, &name, &input)
}
Some(MockResponse::ToolCalls { calls }) => tools_sse(&msg_id, &calls),
Some(MockResponse::Thinking { text, signature }) => {
thinking_sse(&msg_id, &text, &signature)
}
Some(MockResponse::Truncated { content }) => truncated_sse(&msg_id, &content),
Some(MockResponse::Reasoning { content, .. }) => text_sse(&msg_id, &content),
Some(MockResponse::Error { status, message }) => {
let etype = match status {
429 => "rate_limit_error",
529 => "overloaded_error",
_ => "invalid_request_error",
};
vec![(
"error".into(),
serde_json::json!({ "type": etype, "message": message }).to_string(),
)]
}
None => text_sse(&msg_id, "No mock response queued"),
Some(
MockResponse::TextStream { .. }
| MockResponse::ToolCallStream { .. }
| MockResponse::CutStream { .. }
| MockResponse::CutToolCallStream { .. },
) => {
unreachable!()
}
};
sse_from_pairs(pairs)
} else {
match resp {
Some(MockResponse::Text { content }) => {
ResponseKind::Json(axum::Json(text_json(&content)))
}
Some(MockResponse::ToolCall { name, input }) => {
ResponseKind::Json(axum::Json(tool_json(&name, &input)))
}
Some(MockResponse::ToolCalls { calls }) => {
ResponseKind::Json(axum::Json(tools_json(&calls)))
}
Some(MockResponse::Thinking { text, signature }) => {
ResponseKind::Json(axum::Json(thinking_json(&text, &signature)))
}
Some(MockResponse::Truncated { content }) => {
ResponseKind::Json(axum::Json(text_json(&content)))
}
Some(MockResponse::Reasoning { content, .. }) => {
ResponseKind::Json(axum::Json(text_json(&content)))
}
Some(MockResponse::Error { status, message }) => {
let code = StatusCode::from_u16(status)
.unwrap_or(StatusCode::INTERNAL_SERVER_ERROR);
ResponseKind::HttpError(code, axum::Json(error_json(&message)))
}
None => ResponseKind::Json(axum::Json(text_json("No mock response queued"))),
Some(
MockResponse::TextStream { .. }
| MockResponse::ToolCallStream { .. }
| MockResponse::CutStream { .. }
| MockResponse::CutToolCallStream { .. },
) => {
unreachable!()
}
}
}
}
}
}
async fn handle_queue(
State(state): State<Arc<MockState>>,
Json(body): Json<serde_json::Value>,
) -> Json<StatusResponse> {
let delay = body
.get("delayMs")
.and_then(serde_json::Value::as_u64)
.map(std::time::Duration::from_millis);
let response: MockResponse = match serde_json::from_value(body) {
Ok(r) => r,
Err(e) => {
return Json(StatusResponse {
status: "error".into(),
message: Some(e.to_string()),
});
}
};
state.queue.lock().push(QueueEntry {
response,
reached: None,
gate: None,
delay,
});
Json(StatusResponse {
status: "queued".into(),
message: None,
})
}
async fn handle_reset(
State(state): State<Arc<MockState>>,
body: Option<Json<serde_json::Value>>,
) -> Json<StatusResponse> {
state.queue.lock().clear();
state.session_states.lock().clear();
state.session_bindings.lock().clear();
state.captured.lock().clear();
*state.hold_next.lock() = body
.and_then(|Json(b)| b.get("holdFirstMs").and_then(serde_json::Value::as_u64))
.map(std::time::Duration::from_millis);
Json(StatusResponse {
status: "reset".into(),
message: None,
})
}
async fn handle_received(State(state): State<Arc<MockState>>) -> Json<Vec<serde_json::Value>> {
let mut bodies = state.captured.lock().clone();
bodies.reverse();
Json(bodies)
}
async fn handle_load_scenarios(
State(state): State<Arc<MockState>>,
Json(config): Json<ScenarioConfig>,
) -> Json<StatusResponse> {
let count = config.scenarios.len();
state.scenarios.lock().extend(config.scenarios);
Json(StatusResponse {
status: "loaded".into(),
message: Some(format!("{count} scenarios loaded")),
})
}
async fn handle_list_scenarios(State(state): State<Arc<MockState>>) -> Json<ScenariosListResponse> {
let scenarios = state.scenarios.lock();
Json(ScenariosListResponse {
scenarios: scenarios.keys().cloned().collect(),
})
}
async fn handle_register_session(
State(state): State<Arc<MockState>>,
Path((scenario_name, session_id)): Path<(String, String)>,
) -> Result<Json<StatusResponse>, (StatusCode, Json<StatusResponse>)> {
let scenarios = state.scenarios.lock();
if let Some(scenario) = scenarios.get(&scenario_name) {
let scenario_state = ScenarioState::from_scenario(scenario);
drop(scenarios);
state
.session_states
.lock()
.insert(session_id.clone(), scenario_state);
state
.session_bindings
.lock()
.insert(session_id.clone(), scenario_name.clone());
Ok(Json(StatusResponse {
status: "registered".into(),
message: Some(format!(
"Session {session_id} bound to scenario {scenario_name}"
)),
}))
} else {
Err((
StatusCode::NOT_FOUND,
Json(StatusResponse {
status: "error".into(),
message: Some(format!("Scenario '{scenario_name}' not found")),
}),
))
}
}
pub(crate) fn sse_from_pairs(pairs: Vec<(String, String)>) -> ResponseKind {
let events: Vec<Result<Event, Infallible>> = pairs
.into_iter()
.map(|(t, d)| Ok(Event::default().event(t).data(d)))
.collect();
ResponseKind::Sse(Sse::new(futures::stream::iter(events).boxed()))
}
fn text_sse(msg_id: &str, text: &str) -> Vec<(String, String)> {
let tokens = u32::try_from(text.len() / 4).unwrap_or(u32::MAX);
vec![
(
"message_start".into(),
serde_json::json!({"type":"message_start","message":{"id":msg_id,"type":"message","role":"assistant","content":[],"model":"mock-model","stop_reason":null,"usage":{"input_tokens":10,"output_tokens":1}}}).to_string(),
),
(
"content_block_start".into(),
serde_json::json!({"type":"content_block_start","index":0,"content_block":{"type":"text","text":""}}).to_string(),
),
(
"content_block_delta".into(),
serde_json::json!({"type":"content_block_delta","index":0,"delta":{"type":"text_delta","text":text}}).to_string(),
),
(
"content_block_stop".into(),
serde_json::json!({"type":"content_block_stop","index":0}).to_string(),
),
(
"message_delta".into(),
serde_json::json!({"type":"message_delta","delta":{"stop_reason":"end_turn","stop_sequence":null},"usage":{"output_tokens":tokens}}).to_string(),
),
(
"message_stop".into(),
serde_json::json!({"type":"message_stop"}).to_string(),
),
]
}
fn truncated_sse(msg_id: &str, text: &str) -> Vec<(String, String)> {
let mut pairs = text_sse(msg_id, text);
for (event, data) in &mut pairs {
if event == "message_delta" {
*data = serde_json::json!({
"type": "message_delta",
"delta": {"stop_reason": "max_tokens", "stop_sequence": null},
"usage": {"output_tokens": 5}
})
.to_string();
}
}
pairs
}
fn tool_sse(
msg_id: &str,
tool_id: &str,
name: &str,
input: &serde_json::Value,
) -> Vec<(String, String)> {
let input_str = serde_json::to_string(input).unwrap_or_default();
vec![
(
"message_start".into(),
serde_json::json!({"type":"message_start","message":{"id":msg_id,"type":"message","role":"assistant","content":[],"model":"mock-model","stop_reason":null,"usage":{"input_tokens":10,"output_tokens":1}}}).to_string(),
),
(
"content_block_start".into(),
serde_json::json!({"type":"content_block_start","index":0,"content_block":{"type":"tool_use","id":tool_id,"name":name,"input":{}}}).to_string(),
),
(
"content_block_delta".into(),
serde_json::json!({"type":"content_block_delta","index":0,"delta":{"type":"input_json_delta","partial_json":input_str}}).to_string(),
),
(
"content_block_stop".into(),
serde_json::json!({"type":"content_block_stop","index":0}).to_string(),
),
(
"message_delta".into(),
serde_json::json!({"type":"message_delta","delta":{"stop_reason":"tool_use","stop_sequence":null},"usage":{"output_tokens":20}}).to_string(),
),
(
"message_stop".into(),
serde_json::json!({"type":"message_stop"}).to_string(),
),
]
}
fn tools_sse(msg_id: &str, calls: &[(String, serde_json::Value)]) -> Vec<(String, String)> {
let mut out = vec![(
"message_start".to_string(),
serde_json::json!({"type":"message_start","message":{"id":msg_id,"type":"message","role":"assistant","content":[],"model":"mock-model","stop_reason":null,"usage":{"input_tokens":10,"output_tokens":1}}}).to_string(),
)];
for (index, (name, input)) in calls.iter().enumerate() {
let tool_id = format!("toolu_{}", uuid::Uuid::new_v4());
let input_str = serde_json::to_string(input).unwrap_or_default();
out.push((
"content_block_start".to_string(),
serde_json::json!({"type":"content_block_start","index":index,"content_block":{"type":"tool_use","id":tool_id,"name":name,"input":{}}}).to_string(),
));
out.push((
"content_block_delta".to_string(),
serde_json::json!({"type":"content_block_delta","index":index,"delta":{"type":"input_json_delta","partial_json":input_str}}).to_string(),
));
out.push((
"content_block_stop".to_string(),
serde_json::json!({"type":"content_block_stop","index":index}).to_string(),
));
}
out.push((
"message_delta".to_string(),
serde_json::json!({"type":"message_delta","delta":{"stop_reason":"tool_use","stop_sequence":null},"usage":{"output_tokens":20}}).to_string(),
));
out.push((
"message_stop".to_string(),
serde_json::json!({"type":"message_stop"}).to_string(),
));
out
}
fn thinking_sse(msg_id: &str, text: &str, signature: &str) -> Vec<(String, String)> {
vec![
(
"message_start".into(),
serde_json::json!({"type":"message_start","message":{"id":msg_id,"type":"message","role":"assistant","content":[],"model":"mock-model","stop_reason":null,"usage":{"input_tokens":10,"output_tokens":1}}}).to_string(),
),
(
"content_block_start".into(),
serde_json::json!({"type":"content_block_start","index":0,"content_block":{"type":"thinking","thinking":"","signature":""}}).to_string(),
),
(
"content_block_delta".into(),
serde_json::json!({"type":"content_block_delta","index":0,"delta":{"type":"thinking_delta","thinking":text}}).to_string(),
),
(
"content_block_delta".into(),
serde_json::json!({"type":"content_block_delta","index":0,"delta":{"type":"signature_delta","signature":signature}}).to_string(),
),
(
"content_block_stop".into(),
serde_json::json!({"type":"content_block_stop","index":0}).to_string(),
),
(
"message_delta".into(),
serde_json::json!({"type":"message_delta","delta":{"stop_reason":"end_turn","stop_sequence":null},"usage":{"output_tokens":5}}).to_string(),
),
(
"message_stop".into(),
serde_json::json!({"type":"message_stop"}).to_string(),
),
]
}
fn text_stream_sse(chunks: &[String]) -> Vec<(String, String)> {
let msg_id = format!("msg_{}", uuid::Uuid::new_v4());
let mut events = vec![
(
"message_start".into(),
serde_json::json!({"type":"message_start","message":{"id":msg_id,"type":"message","role":"assistant","content":[],"model":"mock-model","stop_reason":null,"usage":{"input_tokens":10,"output_tokens":0}}}).to_string(),
),
(
"content_block_start".into(),
serde_json::json!({"type":"content_block_start","index":0,"content_block":{"type":"text","text":""}}).to_string(),
),
];
for chunk in chunks {
events.push((
"content_block_delta".into(),
serde_json::json!({"type":"content_block_delta","index":0,"delta":{"type":"text_delta","text":chunk}}).to_string(),
));
}
events.push((
"content_block_stop".into(),
serde_json::json!({"type":"content_block_stop","index":0}).to_string(),
));
events.push((
"message_delta".into(),
serde_json::json!({"type":"message_delta","delta":{"stop_reason":"end_turn","stop_sequence":null},"usage":{"output_tokens":chunks.len()}}).to_string(),
));
events.push((
"message_stop".into(),
serde_json::json!({"type":"message_stop"}).to_string(),
));
events
}
fn cut_text_stream_sse(chunks: &[String], after: usize) -> Vec<(String, String)> {
let msg_id = format!("msg_{}", uuid::Uuid::new_v4());
let mut events = vec![
(
"message_start".into(),
serde_json::json!({"type":"message_start","message":{"id":msg_id,"type":"message","role":"assistant","content":[],"model":"mock-model","stop_reason":null,"usage":{"input_tokens":10,"output_tokens":0}}}).to_string(),
),
(
"content_block_start".into(),
serde_json::json!({"type":"content_block_start","index":0,"content_block":{"type":"text","text":""}}).to_string(),
),
];
for chunk in chunks.iter().take(after) {
events.push((
"content_block_delta".into(),
serde_json::json!({"type":"content_block_delta","index":0,"delta":{"type":"text_delta","text":chunk}}).to_string(),
));
}
events
}
fn cut_tool_call_stream_sse(name: &str, id: &str, partial: &str) -> Vec<(String, String)> {
let msg_id = format!("msg_{}", uuid::Uuid::new_v4());
vec![
(
"message_start".into(),
serde_json::json!({"type":"message_start","message":{"id":msg_id,"type":"message","role":"assistant","content":[],"model":"mock-model","stop_reason":null,"usage":{"input_tokens":10,"output_tokens":0}}}).to_string(),
),
(
"content_block_start".into(),
serde_json::json!({"type":"content_block_start","index":0,"content_block":{"type":"tool_use","id":id,"name":name,"input":{}}}).to_string(),
),
(
"content_block_delta".into(),
serde_json::json!({"type":"content_block_delta","index":0,"delta":{"type":"input_json_delta","partial_json":partial}}).to_string(),
),
]
}
fn tool_call_stream_sse(name: &str, id: &str, input: &serde_json::Value) -> Vec<(String, String)> {
let msg_id = format!("msg_{}", uuid::Uuid::new_v4());
let input_str = input.to_string();
let fragments: Vec<String> = input_str
.as_bytes()
.chunks(10)
.map(|c| String::from_utf8_lossy(c).to_string())
.collect();
let mut events = vec![
(
"message_start".into(),
serde_json::json!({"type":"message_start","message":{"id":msg_id,"type":"message","role":"assistant","content":[],"model":"mock-model","stop_reason":null,"usage":{"input_tokens":10,"output_tokens":0}}}).to_string(),
),
(
"content_block_start".into(),
serde_json::json!({"type":"content_block_start","index":0,"content_block":{"type":"tool_use","id":id,"name":name,"input":{}}}).to_string(),
),
];
for frag in &fragments {
events.push((
"content_block_delta".into(),
serde_json::json!({"type":"content_block_delta","index":0,"delta":{"type":"input_json_delta","partial_json":frag}}).to_string(),
));
}
events.push((
"content_block_stop".into(),
serde_json::json!({"type":"content_block_stop","index":0}).to_string(),
));
events.push((
"message_delta".into(),
serde_json::json!({"type":"message_delta","delta":{"stop_reason":"tool_use","stop_sequence":null},"usage":{"output_tokens":20}}).to_string(),
));
events.push((
"message_stop".into(),
serde_json::json!({"type":"message_stop"}).to_string(),
));
events
}
fn text_json(text: &str) -> serde_json::Value {
serde_json::json!({
"type": "message",
"id": format!("msg_{}", uuid::Uuid::new_v4()),
"role": "assistant",
"content": [{"type": "text", "text": text}],
"model": "mock-model",
"stop_reason": "end_turn",
"usage": {"input_tokens": 10, "output_tokens": text.len() / 4}
})
}
fn tool_json(name: &str, input: &serde_json::Value) -> serde_json::Value {
serde_json::json!({
"type": "message",
"id": format!("msg_{}", uuid::Uuid::new_v4()),
"role": "assistant",
"content": [{"type": "tool_use", "id": format!("toolu_{}", uuid::Uuid::new_v4()), "name": name, "input": input}],
"model": "mock-model",
"stop_reason": "tool_use",
"usage": {"input_tokens": 10, "output_tokens": 20}
})
}
fn tools_json(calls: &[(String, serde_json::Value)]) -> serde_json::Value {
let content: Vec<serde_json::Value> = calls
.iter()
.map(|(name, input)| {
serde_json::json!({"type": "tool_use", "id": format!("toolu_{}", uuid::Uuid::new_v4()), "name": name, "input": input})
})
.collect();
serde_json::json!({
"type": "message",
"id": format!("msg_{}", uuid::Uuid::new_v4()),
"role": "assistant",
"content": content,
"model": "mock-model",
"stop_reason": "tool_use",
"usage": {"input_tokens": 10, "output_tokens": 20}
})
}
fn thinking_json(text: &str, signature: &str) -> serde_json::Value {
serde_json::json!({
"type": "message",
"id": format!("msg_{}", uuid::Uuid::new_v4()),
"role": "assistant",
"content": [{"type": "thinking", "thinking": text, "signature": signature}],
"model": "mock-model",
"stop_reason": "end_turn",
"usage": {"input_tokens": 10, "output_tokens": 5}
})
}
fn error_json(message: &str) -> serde_json::Value {
serde_json::json!({"type": "error", "error": {"type": "api_error", "message": message}})
}