F3: Kimi inference pipeline + full grok cloud-surface excision
Sampler / inference (PRD F3):
- kimi_compat.rs: single adaptation point for the Kimi chat/completions
dialect (thinking-field mapping, model_id stripping, empty-content
tool-call message fix, stream_options.include_usage), with kimi-cli
source citations
- Rate-limit handling reworked for Kimi/Moonshot semantics; UA kigi/{version}
- /models replaces the xAI models-v2 endpoint everywhere; idle model
refresh carries X-Msh-* device headers only (X-XAI-Token-Auth and
x-grok-client-mode/CLIENT_MODE_HEADER machinery deleted)
Cloud-surface excision (PRD §5, zero-egress):
- remote/ conversations lane, cli-chat-proxy-types crate, prod/ dir,
share command, credit bar: deleted (single local session lane;
paginate() replaces merge_and_paginate)
- Subscription/tier gate stack deleted end-to-end: AppView
gate/tier/team/ZDR fields, app/subscription.rs watch loop,
dispatch/billing.rs paywall + SuperGrok upsell, free-usage-exhausted
chain, tier-restricted commands, GateInfo, RemoteSettings gate fields,
SettingsUpdateNotification gate fields
- /privacy + coding-data-sharing setting deleted (backed by a dead xAI
RPC; Kigi is zero-egress — nothing to share or retain remotely)
Auth UX correctness (user-reported):
- Device-flow fixtures now mirror the live Kimi payload shape
(https://www.kimi.com/code/authorize_device?user_code=..., verified
against auth.kimi.com); the fabricated auth.kimi.com/device?code=...
URLs are gone
- open_browser_detached is a no-op under cfg(test): unit tests drove
wiremock fixture URLs into the real browser (root cause of the
"garbage mock link" ABCD-1234 tabs)
- Welcome/pager-minimal rebrand: Grok Build -> Kigi, grok.com ->
kimi.com, "Sign in to Grok" -> "Sign in to Kimi"
This commit is contained in:
@@ -87,10 +87,6 @@ fn test_config(base_url: String, model: &str) -> SamplerConfig {
|
||||
idle_timeout_secs: Some(30),
|
||||
reasoning_effort: None,
|
||||
origin_client: None,
|
||||
client_identifier: None,
|
||||
deployment_id: None,
|
||||
user_id: None,
|
||||
client_version: None,
|
||||
attribution_callback: None,
|
||||
bearer_resolver: None,
|
||||
supports_backend_search: false,
|
||||
|
||||
@@ -0,0 +1,564 @@
|
||||
//! Kimi chat/completions wire tests (PRD F3 acceptance).
|
||||
//!
|
||||
//! Exercises the sampler end-to-end against a mock HTTP server:
|
||||
//! * streaming happy path with `reasoning_content` deltas, tool-call deltas,
|
||||
//! and a Kimi-shaped usage chunk (usage riding inside the choice, cache
|
||||
//! hits as top-level `cached_tokens`),
|
||||
//! * the request the wire actually carries: plain `Authorization: Bearer`,
|
||||
//! `User-Agent: kigi/{version}`, no xAI proxy marker headers, and the
|
||||
//! `crate::kimi_compat` body adaptations,
|
||||
//! * 429 honoring the standard `Retry-After` header,
|
||||
//! * mid-stream network drop recovering through the retry loop.
|
||||
//!
|
||||
//! 401-no-retry and the rate-limit retry threshold are covered by
|
||||
//! `test_actor.rs`; this file does not duplicate them.
|
||||
|
||||
use std::net::SocketAddr;
|
||||
use std::sync::atomic::{AtomicU32, Ordering};
|
||||
use std::sync::{Arc, Mutex};
|
||||
use std::time::Duration;
|
||||
|
||||
use axum::Router;
|
||||
use axum::body::Bytes;
|
||||
use axum::http::{HeaderMap, StatusCode};
|
||||
use axum::response::sse::{Event, Sse};
|
||||
use axum::routing::post;
|
||||
use futures_util::stream::{self};
|
||||
use indexmap::IndexMap;
|
||||
use serde_json::{Value, json};
|
||||
use tokio::net::TcpListener;
|
||||
use tokio::sync::{mpsc, oneshot};
|
||||
|
||||
use kigi_sampler::{
|
||||
ApiBackend, RequestId, RetryPolicy, SamplerActor, SamplerConfig, SamplingChannel, SamplingEvent,
|
||||
};
|
||||
use kigi_sampling_types::{
|
||||
AssistantItem, ContentPart, ConversationItem, ConversationRequest, ReasoningEffort, ToolCall,
|
||||
ToolResultItem, ToolSpec, UserItem, synthesized_reasoning_item,
|
||||
};
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// Mock server harness (same shape as test_actor.rs)
|
||||
// ---------------------------------------------------------------------------
|
||||
|
||||
struct MockServer {
|
||||
addr: SocketAddr,
|
||||
shutdown_tx: oneshot::Sender<()>,
|
||||
}
|
||||
|
||||
impl MockServer {
|
||||
async fn spawn(app: Router) -> Self {
|
||||
let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
|
||||
let addr = listener.local_addr().unwrap();
|
||||
let (shutdown_tx, shutdown_rx) = oneshot::channel::<()>();
|
||||
tokio::spawn(async move {
|
||||
let _ = axum::serve(listener, app)
|
||||
.with_graceful_shutdown(async move {
|
||||
let _ = shutdown_rx.await;
|
||||
})
|
||||
.await;
|
||||
});
|
||||
tokio::time::sleep(Duration::from_millis(20)).await;
|
||||
Self { addr, shutdown_tx }
|
||||
}
|
||||
|
||||
fn base_url(&self) -> String {
|
||||
format!("http://{}/v1", self.addr)
|
||||
}
|
||||
|
||||
fn shutdown(self) {
|
||||
let _ = self.shutdown_tx.send(());
|
||||
}
|
||||
}
|
||||
|
||||
fn test_config(base_url: String) -> SamplerConfig {
|
||||
SamplerConfig {
|
||||
api_key: Some("test-kimi-key".into()),
|
||||
base_url,
|
||||
model: "kimi-for-coding".into(),
|
||||
max_completion_tokens: Some(1024),
|
||||
api_backend: ApiBackend::ChatCompletions,
|
||||
extra_headers: IndexMap::new(),
|
||||
context_window: 128_000,
|
||||
max_retries: Some(3),
|
||||
idle_timeout_secs: Some(30),
|
||||
..Default::default()
|
||||
}
|
||||
}
|
||||
|
||||
fn user_request(text: &str) -> ConversationRequest {
|
||||
ConversationRequest {
|
||||
items: vec![ConversationItem::User(UserItem {
|
||||
content: vec![ContentPart::Text {
|
||||
text: Arc::<str>::from(text),
|
||||
}],
|
||||
synthetic_reason: None,
|
||||
..Default::default()
|
||||
})],
|
||||
..Default::default()
|
||||
}
|
||||
}
|
||||
|
||||
fn chunk(delta: Value, finish: Option<&str>, usage: Option<Value>) -> Event {
|
||||
let mut choice = json!({ "index": 0, "delta": delta });
|
||||
choice["finish_reason"] = finish.map(Value::from).unwrap_or(Value::Null);
|
||||
if let Some(u) = usage {
|
||||
// Kimi deviation under test: usage rides INSIDE the choice
|
||||
// (kimi-cli kimi.py:522-533 `extract_usage_from_chunk`).
|
||||
choice["usage"] = u;
|
||||
}
|
||||
let body = json!({
|
||||
"id": "chatcmpl-kimi",
|
||||
"object": "chat.completion.chunk",
|
||||
"created": 0,
|
||||
"model": "kimi-for-coding",
|
||||
"choices": [choice]
|
||||
});
|
||||
Event::default().data(body.to_string())
|
||||
}
|
||||
|
||||
async fn drain_until_terminal(
|
||||
rx: &mut mpsc::UnboundedReceiver<SamplingEvent>,
|
||||
timeout: Duration,
|
||||
) -> Vec<SamplingEvent> {
|
||||
let mut out = Vec::new();
|
||||
let deadline = tokio::time::Instant::now() + timeout;
|
||||
loop {
|
||||
let ev = tokio::time::timeout_at(deadline, rx.recv())
|
||||
.await
|
||||
.expect("timed out waiting for terminal event")
|
||||
.expect("event channel closed before terminal event");
|
||||
let terminal = matches!(
|
||||
ev,
|
||||
SamplingEvent::Completed { .. } | SamplingEvent::Failed { .. }
|
||||
);
|
||||
out.push(ev);
|
||||
if terminal {
|
||||
return out;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// Streaming happy path: reasoning + tool calls + Kimi usage shapes
|
||||
// ---------------------------------------------------------------------------
|
||||
|
||||
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
|
||||
async fn kimi_stream_reasoning_tool_calls_and_choice_usage() {
|
||||
let app = Router::new().route(
|
||||
"/v1/chat/completions",
|
||||
post(|| async {
|
||||
let events = vec![
|
||||
chunk(
|
||||
json!({ "role": "assistant", "reasoning_content": "let me think" }),
|
||||
None,
|
||||
None,
|
||||
),
|
||||
chunk(json!({ "reasoning_content": " harder" }), None, None),
|
||||
chunk(json!({ "content": "Running the tool." }), None, None),
|
||||
// Tool call split across chunks: id+name first, args continue.
|
||||
chunk(
|
||||
json!({ "tool_calls": [{ "index": 0, "id": "call_1", "type": "function",
|
||||
"function": { "name": "read_file", "arguments": "{\"path\":" } }] }),
|
||||
None,
|
||||
None,
|
||||
),
|
||||
chunk(
|
||||
json!({ "tool_calls": [{ "index": 0,
|
||||
"function": { "arguments": "\"a.rs\"}" } }] }),
|
||||
None,
|
||||
None,
|
||||
),
|
||||
// Terminal chunk: finish_reason + usage inside the choice with
|
||||
// Moonshot's top-level `cached_tokens` (kimi.py:427-431).
|
||||
chunk(
|
||||
json!({}),
|
||||
Some("tool_calls"),
|
||||
Some(json!({
|
||||
"prompt_tokens": 100,
|
||||
"completion_tokens": 20,
|
||||
"total_tokens": 120,
|
||||
"cached_tokens": 60
|
||||
})),
|
||||
),
|
||||
];
|
||||
Sse::new(stream::iter(
|
||||
events.into_iter().map(Ok::<_, std::convert::Infallible>),
|
||||
))
|
||||
}),
|
||||
);
|
||||
let server = MockServer::spawn(app).await;
|
||||
let (event_tx, mut event_rx) = mpsc::unbounded_channel();
|
||||
let handle = SamplerActor::spawn(
|
||||
test_config(server.base_url()),
|
||||
RetryPolicy::default(),
|
||||
event_tx,
|
||||
);
|
||||
|
||||
handle.submit(RequestId::from("req-kimi"), user_request("hi"));
|
||||
let events = drain_until_terminal(&mut event_rx, Duration::from_secs(30)).await;
|
||||
server.shutdown();
|
||||
|
||||
// Reasoning tokens stream on the Reasoning channel, text on Text.
|
||||
let reasoning: String = events
|
||||
.iter()
|
||||
.filter_map(|e| match e {
|
||||
SamplingEvent::ChannelToken {
|
||||
channel: SamplingChannel::Reasoning,
|
||||
text,
|
||||
..
|
||||
} => Some(text.as_str()),
|
||||
_ => None,
|
||||
})
|
||||
.collect();
|
||||
assert_eq!(reasoning, "let me think harder");
|
||||
let text: String = events
|
||||
.iter()
|
||||
.filter_map(|e| match e {
|
||||
SamplingEvent::ChannelToken {
|
||||
channel: SamplingChannel::Text,
|
||||
text,
|
||||
..
|
||||
} => Some(text.as_str()),
|
||||
_ => None,
|
||||
})
|
||||
.collect();
|
||||
assert_eq!(text, "Running the tool.");
|
||||
|
||||
// Tool-call deltas surfaced incrementally.
|
||||
assert!(events.iter().any(|e| matches!(
|
||||
e,
|
||||
SamplingEvent::ToolCallDelta { id: Some(id), .. } if id == "call_1"
|
||||
)));
|
||||
|
||||
match events.last().unwrap() {
|
||||
SamplingEvent::Completed { response, .. } => {
|
||||
let calls = response.tool_calls();
|
||||
assert_eq!(calls.len(), 1);
|
||||
assert_eq!(calls[0].id.as_ref(), "call_1");
|
||||
assert_eq!(calls[0].name, "read_file");
|
||||
assert_eq!(calls[0].arguments.as_ref(), "{\"path\":\"a.rs\"}");
|
||||
let reasoning_item = response
|
||||
.reasoning_items()
|
||||
.next()
|
||||
.expect("reasoning sibling preserved");
|
||||
let kigi_sampling_types::rs::SummaryPart::SummaryText(t) = &reasoning_item.summary[0];
|
||||
assert_eq!(t.text, "let me think harder");
|
||||
// Choice-level usage + top-level cached_tokens both absorbed.
|
||||
let usage = response.usage.as_ref().expect("usage from choice");
|
||||
assert_eq!(usage.prompt_tokens, 100);
|
||||
assert_eq!(usage.completion_tokens, 20);
|
||||
assert_eq!(usage.cached_prompt_tokens, 60);
|
||||
}
|
||||
other => panic!("expected Completed, got {other:?}"),
|
||||
}
|
||||
}
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// Request surface: bearer auth, kigi UA, kimi_compat body adaptations
|
||||
// ---------------------------------------------------------------------------
|
||||
|
||||
type Captured = Arc<Mutex<Option<(HeaderMap, Value)>>>;
|
||||
|
||||
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
|
||||
async fn request_carries_bearer_kigi_ua_and_kimi_dialect_body() {
|
||||
let captured: Captured = Arc::new(Mutex::new(None));
|
||||
let captured_handler = Arc::clone(&captured);
|
||||
let app = Router::new().route(
|
||||
"/v1/chat/completions",
|
||||
post(move |headers: HeaderMap, body: Bytes| {
|
||||
let captured = Arc::clone(&captured_handler);
|
||||
async move {
|
||||
let body: Value = serde_json::from_slice(&body).unwrap();
|
||||
*captured.lock().unwrap() = Some((headers, body));
|
||||
let events = vec![chunk(
|
||||
json!({ "role": "assistant", "content": "ok" }),
|
||||
Some("stop"),
|
||||
None,
|
||||
)];
|
||||
Sse::new(stream::iter(
|
||||
events.into_iter().map(Ok::<_, std::convert::Infallible>),
|
||||
))
|
||||
}
|
||||
}),
|
||||
);
|
||||
let server = MockServer::spawn(app).await;
|
||||
let (event_tx, mut event_rx) = mpsc::unbounded_channel();
|
||||
let handle = SamplerActor::spawn(
|
||||
test_config(server.base_url()),
|
||||
RetryPolicy::default(),
|
||||
event_tx,
|
||||
);
|
||||
|
||||
// Multi-turn conversation exercising every request-side adaptation:
|
||||
// reasoning folded onto the assistant, an empty-content tool-call turn,
|
||||
// and an enum-only tool schema property.
|
||||
let request = ConversationRequest {
|
||||
items: vec![
|
||||
ConversationItem::User(UserItem {
|
||||
content: vec![ContentPart::Text {
|
||||
text: Arc::<str>::from("read a.rs"),
|
||||
}],
|
||||
synthetic_reason: None,
|
||||
..Default::default()
|
||||
}),
|
||||
ConversationItem::Reasoning(synthesized_reasoning_item("planning the read")),
|
||||
ConversationItem::Assistant(AssistantItem {
|
||||
content: Arc::<str>::from(""),
|
||||
tool_calls: vec![ToolCall {
|
||||
id: Arc::<str>::from("call_9"),
|
||||
name: "read_file".into(),
|
||||
arguments: Arc::<str>::from("{\"path\":\"a.rs\"}"),
|
||||
}],
|
||||
model_id: Some("kimi-for-coding".into()),
|
||||
model_fingerprint: None,
|
||||
reasoning_effort: None,
|
||||
}),
|
||||
ConversationItem::ToolResult(ToolResultItem {
|
||||
tool_call_id: "call_9".into(),
|
||||
content: Arc::<str>::from("fn main() {}"),
|
||||
images: vec![],
|
||||
}),
|
||||
],
|
||||
tools: vec![ToolSpec {
|
||||
name: "read_file".into(),
|
||||
description: Some("Read a file".into()),
|
||||
parameters: json!({
|
||||
"type": "object",
|
||||
"properties": {
|
||||
// Enum-only property: Moonshot 400s without a `type`.
|
||||
"mode": { "enum": ["full", "head"] },
|
||||
"path": { "type": "string" }
|
||||
}
|
||||
}),
|
||||
}],
|
||||
reasoning_effort: Some(ReasoningEffort::High),
|
||||
..Default::default()
|
||||
};
|
||||
handle.submit(RequestId::from("req-wire"), request);
|
||||
let events = drain_until_terminal(&mut event_rx, Duration::from_secs(30)).await;
|
||||
server.shutdown();
|
||||
assert!(matches!(
|
||||
events.last().unwrap(),
|
||||
SamplingEvent::Completed { .. }
|
||||
));
|
||||
|
||||
let (headers, body) = captured.lock().unwrap().take().expect("request captured");
|
||||
|
||||
// -- Auth: plain bearer, nothing else (PRD F3).
|
||||
assert_eq!(
|
||||
headers.get("authorization").unwrap().to_str().unwrap(),
|
||||
"Bearer test-kimi-key"
|
||||
);
|
||||
for gone in [
|
||||
"x-xai-token-auth",
|
||||
"x-authenticateresponse",
|
||||
"x-grok-conv-id",
|
||||
"x-grok-req-id",
|
||||
"x-grok-model-override",
|
||||
"x-grok-session-id",
|
||||
"x-grok-agent-id",
|
||||
"x-grok-client-identifier",
|
||||
"x-grok-client-version",
|
||||
"x-grok-deployment-id",
|
||||
"x-grok-user-id",
|
||||
"x-grok-client-mode",
|
||||
] {
|
||||
assert!(
|
||||
headers.get(gone).is_none(),
|
||||
"xAI proxy marker header must not be sent: {gone}"
|
||||
);
|
||||
}
|
||||
|
||||
// -- User-Agent: kigi/{version} (os; arch).
|
||||
let ua = headers.get("user-agent").unwrap().to_str().unwrap();
|
||||
let expected_prefix = format!("kigi/{}", kigi_version::VERSION);
|
||||
assert!(
|
||||
ua.starts_with(&expected_prefix),
|
||||
"UA must start with {expected_prefix}, got {ua}"
|
||||
);
|
||||
|
||||
// -- Streaming fields exactly as kimi-cli sends them (kimi.py:174-181).
|
||||
assert_eq!(body["stream"], json!(true));
|
||||
assert_eq!(body["stream_options"], json!({ "include_usage": true }));
|
||||
|
||||
// -- Thinking mapping (kimi.py:214-223): effort → thinking, no
|
||||
// reasoning_effort on the wire.
|
||||
assert_eq!(body["thinking"], json!({ "type": "enabled" }));
|
||||
assert_eq!(body.get("reasoning_effort"), None);
|
||||
|
||||
// -- Message adaptations.
|
||||
let messages = body["messages"].as_array().unwrap();
|
||||
let assistant = messages
|
||||
.iter()
|
||||
.find(|m| m["role"] == "assistant")
|
||||
.expect("assistant turn present");
|
||||
assert_eq!(
|
||||
assistant.get("content"),
|
||||
None,
|
||||
"empty tool-call content dropped (kimi.py:339-350)"
|
||||
);
|
||||
assert_eq!(assistant.get("model_id"), None, "grok extension dropped");
|
||||
assert_eq!(
|
||||
assistant["reasoning_content"],
|
||||
json!("planning the read"),
|
||||
"reasoning folded onto the assistant turn (kimi.py:351-352)"
|
||||
);
|
||||
assert_eq!(assistant["tool_calls"][0]["id"], json!("call_9"));
|
||||
|
||||
// -- Tool schema normalization (kosong jsonschema.py:88-142).
|
||||
let props = &body["tools"][0]["function"]["parameters"]["properties"];
|
||||
assert_eq!(props["mode"]["type"], json!("string"));
|
||||
assert_eq!(props["path"]["type"], json!("string"));
|
||||
}
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// 429 with Retry-After
|
||||
// ---------------------------------------------------------------------------
|
||||
|
||||
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
|
||||
async fn rate_limit_honors_retry_after_then_succeeds() {
|
||||
let counter = Arc::new(AtomicU32::new(0));
|
||||
let counter_handler = Arc::clone(&counter);
|
||||
let app = Router::new().route(
|
||||
"/v1/chat/completions",
|
||||
post(move || {
|
||||
let counter = Arc::clone(&counter_handler);
|
||||
async move {
|
||||
let n = counter.fetch_add(1, Ordering::SeqCst);
|
||||
if n == 0 {
|
||||
// Moonshot-shaped 429 body + standard Retry-After.
|
||||
let mut headers = HeaderMap::new();
|
||||
headers.insert("retry-after", "1".parse().unwrap());
|
||||
Err::<Sse<_>, (StatusCode, HeaderMap, String)>((
|
||||
StatusCode::TOO_MANY_REQUESTS,
|
||||
headers,
|
||||
json!({ "error": {
|
||||
"message": "Your account is rate limited",
|
||||
"type": "rate_limit_reached_error"
|
||||
}})
|
||||
.to_string(),
|
||||
))
|
||||
} else {
|
||||
let events = vec![chunk(
|
||||
json!({ "role": "assistant", "content": "after limit" }),
|
||||
Some("stop"),
|
||||
None,
|
||||
)];
|
||||
Ok(Sse::new(stream::iter(
|
||||
events.into_iter().map(Ok::<_, std::convert::Infallible>),
|
||||
)))
|
||||
}
|
||||
}
|
||||
}),
|
||||
);
|
||||
let server = MockServer::spawn(app).await;
|
||||
let (event_tx, mut event_rx) = mpsc::unbounded_channel();
|
||||
let handle = SamplerActor::spawn(
|
||||
test_config(server.base_url()),
|
||||
RetryPolicy::default(),
|
||||
event_tx,
|
||||
);
|
||||
|
||||
let started = std::time::Instant::now();
|
||||
handle.submit(RequestId::from("req-ra"), user_request("hi"));
|
||||
let events = drain_until_terminal(&mut event_rx, Duration::from_secs(30)).await;
|
||||
let elapsed = started.elapsed();
|
||||
server.shutdown();
|
||||
|
||||
// A Retrying event carried the classified rate-limit.
|
||||
assert!(events.iter().any(|e| matches!(
|
||||
e,
|
||||
SamplingEvent::Retrying { kind, .. }
|
||||
if *kind == kigi_sampler::SamplingErrorKind::RateLimited
|
||||
)));
|
||||
match events.last().unwrap() {
|
||||
SamplingEvent::Completed { response, .. } => {
|
||||
assert_eq!(
|
||||
response.assistant().unwrap().content.as_ref(),
|
||||
"after limit"
|
||||
);
|
||||
}
|
||||
other => panic!("expected Completed after Retry-After wait, got {other:?}"),
|
||||
}
|
||||
assert_eq!(counter.load(Ordering::SeqCst), 2, "exactly one retry");
|
||||
// Retry-After: 1 replaces the ~2s jittered exponential backoff. The wait
|
||||
// must be at least the advertised second (and clearly less than the
|
||||
// exhaust-path 30s timeout).
|
||||
assert!(
|
||||
elapsed >= Duration::from_secs(1),
|
||||
"waited less than Retry-After: {elapsed:?}"
|
||||
);
|
||||
}
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// Mid-stream network drop → retry → recovery
|
||||
// ---------------------------------------------------------------------------
|
||||
|
||||
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
|
||||
async fn mid_stream_drop_recovers_via_retry() {
|
||||
let counter = Arc::new(AtomicU32::new(0));
|
||||
let counter_handler = Arc::clone(&counter);
|
||||
let app = Router::new().route(
|
||||
"/v1/chat/completions",
|
||||
post(move || {
|
||||
let counter = Arc::clone(&counter_handler);
|
||||
async move {
|
||||
let n = counter.fetch_add(1, Ordering::SeqCst);
|
||||
if n == 0 {
|
||||
// First attempt: a partial chunk, then the connection
|
||||
// dies mid-body (simulated network drop).
|
||||
let events: Vec<Result<Event, std::io::Error>> = vec![
|
||||
Ok(chunk(
|
||||
json!({ "role": "assistant", "content": "partial" }),
|
||||
None,
|
||||
None,
|
||||
)),
|
||||
Err(std::io::Error::new(
|
||||
std::io::ErrorKind::ConnectionReset,
|
||||
"connection reset by peer",
|
||||
)),
|
||||
];
|
||||
Sse::new(stream::iter(events))
|
||||
} else {
|
||||
let events: Vec<Result<Event, std::io::Error>> = vec![Ok(chunk(
|
||||
json!({ "role": "assistant", "content": "recovered" }),
|
||||
Some("stop"),
|
||||
None,
|
||||
))];
|
||||
Sse::new(stream::iter(events))
|
||||
}
|
||||
}
|
||||
}),
|
||||
);
|
||||
let server = MockServer::spawn(app).await;
|
||||
let (event_tx, mut event_rx) = mpsc::unbounded_channel();
|
||||
let handle = SamplerActor::spawn(
|
||||
test_config(server.base_url()),
|
||||
RetryPolicy::default(),
|
||||
event_tx,
|
||||
);
|
||||
|
||||
handle.submit(RequestId::from("req-drop"), user_request("hi"));
|
||||
let events = drain_until_terminal(&mut event_rx, Duration::from_secs(60)).await;
|
||||
server.shutdown();
|
||||
|
||||
assert!(
|
||||
events
|
||||
.iter()
|
||||
.any(|e| matches!(e, SamplingEvent::Retrying { .. })),
|
||||
"mid-stream drop must go through the retry loop"
|
||||
);
|
||||
match events.last().unwrap() {
|
||||
SamplingEvent::Completed { response, .. } => {
|
||||
// The poisoned partial attempt is discarded; only the fresh
|
||||
// attempt's content survives.
|
||||
assert_eq!(response.assistant().unwrap().content.as_ref(), "recovered");
|
||||
}
|
||||
other => panic!("expected Completed after recovery, got {other:?}"),
|
||||
}
|
||||
assert!(
|
||||
counter.load(Ordering::SeqCst) >= 2,
|
||||
"server hit at least twice"
|
||||
);
|
||||
}
|
||||
Reference in New Issue
Block a user