1442 lines
51 KiB
Rust
1442 lines
51 KiB
Rust
//! Integration tests for the sampling client with mock HTTP servers.
|
|
//!
|
|
//! These tests verify the full streaming flow for both:
|
|
//! - Chat Completions API (`/v1/chat/completions`)
|
|
//! - Responses API (`/v1/responses`)
|
|
//!
|
|
//! Each test spawns a temporary mock server that returns SSE streams,
|
|
//! allowing us to test the client without real API credentials.
|
|
|
|
use futures_util::stream::StreamExt;
|
|
use reqwest::StatusCode;
|
|
use serde_json::{Value, json};
|
|
|
|
use xai_grok_shell::sampling::{
|
|
ApiBackend, Client, ConversationItem, ConversationRequest, ConversationToolChoice,
|
|
SamplingError, ToolCall, ToolSpec, rs,
|
|
};
|
|
use xai_grok_shell::session::storage::JsonlStorageAdapter;
|
|
use xai_grok_test_support::sse::responses_api_reasoning_and_text_events;
|
|
use xai_grok_test_support::{MockInferenceServer, ScriptedResponse, SseEvent};
|
|
|
|
mod common;
|
|
|
|
use common::{create_test_client, create_test_client_with_extra_headers, test_sampler_config};
|
|
|
|
// ============================================================================
|
|
// Mock Response Generators
|
|
// ============================================================================
|
|
|
|
/// Generate a Chat Completions SSE stream with tool calls.
|
|
fn chat_completion_tool_call_stream(tool_calls: Vec<Value>, model: &str) -> Vec<SseEvent> {
|
|
let mut events = Vec::new();
|
|
|
|
// First chunk with role
|
|
let first_chunk = json!({
|
|
"id": "chatcmpl-test123",
|
|
"object": "chat.completion.chunk",
|
|
"created": 1234567890,
|
|
"model": model,
|
|
"choices": [{
|
|
"index": 0,
|
|
"delta": {
|
|
"role": "assistant",
|
|
"content": null,
|
|
"tool_calls": tool_calls
|
|
},
|
|
"finish_reason": null
|
|
}]
|
|
});
|
|
events.push(SseEvent::data(first_chunk.to_string()));
|
|
|
|
// Final chunk with finish reason
|
|
let final_chunk = json!({
|
|
"id": "chatcmpl-test123",
|
|
"object": "chat.completion.chunk",
|
|
"created": 1234567890,
|
|
"model": model,
|
|
"choices": [{
|
|
"index": 0,
|
|
"delta": {},
|
|
"finish_reason": "tool_calls"
|
|
}],
|
|
"usage": {
|
|
"prompt_tokens": 10,
|
|
"completion_tokens": 20,
|
|
"total_tokens": 30
|
|
}
|
|
});
|
|
events.push(SseEvent::data(final_chunk.to_string()));
|
|
events.push(SseEvent::data("[DONE]"));
|
|
|
|
events
|
|
}
|
|
|
|
/// Generate a Chat Completions SSE stream with reasoning content.
|
|
fn chat_completion_with_reasoning_stream(
|
|
reasoning: &str,
|
|
content: &str,
|
|
model: &str,
|
|
) -> Vec<SseEvent> {
|
|
let mut events = Vec::new();
|
|
|
|
// Reasoning chunks
|
|
for word in reasoning.split_whitespace() {
|
|
let chunk = json!({
|
|
"id": "chatcmpl-test123",
|
|
"object": "chat.completion.chunk",
|
|
"created": 1234567890,
|
|
"model": model,
|
|
"choices": [{
|
|
"index": 0,
|
|
"delta": {
|
|
"reasoning_content": format!("{} ", word)
|
|
},
|
|
"finish_reason": null
|
|
}]
|
|
});
|
|
events.push(SseEvent::data(chunk.to_string()));
|
|
}
|
|
|
|
// Content chunks
|
|
for word in content.split_whitespace() {
|
|
let chunk = json!({
|
|
"id": "chatcmpl-test123",
|
|
"object": "chat.completion.chunk",
|
|
"created": 1234567890,
|
|
"model": model,
|
|
"choices": [{
|
|
"index": 0,
|
|
"delta": {
|
|
"content": format!("{} ", word)
|
|
},
|
|
"finish_reason": null
|
|
}]
|
|
});
|
|
events.push(SseEvent::data(chunk.to_string()));
|
|
}
|
|
|
|
// Final chunk
|
|
let final_chunk = json!({
|
|
"id": "chatcmpl-test123",
|
|
"object": "chat.completion.chunk",
|
|
"created": 1234567890,
|
|
"model": model,
|
|
"choices": [{
|
|
"index": 0,
|
|
"delta": {},
|
|
"finish_reason": "stop"
|
|
}],
|
|
"usage": {
|
|
"prompt_tokens": 10,
|
|
"completion_tokens": 20,
|
|
"total_tokens": 30
|
|
}
|
|
});
|
|
events.push(SseEvent::data(final_chunk.to_string()));
|
|
events.push(SseEvent::data("[DONE]"));
|
|
|
|
events
|
|
}
|
|
|
|
/// Generate a Responses API SSE stream with function calls.
|
|
fn responses_api_tool_call_stream(
|
|
call_id: &str,
|
|
name: &str,
|
|
arguments: &str,
|
|
model: &str,
|
|
) -> Vec<SseEvent> {
|
|
let mut events = Vec::new();
|
|
let mut seq = 0;
|
|
|
|
// response.created event
|
|
let created = json!({
|
|
"type": "response.created",
|
|
"sequence_number": seq,
|
|
"response": {
|
|
"id": "resp_test123",
|
|
"object": "response",
|
|
"created_at": 1234567890,
|
|
"model": model,
|
|
"status": "in_progress",
|
|
"output": []
|
|
}
|
|
});
|
|
events.push(SseEvent::data(created.to_string()));
|
|
seq += 1;
|
|
|
|
// Function call arguments delta
|
|
let args_delta = json!({
|
|
"type": "response.function_call_arguments.delta",
|
|
"sequence_number": seq,
|
|
"item_id": call_id,
|
|
"output_index": 0,
|
|
"delta": arguments
|
|
});
|
|
events.push(SseEvent::data(args_delta.to_string()));
|
|
seq += 1;
|
|
|
|
// response.completed event with function call
|
|
let completed = json!({
|
|
"type": "response.completed",
|
|
"sequence_number": seq,
|
|
"response": {
|
|
"id": "resp_test123",
|
|
"object": "response",
|
|
"created_at": 1234567890,
|
|
"model": model,
|
|
"status": "completed",
|
|
"output": [{
|
|
"type": "function_call",
|
|
"call_id": call_id,
|
|
"name": name,
|
|
"arguments": arguments
|
|
}],
|
|
"usage": {
|
|
"input_tokens": 10,
|
|
"output_tokens": 20,
|
|
"total_tokens": 30,
|
|
"input_tokens_details": {
|
|
"cached_tokens": 0
|
|
},
|
|
"output_tokens_details": {
|
|
"reasoning_tokens": 0
|
|
}
|
|
}
|
|
}
|
|
});
|
|
events.push(SseEvent::data(completed.to_string()));
|
|
events.push(SseEvent::data("[DONE]"));
|
|
|
|
events
|
|
}
|
|
|
|
// ============================================================================
|
|
// Chat Completions API Tests
|
|
// ============================================================================
|
|
|
|
#[tokio::test]
|
|
async fn test_chat_completions_streaming_text() {
|
|
let server = MockInferenceServer::start().await.unwrap();
|
|
server.set_response("Hello world from the assistant!");
|
|
let client = create_test_client(&server.url(), ApiBackend::ChatCompletions);
|
|
|
|
let request = ConversationRequest::from_items(vec![
|
|
ConversationItem::system("You are a helpful assistant."),
|
|
ConversationItem::user("Say hello"),
|
|
]);
|
|
|
|
let (mut stream, _metadata) = client.conversation_stream(request).await.unwrap();
|
|
|
|
let mut content = String::new();
|
|
while let Some(chunk_result) = stream.next().await {
|
|
let chunk = chunk_result.unwrap();
|
|
for choice in chunk.choices {
|
|
if let Some(ref text) = choice.delta.content {
|
|
content.push_str(text);
|
|
}
|
|
}
|
|
}
|
|
|
|
assert_eq!(content, "Hello world from the assistant!");
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn test_chat_completions_streaming_tool_calls() {
|
|
let tool_calls = vec![json!({
|
|
"id": "call_abc123",
|
|
"type": "function",
|
|
"function": {
|
|
"name": "read_file",
|
|
"arguments": r#"{"path": "/test.txt"}"#
|
|
}
|
|
})];
|
|
|
|
let server = MockInferenceServer::start().await.unwrap();
|
|
server.enqueue_response(
|
|
"/v1/chat/completions",
|
|
ScriptedResponse::sse(chat_completion_tool_call_stream(tool_calls, "grok-test")),
|
|
);
|
|
let client = create_test_client(&server.url(), ApiBackend::ChatCompletions);
|
|
|
|
let request = ConversationRequest::from_items(vec![ConversationItem::user("Read a file")])
|
|
.with_tools(vec![ToolSpec {
|
|
name: "read_file".to_string(),
|
|
description: Some("Read a file".to_string()),
|
|
parameters: json!({"type": "object"}),
|
|
}]);
|
|
|
|
let (mut stream, _metadata) = client.conversation_stream(request).await.unwrap();
|
|
|
|
let mut tool_calls_received = Vec::new();
|
|
while let Some(chunk_result) = stream.next().await {
|
|
let chunk = chunk_result.unwrap();
|
|
for choice in chunk.choices {
|
|
for tc in choice.delta.tool_calls {
|
|
tool_calls_received.push(tc);
|
|
}
|
|
}
|
|
}
|
|
|
|
assert_eq!(tool_calls_received.len(), 1);
|
|
assert_eq!(
|
|
tool_calls_received[0]
|
|
.function
|
|
.as_ref()
|
|
.and_then(|f| f.name.as_deref()),
|
|
Some("read_file")
|
|
);
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn test_chat_completions_with_reasoning() {
|
|
let server = MockInferenceServer::start().await.unwrap();
|
|
server.enqueue_response(
|
|
"/v1/chat/completions",
|
|
ScriptedResponse::sse(chat_completion_with_reasoning_stream(
|
|
"Let me think about this",
|
|
"The answer is 42",
|
|
"grok-test",
|
|
)),
|
|
);
|
|
let client = create_test_client(&server.url(), ApiBackend::ChatCompletions);
|
|
|
|
let request = ConversationRequest::from_items(vec![ConversationItem::user(
|
|
"What is the meaning of life?",
|
|
)]);
|
|
|
|
let (mut stream, _metadata) = client.conversation_stream(request).await.unwrap();
|
|
|
|
let mut content = String::new();
|
|
let mut reasoning = String::new();
|
|
while let Some(chunk_result) = stream.next().await {
|
|
let chunk = chunk_result.unwrap();
|
|
for choice in chunk.choices {
|
|
if let Some(ref text) = choice.delta.content {
|
|
content.push_str(text);
|
|
}
|
|
if let Some(ref thought) = choice.delta.reasoning_content {
|
|
reasoning.push_str(thought);
|
|
}
|
|
}
|
|
}
|
|
|
|
assert!(reasoning.contains("Let me think"));
|
|
assert!(content.contains("42"));
|
|
}
|
|
|
|
// ============================================================================
|
|
// Reasoning-as-sibling — chat completions, both paths
|
|
// ============================================================================
|
|
|
|
/// All-new path: a chat-completions stream carrying `reasoning_content`
|
|
/// deltas must be collected into a sibling `ConversationItem::Reasoning`
|
|
/// that *precedes* the assistant — the exact shape currently persisted to
|
|
/// chat_history.jsonl. Unlike `test_chat_completions_with_reasoning`
|
|
/// (which only inspects raw SSE deltas), this drives the high-level
|
|
/// `conversation_collect` path so it exercises
|
|
/// stream_chat_completions → collect_response → `ConversationResponse.items`.
|
|
#[tokio::test]
|
|
async fn chat_completions_collect_synthesizes_reasoning_sibling() {
|
|
let server = MockInferenceServer::start().await.unwrap();
|
|
server.enqueue_response(
|
|
"/v1/chat/completions",
|
|
ScriptedResponse::sse(chat_completion_with_reasoning_stream(
|
|
"Let me think about this",
|
|
"The answer is 42",
|
|
"grok-test",
|
|
)),
|
|
);
|
|
let client = create_test_client(&server.url(), ApiBackend::ChatCompletions);
|
|
|
|
let request = ConversationRequest::from_items(vec![ConversationItem::user(
|
|
"What is the meaning of life?",
|
|
)]);
|
|
|
|
let response = client.conversation_collect(request).await.unwrap();
|
|
|
|
// Items must be exactly [Reasoning, Assistant] — the sibling shape.
|
|
assert!(
|
|
matches!(
|
|
response.items.as_slice(),
|
|
[
|
|
ConversationItem::Reasoning(_),
|
|
ConversationItem::Assistant(_)
|
|
]
|
|
),
|
|
"expected [Reasoning, Assistant], got {:?}",
|
|
response.items
|
|
);
|
|
|
|
let reasoning = response
|
|
.reasoning_items()
|
|
.next()
|
|
.expect("a reasoning sibling must be synthesized from reasoning_content");
|
|
let rs::SummaryPart::SummaryText(summary) = &reasoning.summary[0];
|
|
assert!(
|
|
summary.text.contains("Let me think"),
|
|
"reasoning sibling should carry the streamed reasoning_content, got: {:?}",
|
|
summary.text
|
|
);
|
|
|
|
let assistant = response.assistant().expect("assistant item present");
|
|
assert!(
|
|
assistant.content.contains("42"),
|
|
"assistant content should carry the streamed text, got: {:?}",
|
|
assistant.content
|
|
);
|
|
}
|
|
|
|
/// Upgrade path: a legacy chat-completions session on disk —
|
|
/// an assistant carrying inline `reasoning: {text}` — must, when loaded,
|
|
/// reconstruct a sibling Reasoning item, which then folds into
|
|
/// `reasoning_content` on the *correct* assistant message in the outgoing
|
|
/// chat-completions request body. This ties the whole chain together:
|
|
/// read_chat_history_sync (upgrade_legacy_reasoning) → ConversationRequest →
|
|
/// `From<ConversationRequest> for ChatCompletionRequest`
|
|
/// (conversation_to_chat_messages) → wire.
|
|
#[tokio::test]
|
|
async fn chat_completions_upgrade_folds_reconstructed_reasoning_into_request() {
|
|
// 1. Seed a legacy chat-completions chat_history.jsonl (inline reasoning
|
|
// on the assistant — the shape an older binary wrote).
|
|
let dir = tempfile::tempdir().unwrap();
|
|
std::fs::write(
|
|
dir.path().join("chat_history.jsonl"),
|
|
concat!(
|
|
r#"{"type":"system","content":"You are helpful."}"#,
|
|
"\n",
|
|
r#"{"type":"user","content":[{"type":"text","text":"q1"}]}"#,
|
|
"\n",
|
|
r#"{"type":"assistant","content":"a1","reasoning":{"text":"legacy chain of thought","id":"rs_legacy"}}"#,
|
|
"\n",
|
|
),
|
|
)
|
|
.unwrap();
|
|
|
|
// 2. Load through the real adapter — applies the in-memory upgrade.
|
|
let adapter = JsonlStorageAdapter::with_root(dir.path().to_path_buf());
|
|
let mut items = adapter.load_chat_history_from_dir(dir.path()).unwrap();
|
|
assert!(
|
|
items
|
|
.iter()
|
|
.any(|i| matches!(i, ConversationItem::Reasoning(_))),
|
|
"legacy inline reasoning must be reconstructed as a sibling on load, got {:?}",
|
|
items
|
|
);
|
|
|
|
// 3. Continue the conversation and send it over chat-completions,
|
|
// capturing the outgoing request body.
|
|
items.push(ConversationItem::user("q2"));
|
|
|
|
let server = MockInferenceServer::start().await.unwrap();
|
|
server.set_response("ok");
|
|
let client = create_test_client(&server.url(), ApiBackend::ChatCompletions);
|
|
|
|
let _ = client
|
|
.conversation_collect(ConversationRequest::from_items(items))
|
|
.await
|
|
.unwrap();
|
|
|
|
// 4. The reconstructed reasoning must land on the assistant's
|
|
// reasoning_content in the wire request — not be dropped.
|
|
let body = server.request_bodies().pop().unwrap();
|
|
let messages = body.get("messages").unwrap().as_array().unwrap();
|
|
let assistant = messages
|
|
.iter()
|
|
.find(|m| m.get("role").and_then(Value::as_str) == Some("assistant"))
|
|
.expect("assistant message present in request");
|
|
assert_eq!(
|
|
assistant.get("reasoning_content").and_then(Value::as_str),
|
|
Some("legacy chain of thought"),
|
|
"reconstructed reasoning must fold onto the assistant's reasoning_content; \
|
|
full messages: {messages:#?}"
|
|
);
|
|
}
|
|
|
|
/// Upgrade path, grok-build / Responses API: a legacy session whose
|
|
/// assistant carries inline `reasoning: {text, encrypted, id}` must, on
|
|
/// load, reconstruct a sibling Reasoning item that round-trips back to
|
|
/// the Responses API as a **typed** `reasoning` input item — `summary`,
|
|
/// `encrypted_content`, and `id` all preserved — NOT flattened to a
|
|
/// string. This is the byte-stable SGLang-prefix path; it must not go
|
|
/// through `reasoning_item_text`.
|
|
#[tokio::test]
|
|
async fn responses_upgrade_roundtrips_reconstructed_reasoning_as_typed_input() {
|
|
// 1. Seed a legacy grok-build chat_history.jsonl (inline reasoning
|
|
// with encrypted_content + id — the older shape).
|
|
let dir = tempfile::tempdir().unwrap();
|
|
std::fs::write(
|
|
dir.path().join("chat_history.jsonl"),
|
|
concat!(
|
|
r#"{"type":"system","content":"You are helpful."}"#,
|
|
"\n",
|
|
r#"{"type":"user","content":[{"type":"text","text":"q1"}]}"#,
|
|
"\n",
|
|
r#"{"type":"assistant","content":"a1","reasoning":{"text":"legacy grok-build reasoning","encrypted":"ENC_BLOB_xyz","id":"rs_grokbuild_legacy"},"model_id":"grok-build"}"#,
|
|
"\n",
|
|
),
|
|
)
|
|
.unwrap();
|
|
|
|
// 2. Load through the real adapter — applies the upgrade.
|
|
let adapter = JsonlStorageAdapter::with_root(dir.path().to_path_buf());
|
|
let mut items = adapter.load_chat_history_from_dir(dir.path()).unwrap();
|
|
assert!(
|
|
items
|
|
.iter()
|
|
.any(|i| matches!(i, ConversationItem::Reasoning(_))),
|
|
"legacy inline reasoning must be reconstructed as a sibling on load, got {items:?}"
|
|
);
|
|
|
|
// 3. Continue and send over the Responses API, capturing the body.
|
|
items.push(ConversationItem::user("q2"));
|
|
|
|
let server = MockInferenceServer::start().await.unwrap();
|
|
server.set_response("ok");
|
|
let client = create_test_client(&server.url(), ApiBackend::Responses);
|
|
|
|
let _ = client
|
|
.conversation_collect(ConversationRequest::from_items(items))
|
|
.await
|
|
.unwrap();
|
|
|
|
// 4. The reconstructed reasoning must appear as a typed `reasoning`
|
|
// input item with its fields preserved verbatim — the byte-stable
|
|
// typed round-trip, not a flattened string.
|
|
let body = server.request_bodies().pop().unwrap();
|
|
let input = body.get("input").unwrap().as_array().unwrap();
|
|
let reasoning = input
|
|
.iter()
|
|
.find(|i| i.get("type").and_then(Value::as_str) == Some("reasoning"))
|
|
.unwrap_or_else(|| {
|
|
panic!("reconstructed reasoning must round-trip as a typed Responses input item; input: {input:#?}")
|
|
});
|
|
assert_eq!(
|
|
reasoning.get("id").and_then(Value::as_str),
|
|
Some("rs_grokbuild_legacy"),
|
|
"reasoning id preserved"
|
|
);
|
|
assert_eq!(
|
|
reasoning.get("encrypted_content").and_then(Value::as_str),
|
|
Some("ENC_BLOB_xyz"),
|
|
"encrypted_content preserved verbatim (this is what restores exact tokens server-side)"
|
|
);
|
|
let summary = reasoning.get("summary").and_then(Value::as_array).unwrap();
|
|
assert_eq!(
|
|
summary[0].get("type").and_then(Value::as_str),
|
|
Some("summary_text")
|
|
);
|
|
assert_eq!(
|
|
summary[0].get("text").and_then(Value::as_str),
|
|
Some("legacy grok-build reasoning")
|
|
);
|
|
}
|
|
|
|
/// Upgrade path, Anthropic Messages API: a legacy session whose assistant
|
|
/// carries inline `reasoning: {text, encrypted, id}` (text = thinking,
|
|
/// encrypted = signature) must, on load, reconstruct a sibling Reasoning
|
|
/// item that emits a Anthropic Messages `thinking` content block (with `thinking`
|
|
/// + `signature`) on the outgoing `/v1/messages` request.
|
|
#[tokio::test]
|
|
async fn messages_upgrade_emits_reconstructed_reasoning_as_thinking_block() {
|
|
// 1. Seed a legacy Anthropic Messages-origin chat_history.jsonl. Anthropic Messages
|
|
// thinking blocks never carried an id (stream/messages.rs sets
|
|
// id=""), and the signature lives in `encrypted`.
|
|
let dir = tempfile::tempdir().unwrap();
|
|
std::fs::write(
|
|
dir.path().join("chat_history.jsonl"),
|
|
concat!(
|
|
r#"{"type":"system","content":"You are helpful."}"#,
|
|
"\n",
|
|
r#"{"type":"user","content":[{"type":"text","text":"q1"}]}"#,
|
|
"\n",
|
|
r#"{"type":"assistant","content":"a1","reasoning":{"text":"legacy anthropic thinking","encrypted":"SIGNATURE_abc","id":""},"model_id":"grok-4.5"}"#,
|
|
"\n",
|
|
),
|
|
)
|
|
.unwrap();
|
|
|
|
// 2. Load + upgrade.
|
|
let adapter = JsonlStorageAdapter::with_root(dir.path().to_path_buf());
|
|
let mut items = adapter.load_chat_history_from_dir(dir.path()).unwrap();
|
|
assert!(
|
|
items
|
|
.iter()
|
|
.any(|i| matches!(i, ConversationItem::Reasoning(_))),
|
|
"legacy inline reasoning must be reconstructed as a sibling on load, got {items:?}"
|
|
);
|
|
|
|
// 3. Continue and send over the Messages API, capturing the body.
|
|
items.push(ConversationItem::user("q2"));
|
|
|
|
let server = MockInferenceServer::start().await.unwrap();
|
|
server.set_response("ok");
|
|
let client = create_test_client(&server.url(), ApiBackend::Messages);
|
|
|
|
let _ = client
|
|
.conversation_collect(ConversationRequest::from_items(items))
|
|
.await
|
|
.unwrap();
|
|
|
|
// 4. The reconstructed reasoning must emit a Anthropic Messages `thinking`
|
|
// content block carrying the thinking text + signature.
|
|
let body = server.request_bodies().pop().unwrap();
|
|
let messages = body.get("messages").unwrap().as_array().unwrap();
|
|
let thinking_block = messages
|
|
.iter()
|
|
.filter(|m| m.get("role").and_then(Value::as_str) == Some("assistant"))
|
|
.flat_map(|m| {
|
|
m.get("content")
|
|
.and_then(Value::as_array)
|
|
.cloned()
|
|
.unwrap_or_default()
|
|
})
|
|
.find(|b| b.get("type").and_then(Value::as_str) == Some("thinking"))
|
|
.unwrap_or_else(|| {
|
|
panic!("reconstructed reasoning must emit an Anthropic thinking block; messages: {messages:#?}")
|
|
});
|
|
assert_eq!(
|
|
thinking_block.get("thinking").and_then(Value::as_str),
|
|
Some("legacy anthropic thinking"),
|
|
"thinking text preserved"
|
|
);
|
|
assert_eq!(
|
|
thinking_block.get("signature").and_then(Value::as_str),
|
|
Some("SIGNATURE_abc"),
|
|
"signature (encrypted) preserved — required to reuse the thought server-side"
|
|
);
|
|
}
|
|
|
|
// ============================================================================
|
|
// Responses API Tests
|
|
// ============================================================================
|
|
|
|
/// Generate a Responses API SSE stream with reasoning (including encrypted content).
|
|
fn responses_api_with_reasoning_stream(
|
|
reasoning_summary: &str,
|
|
encrypted_content: Option<&str>,
|
|
content: &str,
|
|
model: &str,
|
|
) -> Vec<SseEvent> {
|
|
let mut events = Vec::new();
|
|
let mut seq = 0;
|
|
|
|
// response.created event
|
|
let created = json!({
|
|
"type": "response.created",
|
|
"sequence_number": seq,
|
|
"response": {
|
|
"id": "resp_reason123",
|
|
"object": "response",
|
|
"created_at": 1234567890,
|
|
"model": model,
|
|
"status": "in_progress",
|
|
"output": []
|
|
}
|
|
});
|
|
events.push(SseEvent::data(created.to_string()));
|
|
seq += 1;
|
|
|
|
// Build reasoning output item
|
|
let mut reasoning_item = json!({
|
|
"type": "reasoning",
|
|
"id": "reasoning_item_1",
|
|
"summary": [{
|
|
"type": "summary_text",
|
|
"text": reasoning_summary
|
|
}],
|
|
"status": "completed"
|
|
});
|
|
|
|
if let Some(enc) = encrypted_content {
|
|
reasoning_item["encrypted_content"] = json!(enc);
|
|
}
|
|
|
|
// response.completed event with reasoning
|
|
let completed = json!({
|
|
"type": "response.completed",
|
|
"sequence_number": seq,
|
|
"response": {
|
|
"id": "resp_reason123",
|
|
"object": "response",
|
|
"created_at": 1234567890,
|
|
"model": model,
|
|
"status": "completed",
|
|
"output": [
|
|
reasoning_item,
|
|
{
|
|
"type": "message",
|
|
"id": "msg_reason123",
|
|
"role": "assistant",
|
|
"status": "completed",
|
|
"content": [{
|
|
"type": "output_text",
|
|
"text": content,
|
|
"annotations": []
|
|
}]
|
|
}
|
|
],
|
|
"usage": {
|
|
"input_tokens": 10,
|
|
"output_tokens": 20,
|
|
"total_tokens": 30,
|
|
"input_tokens_details": {
|
|
"cached_tokens": 0
|
|
},
|
|
"output_tokens_details": {
|
|
"reasoning_tokens": 50
|
|
}
|
|
}
|
|
}
|
|
});
|
|
events.push(SseEvent::data(completed.to_string()));
|
|
events.push(SseEvent::data("[DONE]"));
|
|
|
|
events
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn test_responses_api_streaming_text() {
|
|
let server = MockInferenceServer::start().await.unwrap();
|
|
server.set_response("Hello from Responses API!");
|
|
let client = create_test_client(&server.url(), ApiBackend::Responses);
|
|
|
|
let request = ConversationRequest::from_items(vec![
|
|
ConversationItem::system("You are a helpful assistant."),
|
|
ConversationItem::user("Say hello"),
|
|
]);
|
|
|
|
let (mut stream, _metadata, _) = client.conversation_stream_responses(request).await.unwrap();
|
|
|
|
let mut content = String::new();
|
|
let mut completed = false;
|
|
while let Some(event_result) = stream.next().await {
|
|
let event = event_result.unwrap();
|
|
use xai_grok_shell::sampling::rs::ResponseStreamEvent;
|
|
match event {
|
|
ResponseStreamEvent::ResponseOutputTextDelta(delta) => {
|
|
content.push_str(&delta.delta);
|
|
}
|
|
ResponseStreamEvent::ResponseCompleted(_) => {
|
|
completed = true;
|
|
}
|
|
_ => {}
|
|
}
|
|
}
|
|
|
|
assert!(content.contains("Hello"));
|
|
assert!(content.contains("Responses"));
|
|
assert!(completed);
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn test_responses_api_streaming_tool_call() {
|
|
let server = MockInferenceServer::start().await.unwrap();
|
|
server.enqueue_response(
|
|
"/v1/responses",
|
|
ScriptedResponse::sse(responses_api_tool_call_stream(
|
|
"call_xyz789",
|
|
"bash",
|
|
r#"{"command": "ls -la"}"#,
|
|
"grok-test",
|
|
)),
|
|
);
|
|
let client = create_test_client(&server.url(), ApiBackend::Responses);
|
|
|
|
let request = ConversationRequest::from_items(vec![ConversationItem::user("List files")])
|
|
.with_tools(vec![ToolSpec {
|
|
name: "bash".to_string(),
|
|
description: Some("Run a command".to_string()),
|
|
parameters: json!({"type": "object"}),
|
|
}]);
|
|
|
|
let (mut stream, _metadata, _) = client.conversation_stream_responses(request).await.unwrap();
|
|
|
|
let mut function_call_found = false;
|
|
while let Some(event_result) = stream.next().await {
|
|
let event = event_result.unwrap();
|
|
use xai_grok_shell::sampling::rs::ResponseStreamEvent;
|
|
if let ResponseStreamEvent::ResponseCompleted(completed) = event {
|
|
for output in completed.response.output {
|
|
use xai_grok_shell::sampling::rs::OutputItem;
|
|
if let OutputItem::FunctionCall(fc) = output {
|
|
assert_eq!(fc.call_id, "call_xyz789");
|
|
assert_eq!(fc.name, "bash");
|
|
assert!(fc.arguments.contains("ls -la"));
|
|
function_call_found = true;
|
|
}
|
|
}
|
|
}
|
|
}
|
|
|
|
assert!(function_call_found);
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn test_responses_api_with_reasoning_and_encrypted_content() {
|
|
let server = MockInferenceServer::start().await.unwrap();
|
|
server.enqueue_response(
|
|
"/v1/responses",
|
|
ScriptedResponse::sse(responses_api_with_reasoning_stream(
|
|
"Let me think step by step about this problem.",
|
|
Some("enc_base64_encrypted_reasoning_chain_data"),
|
|
"The answer based on my reasoning is 42.",
|
|
"grok-test",
|
|
)),
|
|
);
|
|
let client = create_test_client(&server.url(), ApiBackend::Responses);
|
|
|
|
let request = ConversationRequest::from_items(vec![ConversationItem::user(
|
|
"What is the meaning of life?",
|
|
)]);
|
|
|
|
let (mut stream, _metadata, _) = client.conversation_stream_responses(request).await.unwrap();
|
|
|
|
let mut found_reasoning = false;
|
|
let mut found_encrypted = false;
|
|
let mut found_content = false;
|
|
|
|
while let Some(event_result) = stream.next().await {
|
|
let event = event_result.unwrap();
|
|
use xai_grok_shell::sampling::rs::ResponseStreamEvent;
|
|
if let ResponseStreamEvent::ResponseCompleted(completed) = event {
|
|
for output in &completed.response.output {
|
|
use xai_grok_shell::sampling::rs::OutputItem;
|
|
match output {
|
|
OutputItem::Reasoning(r) => {
|
|
// Check summary text
|
|
if !r.summary.is_empty() {
|
|
found_reasoning = true;
|
|
}
|
|
// Check encrypted content
|
|
if let Some(encrypted_content) = &r.encrypted_content {
|
|
found_encrypted = true;
|
|
assert!(encrypted_content.contains("enc_base64"));
|
|
}
|
|
}
|
|
OutputItem::Message(msg) => {
|
|
for content in &msg.content {
|
|
use xai_grok_shell::sampling::rs::OutputMessageContent;
|
|
if let OutputMessageContent::OutputText(text) = content
|
|
&& text.text.contains("42")
|
|
{
|
|
found_content = true;
|
|
}
|
|
}
|
|
}
|
|
_ => {}
|
|
}
|
|
}
|
|
}
|
|
}
|
|
|
|
assert!(found_reasoning, "Should have found reasoning summary");
|
|
assert!(
|
|
found_encrypted,
|
|
"Should have found encrypted reasoning content"
|
|
);
|
|
assert!(found_content, "Should have found message content");
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn test_responses_api_reasoning_without_encrypted() {
|
|
// Test reasoning with only visible summary, no encrypted content
|
|
let server = MockInferenceServer::start().await.unwrap();
|
|
server.enqueue_response(
|
|
"/v1/responses",
|
|
ScriptedResponse::sse(responses_api_with_reasoning_stream(
|
|
"I need to analyze the code carefully.",
|
|
None, // No encrypted content
|
|
"Here is my analysis.",
|
|
"grok-test",
|
|
)),
|
|
);
|
|
let client = create_test_client(&server.url(), ApiBackend::Responses);
|
|
|
|
let request =
|
|
ConversationRequest::from_items(vec![ConversationItem::user("Analyze this code")]);
|
|
|
|
let (mut stream, _metadata, _) = client.conversation_stream_responses(request).await.unwrap();
|
|
|
|
let mut found_reasoning = false;
|
|
while let Some(event_result) = stream.next().await {
|
|
let event = event_result.unwrap();
|
|
use xai_grok_shell::sampling::rs::ResponseStreamEvent;
|
|
if let ResponseStreamEvent::ResponseCompleted(completed) = event {
|
|
for output in &completed.response.output {
|
|
use xai_grok_shell::sampling::rs::OutputItem;
|
|
if let OutputItem::Reasoning(r) = output {
|
|
found_reasoning = true;
|
|
// Should have summary but no encrypted content
|
|
assert!(!r.summary.is_empty());
|
|
assert!(r.encrypted_content.is_none());
|
|
}
|
|
}
|
|
}
|
|
}
|
|
|
|
assert!(found_reasoning);
|
|
}
|
|
|
|
// ============================================================================
|
|
// Error Handling Tests
|
|
// ============================================================================
|
|
|
|
#[tokio::test]
|
|
async fn test_chat_completions_401_unauthorized() {
|
|
let server = MockInferenceServer::start().await.unwrap();
|
|
server.enqueue_response(
|
|
"/v1/chat/completions",
|
|
ScriptedResponse::text(401, "Unauthorized"),
|
|
);
|
|
let client = create_test_client(&server.url(), ApiBackend::ChatCompletions);
|
|
|
|
let request = ConversationRequest::from_items(vec![ConversationItem::user("Hello")]);
|
|
|
|
let result = client.conversation_stream(request).await;
|
|
assert!(result.is_err());
|
|
|
|
if let Err(SamplingError::Auth(_)) = result {
|
|
// Expected
|
|
} else {
|
|
panic!("Expected Auth error");
|
|
}
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn test_chat_completions_500_server_error() {
|
|
let server = MockInferenceServer::start().await.unwrap();
|
|
server.enqueue_response(
|
|
"/v1/chat/completions",
|
|
ScriptedResponse::json(500, json!({"error": {"message": "Internal server error"}})),
|
|
);
|
|
let client = create_test_client(&server.url(), ApiBackend::ChatCompletions);
|
|
|
|
let request = ConversationRequest::from_items(vec![ConversationItem::user("Hello")]);
|
|
|
|
let result = client.conversation_stream(request).await;
|
|
assert!(result.is_err());
|
|
|
|
if let Err(SamplingError::Api {
|
|
status, message, ..
|
|
}) = result
|
|
{
|
|
assert_eq!(status, StatusCode::INTERNAL_SERVER_ERROR);
|
|
assert!(message.contains("Internal server error"));
|
|
} else {
|
|
panic!("Expected Api error");
|
|
}
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn test_responses_api_401_unauthorized() {
|
|
let server = MockInferenceServer::start().await.unwrap();
|
|
server.enqueue_response("/v1/responses", ScriptedResponse::text(401, "Unauthorized"));
|
|
let client = create_test_client(&server.url(), ApiBackend::Responses);
|
|
|
|
let request = ConversationRequest::from_items(vec![ConversationItem::user("Hello")]);
|
|
|
|
let result = client.conversation_stream_responses(request).await;
|
|
assert!(result.is_err());
|
|
|
|
if let Err(SamplingError::Auth(_)) = result {
|
|
// Expected
|
|
} else {
|
|
panic!("Expected Auth error");
|
|
}
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn test_stream_error_during_streaming() {
|
|
// Simulate a stream error mid-response
|
|
let events = vec![
|
|
SseEvent::data(
|
|
json!({
|
|
"id": "chatcmpl-test123",
|
|
"object": "chat.completion.chunk",
|
|
"created": 1234567890,
|
|
"model": "grok-test",
|
|
"choices": [{
|
|
"index": 0,
|
|
"delta": {"role": "assistant", "content": "Hello"},
|
|
"finish_reason": null
|
|
}]
|
|
})
|
|
.to_string(),
|
|
),
|
|
// Stream error
|
|
SseEvent::data(
|
|
json!({
|
|
"error": {
|
|
"message": "Stream interrupted",
|
|
"type": "stream_error"
|
|
}
|
|
})
|
|
.to_string(),
|
|
),
|
|
];
|
|
|
|
let server = MockInferenceServer::start().await.unwrap();
|
|
server.enqueue_response("/v1/chat/completions", ScriptedResponse::sse(events));
|
|
let client = create_test_client(&server.url(), ApiBackend::ChatCompletions);
|
|
|
|
let request = ConversationRequest::from_items(vec![ConversationItem::user("Hello")]);
|
|
|
|
let (mut stream, _metadata) = client.conversation_stream(request).await.unwrap();
|
|
|
|
let mut got_content = false;
|
|
let mut got_error = false;
|
|
|
|
while let Some(result) = stream.next().await {
|
|
match result {
|
|
Ok(chunk) => {
|
|
for choice in chunk.choices {
|
|
if choice.delta.content.is_some() {
|
|
got_content = true;
|
|
}
|
|
}
|
|
}
|
|
Err(SamplingError::StreamError { .. }) => {
|
|
got_error = true;
|
|
}
|
|
Err(e) => {
|
|
panic!("Unexpected error: {:?}", e);
|
|
}
|
|
}
|
|
}
|
|
|
|
assert!(got_content, "Should have received content before error");
|
|
assert!(got_error, "Should have received stream error");
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn test_stream_error_during_responses_streaming() {
|
|
// Simulate a stream error mid-response on the Responses API path.
|
|
// This mirrors test_stream_error_during_streaming but exercises the
|
|
// Responses API stream-error detection (the second call site for the
|
|
// fast-path contains("error") guard).
|
|
let events = vec![
|
|
SseEvent::with_event(
|
|
"response.output_text.delta",
|
|
json!({
|
|
"type": "response.output_text.delta",
|
|
"sequence_number": 0,
|
|
"item_id": "item_test",
|
|
"output_index": 0,
|
|
"content_index": 0,
|
|
"delta": "Hello"
|
|
})
|
|
.to_string(),
|
|
),
|
|
// Stream error
|
|
SseEvent::data(
|
|
json!({
|
|
"error": {
|
|
"message": "Stream interrupted",
|
|
"type": "stream_error"
|
|
}
|
|
})
|
|
.to_string(),
|
|
),
|
|
];
|
|
|
|
let server = MockInferenceServer::start().await.unwrap();
|
|
server.enqueue_response("/v1/responses", ScriptedResponse::sse(events));
|
|
let client = create_test_client(&server.url(), ApiBackend::Responses);
|
|
|
|
let request = ConversationRequest::from_items(vec![ConversationItem::user("Hello")]);
|
|
|
|
let (mut stream, _metadata, _) = client.conversation_stream_responses(request).await.unwrap();
|
|
|
|
let mut got_event = false;
|
|
let mut got_error = false;
|
|
|
|
while let Some(result) = stream.next().await {
|
|
match result {
|
|
Ok(_) => {
|
|
got_event = true;
|
|
}
|
|
Err(SamplingError::StreamError { .. }) => {
|
|
got_error = true;
|
|
}
|
|
Err(e) => {
|
|
panic!("Unexpected error: {:?}", e);
|
|
}
|
|
}
|
|
}
|
|
|
|
assert!(got_event, "Should have received an event before error");
|
|
assert!(got_error, "Should have received stream error");
|
|
}
|
|
|
|
// ============================================================================
|
|
// Request Validation Tests
|
|
// ============================================================================
|
|
|
|
#[tokio::test]
|
|
async fn test_request_includes_headers() {
|
|
let server = MockInferenceServer::start().await.unwrap();
|
|
server.set_response("OK");
|
|
let client = create_test_client(&server.url(), ApiBackend::ChatCompletions);
|
|
|
|
let request = ConversationRequest::from_items(vec![ConversationItem::user("Hello")])
|
|
.with_conv_id("conv-12345")
|
|
.with_req_id("req-67890");
|
|
|
|
let (mut stream, _metadata) = client.conversation_stream(request).await.unwrap();
|
|
while stream.next().await.is_some() {}
|
|
|
|
let request = server.requests().pop().unwrap();
|
|
|
|
assert_eq!(request.header("authorization"), Some("Bearer test-api-key"));
|
|
assert_eq!(request.header("x-grok-conv-id"), Some("conv-12345"));
|
|
assert_eq!(request.header("x-grok-req-id"), Some("req-67890"));
|
|
}
|
|
|
|
/// The session writes the resolved `x-compaction-at` value into
|
|
/// `SamplerConfig::extra_headers`; this proves the client forwards it on the wire.
|
|
#[tokio::test]
|
|
async fn test_request_forwards_compaction_at_header() {
|
|
let server = MockInferenceServer::start().await.unwrap();
|
|
server.set_response("OK");
|
|
let client = create_test_client_with_extra_headers(
|
|
&server.url(),
|
|
ApiBackend::ChatCompletions,
|
|
&[("x-compaction-at", "217600")],
|
|
);
|
|
|
|
let request = ConversationRequest::from_items(vec![ConversationItem::user("Hello")]);
|
|
let (mut stream, _metadata) = client.conversation_stream(request).await.unwrap();
|
|
while stream.next().await.is_some() {}
|
|
|
|
let request = server.requests().pop().unwrap();
|
|
assert_eq!(request.header("x-compaction-at"), Some("217600"));
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn test_request_includes_tools() {
|
|
let server = MockInferenceServer::start().await.unwrap();
|
|
server.set_response("OK");
|
|
let client = create_test_client(&server.url(), ApiBackend::ChatCompletions);
|
|
|
|
let request = ConversationRequest::from_items(vec![ConversationItem::user("Hello")])
|
|
.with_tools(vec![
|
|
ToolSpec {
|
|
name: "read_file".to_string(),
|
|
description: Some("Read a file".to_string()),
|
|
parameters: json!({"type": "object", "properties": {"path": {"type": "string"}}}),
|
|
},
|
|
ToolSpec {
|
|
name: "bash".to_string(),
|
|
description: Some("Run a command".to_string()),
|
|
parameters: json!({"type": "object"}),
|
|
},
|
|
])
|
|
.with_tool_choice(ConversationToolChoice::Auto);
|
|
|
|
let (mut stream, _metadata) = client.conversation_stream(request).await.unwrap();
|
|
while stream.next().await.is_some() {}
|
|
|
|
let body = server.request_bodies().pop().unwrap();
|
|
|
|
// Verify tools were included
|
|
let tools = body.get("tools").unwrap().as_array().unwrap();
|
|
assert_eq!(tools.len(), 2);
|
|
assert_eq!(tools[0]["function"]["name"], "read_file");
|
|
assert_eq!(tools[1]["function"]["name"], "bash");
|
|
|
|
// Verify tool_choice
|
|
assert_eq!(body.get("tool_choice").unwrap(), "auto");
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn test_responses_api_request_format() {
|
|
let server = MockInferenceServer::start().await.unwrap();
|
|
server.set_response("OK");
|
|
let client = create_test_client(&server.url(), ApiBackend::Responses);
|
|
|
|
let request = ConversationRequest::from_items(vec![
|
|
ConversationItem::system("You are helpful"),
|
|
ConversationItem::user("Hello"),
|
|
])
|
|
.with_temperature(0.5)
|
|
.with_max_output_tokens(500);
|
|
|
|
let (mut stream, _metadata, _) = client.conversation_stream_responses(request).await.unwrap();
|
|
while stream.next().await.is_some() {}
|
|
|
|
let body = server.request_bodies().pop().unwrap();
|
|
|
|
// Verify Responses API format
|
|
assert!(body.get("input").is_some());
|
|
assert_eq!(body.get("temperature").unwrap(), 0.5);
|
|
assert_eq!(body.get("max_output_tokens").unwrap(), 500);
|
|
assert_eq!(body.get("stream").unwrap(), true);
|
|
|
|
// Verify input items format
|
|
let input = body.get("input").unwrap().as_array().unwrap();
|
|
assert!(input.len() >= 2);
|
|
}
|
|
|
|
/// The sampler owns the doom-loop opt-in: setting
|
|
/// `SamplerConfig::doom_loop_recovery` puts `x-grok-doom-loop-check` on the
|
|
/// wire AND arms the collector, and the server's named check event is
|
|
/// absorbed mid-stream without disturbing the typed event flow.
|
|
#[tokio::test]
|
|
async fn test_doom_loop_check_enabled_sends_header_and_absorbs_check_event() {
|
|
use xai_grok_sampling_types::doom_loop::{DOOM_LOOP_CHECK_EVENT_TYPE, SAMPLE_CHECK_EVENT_DATA};
|
|
|
|
let server = MockInferenceServer::start().await.unwrap();
|
|
let mut events =
|
|
responses_api_reasoning_and_text_events("pondering", "Hello there", "test-model");
|
|
// Splice the named check event right after `response.created`.
|
|
events.insert(
|
|
1,
|
|
SseEvent::with_event(DOOM_LOOP_CHECK_EVENT_TYPE, SAMPLE_CHECK_EVENT_DATA),
|
|
);
|
|
server.enqueue_response("/v1/responses", ScriptedResponse::sse(events));
|
|
|
|
let mut config = test_sampler_config(&server.url(), ApiBackend::Responses, &[]);
|
|
config.doom_loop_recovery = Some(Default::default());
|
|
let client = Client::new(config).unwrap();
|
|
|
|
let request = ConversationRequest::from_items(vec![ConversationItem::user("Hello")]);
|
|
let (mut stream, _metadata, collector) =
|
|
client.conversation_stream_responses(request).await.unwrap();
|
|
assert!(collector.is_some(), "set policy must arm a collector");
|
|
|
|
let mut completed = false;
|
|
while let Some(event_result) = stream.next().await {
|
|
let event = event_result.expect("absorbed check event must not fail the typed stream");
|
|
if matches!(event, rs::ResponseStreamEvent::ResponseCompleted(_)) {
|
|
completed = true;
|
|
}
|
|
}
|
|
assert!(completed);
|
|
|
|
let logged = server.requests().pop().unwrap();
|
|
assert!(logged.path.contains("/responses"));
|
|
assert_eq!(logged.header("x-grok-doom-loop-check"), Some("true"));
|
|
}
|
|
|
|
/// With the check disabled no header goes on the wire, and check frames from
|
|
/// a misbehaving server (rollout skew) are dropped instead of failing the
|
|
/// typed stream — a named frame even with a garbage payload, and an unnamed
|
|
/// frame identified only by its payload `type` tag.
|
|
#[tokio::test]
|
|
async fn test_doom_loop_check_disabled_sends_no_header_and_drops_check_frames() {
|
|
use xai_grok_sampling_types::doom_loop::{DOOM_LOOP_CHECK_EVENT_TYPE, SAMPLE_CHECK_EVENT_DATA};
|
|
|
|
let server = MockInferenceServer::start().await.unwrap();
|
|
let mut events =
|
|
responses_api_reasoning_and_text_events("pondering", "Hello there", "test-model");
|
|
events.insert(
|
|
1,
|
|
SseEvent::with_event(DOOM_LOOP_CHECK_EVENT_TYPE, "this is not json"),
|
|
);
|
|
events.insert(2, SseEvent::data(SAMPLE_CHECK_EVENT_DATA));
|
|
server.enqueue_response("/v1/responses", ScriptedResponse::sse(events));
|
|
|
|
let client = create_test_client(&server.url(), ApiBackend::Responses);
|
|
let request = ConversationRequest::from_items(vec![ConversationItem::user("Hello")]);
|
|
let (mut stream, _metadata, collector) =
|
|
client.conversation_stream_responses(request).await.unwrap();
|
|
assert!(
|
|
collector.is_none(),
|
|
"disabled policy must not arm a collector"
|
|
);
|
|
|
|
let mut completed = false;
|
|
while let Some(event_result) = stream.next().await {
|
|
let event = event_result.expect("check frames must be dropped, not fail the stream");
|
|
if matches!(event, rs::ResponseStreamEvent::ResponseCompleted(_)) {
|
|
completed = true;
|
|
}
|
|
}
|
|
assert!(completed);
|
|
|
|
let logged = server.requests().pop().unwrap();
|
|
assert!(logged.path.contains("/responses"));
|
|
assert_eq!(logged.header("x-grok-doom-loop-check"), None);
|
|
}
|
|
|
|
// ============================================================================
|
|
// Multi-turn Conversation Tests
|
|
// ============================================================================
|
|
|
|
#[tokio::test]
|
|
async fn test_multi_turn_conversation_with_tool_calls() {
|
|
let server = MockInferenceServer::start().await.unwrap();
|
|
server.set_response("I've read the file for you.");
|
|
let client = create_test_client(&server.url(), ApiBackend::ChatCompletions);
|
|
|
|
// Simulate a multi-turn conversation with tool call and result
|
|
let request = ConversationRequest::from_items(vec![
|
|
ConversationItem::system("You are a helpful assistant."),
|
|
ConversationItem::user("Read the README file"),
|
|
ConversationItem::assistant_tool_calls(vec![ToolCall {
|
|
id: "call_1".into(),
|
|
name: "read_file".to_string(),
|
|
arguments: r#"{"path": "README.md"}"#.into(),
|
|
}]),
|
|
ConversationItem::tool_result("call_1", "# My Project\n\nThis is a test project."),
|
|
// The model should now respond based on the file content
|
|
]);
|
|
|
|
let (mut stream, _metadata) = client.conversation_stream(request).await.unwrap();
|
|
|
|
let mut content = String::new();
|
|
while let Some(chunk_result) = stream.next().await {
|
|
let chunk = chunk_result.unwrap();
|
|
for choice in chunk.choices {
|
|
if let Some(ref text) = choice.delta.content {
|
|
content.push_str(text);
|
|
}
|
|
}
|
|
}
|
|
|
|
assert!(!content.is_empty());
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn test_responses_api_multi_turn_with_tool_calls() {
|
|
let server = MockInferenceServer::start().await.unwrap();
|
|
server.set_response("Done with the file.");
|
|
let client = create_test_client(&server.url(), ApiBackend::Responses);
|
|
|
|
// Multi-turn with tool call and result
|
|
let request = ConversationRequest::from_items(vec![
|
|
ConversationItem::user("Read the config"),
|
|
ConversationItem::assistant_tool_calls(vec![ToolCall {
|
|
id: "call_abc".into(),
|
|
name: "read_file".to_string(),
|
|
arguments: r#"{"path": "config.json"}"#.into(),
|
|
}]),
|
|
ConversationItem::tool_result("call_abc", r#"{"key": "value"}"#),
|
|
]);
|
|
|
|
let (mut stream, _metadata, _) = client.conversation_stream_responses(request).await.unwrap();
|
|
|
|
let mut completed = false;
|
|
while let Some(event_result) = stream.next().await {
|
|
let event = event_result.unwrap();
|
|
use xai_grok_shell::sampling::rs::ResponseStreamEvent;
|
|
if let ResponseStreamEvent::ResponseCompleted(_) = event {
|
|
completed = true;
|
|
}
|
|
}
|
|
|
|
assert!(completed);
|
|
}
|
|
|
|
// ============================================================================
|
|
// Request Counter Tests (verify retry behavior, etc.)
|
|
// ============================================================================
|
|
|
|
#[tokio::test]
|
|
async fn test_single_request_per_stream() {
|
|
let server = MockInferenceServer::start().await.unwrap();
|
|
server.set_response("Hello");
|
|
let client = create_test_client(&server.url(), ApiBackend::ChatCompletions);
|
|
|
|
let request = ConversationRequest::from_items(vec![ConversationItem::user("Hello")]);
|
|
|
|
let (mut stream, _metadata) = client.conversation_stream(request).await.unwrap();
|
|
while stream.next().await.is_some() {}
|
|
|
|
assert_eq!(server.request_count(), 1);
|
|
}
|
|
|
|
// ============================================================================
|
|
// API Backend Routing Tests
|
|
// ============================================================================
|
|
|
|
#[tokio::test]
|
|
async fn test_api_backend_getter_returns_configured_value() {
|
|
// Verify that the client correctly reports its configured API backend
|
|
let client_responses = create_test_client("http://localhost/v1", ApiBackend::Responses);
|
|
assert_eq!(client_responses.api_backend(), ApiBackend::Responses);
|
|
|
|
let client_chat = create_test_client("http://localhost/v1", ApiBackend::ChatCompletions);
|
|
assert_eq!(client_chat.api_backend(), ApiBackend::ChatCompletions);
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn test_responses_backend_hits_responses_endpoint_not_chat_completions() {
|
|
// This test verifies that when ApiBackend::Responses is configured,
|
|
// the client hits /v1/responses and NOT /v1/chat/completions.
|
|
// The session-level dispatch in `xai-grok-sampler` selects the
|
|
// backend stream based on `SamplingClient::api_backend()`.
|
|
|
|
let server = MockInferenceServer::start().await.unwrap();
|
|
// A request to the wrong endpoint fails loudly instead of streaming.
|
|
server.enqueue_response(
|
|
"/v1/chat/completions",
|
|
ScriptedResponse::text(500, "Wrong endpoint!"),
|
|
);
|
|
server.set_response("OK");
|
|
let client = create_test_client(&server.url(), ApiBackend::Responses);
|
|
|
|
// Simulate the routing logic from acp_session.rs
|
|
match client.api_backend() {
|
|
ApiBackend::Responses => {
|
|
let request = ConversationRequest::from_items(vec![ConversationItem::user("Hello")]);
|
|
let (mut stream, _metadata, _) =
|
|
client.conversation_stream_responses(request).await.unwrap();
|
|
while stream.next().await.is_some() {}
|
|
}
|
|
ApiBackend::ChatCompletions => {
|
|
panic!("Expected Responses backend but got ChatCompletions");
|
|
}
|
|
ApiBackend::Messages => {
|
|
panic!("Expected Responses backend but got Messages");
|
|
}
|
|
}
|
|
|
|
assert!(
|
|
!server.has_chat_completion_request(),
|
|
"Should NOT have called /v1/chat/completions"
|
|
);
|
|
assert!(
|
|
server.has_responses_request(),
|
|
"Should have called /v1/responses"
|
|
);
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn test_chat_completions_backend_hits_chat_endpoint_not_responses() {
|
|
// Verify the inverse: ChatCompletions backend hits /v1/chat/completions
|
|
|
|
let server = MockInferenceServer::start().await.unwrap();
|
|
// A request to the wrong endpoint fails loudly instead of streaming.
|
|
server.enqueue_response(
|
|
"/v1/responses",
|
|
ScriptedResponse::text(500, "Wrong endpoint!"),
|
|
);
|
|
server.set_response("OK");
|
|
let client = create_test_client(&server.url(), ApiBackend::ChatCompletions);
|
|
|
|
// Simulate the routing logic from acp_session.rs
|
|
match client.api_backend() {
|
|
ApiBackend::ChatCompletions => {
|
|
let request = ConversationRequest::from_items(vec![ConversationItem::user("Hello")]);
|
|
let (mut stream, _metadata) = client.conversation_stream(request).await.unwrap();
|
|
while stream.next().await.is_some() {}
|
|
}
|
|
ApiBackend::Responses => {
|
|
panic!("Expected ChatCompletions backend but got Responses");
|
|
}
|
|
ApiBackend::Messages => {
|
|
panic!("Expected ChatCompletions backend but got Messages");
|
|
}
|
|
}
|
|
|
|
assert!(
|
|
server.has_chat_completion_request(),
|
|
"Should have called /v1/chat/completions"
|
|
);
|
|
assert!(
|
|
!server.has_responses_request(),
|
|
"Should NOT have called /v1/responses"
|
|
);
|
|
}
|