ZacharyZhang-NY/Kigi-CLI
1
0
Fork 0
Code Issues Pull Requests Actions 13 Packages Projects Releases Wiki Activity

F3: Kimi inference pipeline + full grok cloud-surface excision

Browse Source 
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:
ZacharyZhang-NY
2026-07-17 16:05:51 -04:00
parent fe1f885bb3
commit ea0ce9d15f
231 changed files with 4730 additions and 26358 deletions
Download Patch File Download Diff File Expand all files Collapse all files
crates/codegen/kigi-sampler/Cargo.toml
+1 -1
View File
@@ -3,7 +3,7 @@ license = "Apache-2.0"
name = "kigi-sampler"
version.workspace = true
edition.workspace = true
description = "Actor-based sampling/inference layer for xAI grok (HTTP streaming + retry, no shell coupling)"
description = "Actor-based sampling/inference layer for the Kimi APIs (HTTP streaming + retry, no shell coupling)"
[dependencies]
# Internal
crates/codegen/kigi-sampler/src/actor/request_task.rs
+6 -10
View File
@@ -373,16 +373,13 @@ async fn apply_retry_decision(
RetryDecision::Fatal(fatal_err) => {
// Emit only on true budget exhaustion (hit the retry / rate-limit
// cap), mirroring `classify_error`'s Fatal conditions — NOT on a
// server `x-should-retry: false` or a non-retryable error, which
// are also Fatal but are not "exhausted".
// non-retryable error, which is also Fatal but not "exhausted".
let next_attempt = *retry_count + 1;
let server_said_stop = matches!(err.should_retry_header(), Some(false));
let budget_exhausted = !server_said_stop
&& if err.is_rate_limited() {
next_attempt >= max_retries.min(rate_limit_threshold)
} else {
err.is_retryable() && next_attempt >= max_retries
};
let budget_exhausted = if err.is_rate_limited() {
next_attempt >= max_retries.min(rate_limit_threshold)
} else {
err.is_retryable() && next_attempt >= max_retries
};
if budget_exhausted {
let exhausted_span = tracing::info_span!(
"http.retries_exhausted",
@@ -609,7 +606,6 @@ fn synthesize_from_info(info: &SamplingErrorInfo) -> SamplingError {
message: info.message.clone(),
model_metadata: info.model_metadata.clone(),
retry_after_secs: info.retry_after_secs,
should_retry: None,
}
}
SamplingErrorKind::EmptyResponse => {
crates/codegen/kigi-sampler/src/actor/state.rs
-4
View File
@@ -97,10 +97,6 @@ mod tests {
idle_timeout_secs: None,
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,
crates/codegen/kigi-sampler/src/client.rs
+58 -263
View File
@@ -1,16 +1,20 @@
//! HTTP client for the xAI sampling APIs.
//! HTTP client for the sampling APIs.
//!
//! Owns the `reqwest::Client`, default request headers, and per-method
//! defaults. Talks to three backend shapes:
//!
//! * Chat Completions (`/chat/completions`)
//! * Responses API (`/responses`)
//! * Anthropic Messages API (`/messages`)
//! * Chat Completions (`/chat/completions`) — the Kimi dialect; both
//! product channels (subscription OAuth, Moonshot API key) ride it.
//! Kimi-specific request deviations are absorbed by [`crate::kimi_compat`].
//! * Responses API (`/responses`) — kept for custom providers.
//! * Anthropic Messages API (`/messages`) — kept for custom providers.
//!
//! All trace-upload and URL-based header injection is intentionally
//! *not* here. The session is responsible for putting any per-request
//! headers (proxy auth, OTel context, etc.)
//! into [`SamplerConfig::extra_headers`] before constructing the client.
//! Auth on the wire is a plain `Authorization: Bearer {token}` (or
//! `x-api-key` for Anthropic-scheme custom providers). All URL-based
//! header injection is intentionally *not* here. The session is
//! responsible for putting any per-request headers (device identity,
//! OTel context, etc.) into [`SamplerConfig::extra_headers`] before
//! constructing the client.
use eventsource_stream::Eventsource;
use futures_util::StreamExt;
@@ -33,46 +37,10 @@ use crate::config::{AuthScheme, OriginClientInfo, SamplerConfig};
// Re-export ApiBackend from the shared types crate for downstream callers.
pub use kigi_sampling_types::ApiBackend;
/// Process-level fallback for the `x-grok-client-identifier` header.
const DEFAULT_CLIENT_IDENTIFIER: &str = "grok-shell";
/// Product identifier baked into User-Agent strings.
const AGENT_PRODUCT: &str = "grok-shell";
/// Product identifier baked into User-Agent strings (PRD F3: `kigi/{version}`).
const AGENT_PRODUCT: &str = "kigi";
const ANTHROPIC_DEFAULT_MAX_TOKENS: u32 = 128_000;
/// Per-request `x-grok-*` headers. Optional fields are skipped when empty/`None`.
struct GrokRequestHeaders<'a> {
conv_id: &'a str,
req_id: &'a str,
model_id: &'a str,
session_id: &'a str,
turn_idx: Option<&'a str>,
agent_id: &'a str,
deployment_id: Option<&'a str>,
user_id: Option<&'a str>,
}
impl GrokRequestHeaders<'_> {
fn apply(&self, builder: reqwest::RequestBuilder) -> reqwest::RequestBuilder {
let mut b = builder
.header("x-grok-conv-id", self.conv_id)
.header("x-grok-req-id", self.req_id)
.header("x-grok-model-override", self.model_id)
.header("x-grok-session-id", self.session_id)
.header("x-grok-agent-id", self.agent_id);
if let Some(idx) = self.turn_idx {
b = b.header("x-grok-turn-idx", idx);
}
if let Some(id) = self.deployment_id.filter(|s| !s.is_empty()) {
b = b.header("x-grok-deployment-id", id);
}
if let Some(id) = self.user_id.filter(|s| !s.is_empty()) {
b = b.header("x-grok-user-id", id);
}
b
}
}
/// Parse the `Retry-After` response header as delta-seconds.
/// Our inference backends only emit integer seconds (never HTTP-date),
/// so we only handle that form. HTTP-dates silently return `None` and
@@ -211,21 +179,6 @@ fn extract_retry_after(headers: &reqwest::header::HeaderMap) -> Option<u64> {
.map(|s| s.min(120))
}
fn extract_should_retry(headers: &reqwest::header::HeaderMap) -> Option<bool> {
headers
.get("x-should-retry")
.and_then(|v| v.to_str().ok())
.and_then(|s| {
if s.eq_ignore_ascii_case("true") {
Some(true)
} else if s.eq_ignore_ascii_case("false") {
Some(false)
} else {
None // unknown value — treat as absent
}
})
}
fn extract_model_metadata(headers: &reqwest::header::HeaderMap) -> Option<ResponseModelMetadata> {
let context_window = headers
.get("x-grok-context-window")
@@ -444,45 +397,11 @@ impl SamplingClient {
headers.insert(header_name, header_value);
}
// Add x-grok-client-version header for version gating at the proxy.
if let Some(client_version) = config.client_version.as_ref()
&& let Ok(header_value) = HeaderValue::from_str(client_version)
{
headers.insert(
HeaderName::from_static("x-grok-client-version"),
header_value,
);
}
if let Some(deployment_id) = config.deployment_id.as_ref()
&& let Ok(header_value) = HeaderValue::from_str(deployment_id)
{
headers.insert(
HeaderName::from_static("x-grok-deployment-id"),
header_value,
);
}
if let Some(user_id) = config.user_id.as_ref()
&& let Ok(header_value) = HeaderValue::from_str(user_id)
{
headers.insert(HeaderName::from_static("x-grok-user-id"), header_value);
}
{
let client_id = config
.client_identifier
.clone()
.unwrap_or_else(|| DEFAULT_CLIENT_IDENTIFIER.to_string());
if let Ok(header_value) = HeaderValue::from_str(&client_id) {
headers.insert(
HeaderName::from_static("x-grok-client-identifier"),
header_value,
);
}
}
// Always set User-Agent: per-session origin if available, else fallback.
// This and `extra_headers` are the only client-identity signals on the
// wire — the old xAI proxy's `x-grok-*` marker headers are gone
// (PRD F3: auth is a plain bearer; kimi-cli sends only User-Agent
// plus the OAuth device headers, src/kimi_cli/llm.py:317-323).
{
let ua_string = match config.origin_client.as_ref() {
Some(origin) => user_agent_string_for(origin),
@@ -689,13 +608,7 @@ impl SamplingClient {
}
/// Build request headers string for error messages (redacting sensitive values).
fn format_request_headers(
&self,
x_grok_conv_id: &str,
x_grok_req_id: &str,
model_id: &str,
include_accept: bool,
) -> Vec<String> {
fn format_request_headers(&self, include_accept: bool) -> Vec<String> {
let mut req_headers: Vec<String> = self
.default_headers
.iter()
@@ -704,9 +617,6 @@ impl SamplingClient {
})
.collect();
req_headers.push(Self::format_header("x-grok-conv-id", x_grok_conv_id));
req_headers.push(Self::format_header("x-grok-req-id", x_grok_req_id));
req_headers.push(Self::format_header("x-grok-model-override", model_id));
if include_accept {
req_headers.push(Self::format_header("accept", "text/event-stream"));
}
@@ -799,7 +709,6 @@ impl SamplingClient {
let status = response.status();
let model_metadata = extract_model_metadata(response.headers());
let retry_after_secs = extract_retry_after(response.headers());
let should_retry = extract_should_retry(response.headers());
let bytes = response.bytes().await?;
if !status.is_success() {
@@ -816,7 +725,6 @@ impl SamplingClient {
message,
model_metadata,
retry_after_secs,
should_retry,
});
}
@@ -841,8 +749,6 @@ impl SamplingClient {
request: ChatCompletionRequest,
) -> Result<ChatCompletionResponse> {
let payload = self.apply_defaults(request)?;
let x_grok_conv_id = &payload.x_grok_conv_id.clone().unwrap_or_default();
let x_grok_req_id = &payload.x_grok_req_id.clone().unwrap_or_default();
let model_id = payload.model.clone().unwrap_or_default();
tracing::debug!(
@@ -851,19 +757,18 @@ impl SamplingClient {
"Sending chat completion request"
);
let grok_headers = GrokRequestHeaders {
conv_id: x_grok_conv_id,
req_id: x_grok_req_id,
model_id: &model_id,
session_id: payload.x_grok_session_id.as_deref().unwrap_or_default(),
turn_idx: payload.x_grok_turn_idx.as_deref(),
agent_id: payload.x_grok_agent_id.as_deref().unwrap_or_default(),
deployment_id: payload.x_grok_deployment_id.as_deref(),
user_id: payload.x_grok_user_id.as_deref(),
};
let http_request = grok_headers
.apply(self.post(self.endpoint("chat/completions")))
.json(&payload);
// Serialize, then run the Kimi dialect adaptations (single
// adaptation point for every request-side deviation; see
// `crate::kimi_compat`).
let mut request_body = serde_json::to_value(&payload).map_err(|e| {
tracing::error!("Failed to serialize chat/completions request: {}", e);
SamplingError::Serialization(e)
})?;
crate::kimi_compat::adapt_chat_completions_body(&mut request_body);
let http_request = self
.post(self.endpoint("chat/completions"))
.json(&request_body);
let response = http_request.send().await.map_err(|e| {
// Log at debug level; errors are surfaced to the caller.
@@ -894,13 +799,14 @@ impl SamplingClient {
Option<ResponseModelMetadata>,
)> {
let payload = self.apply_defaults(request)?;
let x_grok_conv_id = &payload.x_grok_conv_id.clone().unwrap_or_default();
let x_grok_req_id = &payload.x_grok_req_id.clone().unwrap_or_default();
let model_id = payload.model.clone().unwrap_or_default();
// Wrap the request with streaming fields and serialize once.
// Previously this path serialized twice: first to serde_json::Value
// (to inject `stream` and `stream_options`), then to HTTP body bytes.
// Wrap the request with the streaming fields the Kimi API expects
// (`stream: true` + `stream_options.include_usage: true`, exactly
// what kimi-cli sends —
// packages/kosong/src/kosong/chat_provider/kimi.py:174-181), then
// serialize and run the Kimi dialect adaptations (single adaptation
// point for every request-side deviation; see `crate::kimi_compat`).
let streaming_request = StreamingChatRequest {
inner: &payload,
stream: true,
@@ -908,21 +814,16 @@ impl SamplingClient {
include_usage: true,
},
};
let mut request_body = serde_json::to_value(&streaming_request).map_err(|e| {
tracing::error!("Failed to serialize chat/completions request: {}", e);
SamplingError::Serialization(e)
})?;
crate::kimi_compat::adapt_chat_completions_body(&mut request_body);
let grok_headers = GrokRequestHeaders {
conv_id: x_grok_conv_id,
req_id: x_grok_req_id,
model_id: &model_id,
session_id: payload.x_grok_session_id.as_deref().unwrap_or_default(),
turn_idx: payload.x_grok_turn_idx.as_deref(),
agent_id: payload.x_grok_agent_id.as_deref().unwrap_or_default(),
deployment_id: payload.x_grok_deployment_id.as_deref(),
user_id: payload.x_grok_user_id.as_deref(),
};
let http_request = grok_headers
.apply(self.post(self.endpoint("chat/completions")))
let http_request = self
.post(self.endpoint("chat/completions"))
.header(ACCEPT, HeaderValue::from_static("text/event-stream"))
.json(&streaming_request);
.json(&request_body);
let built_request = http_request.build().map_err(|e| {
tracing::error!("Failed to build HTTP request: {}", e);
@@ -948,7 +849,6 @@ impl SamplingClient {
span.record("success", status.is_success());
let model_metadata = extract_model_metadata(response.headers());
let retry_after_secs = extract_retry_after(response.headers());
let should_retry = extract_should_retry(response.headers());
if !status.is_success() {
if status == reqwest::StatusCode::UNAUTHORIZED {
span.record("error", "unauthorized (401)");
@@ -962,8 +862,7 @@ impl SamplingClient {
)));
}
let req_headers =
self.format_request_headers(x_grok_conv_id, x_grok_req_id, &model_id, true);
let req_headers = self.format_request_headers(true);
let resp_headers = Self::format_response_headers(&response);
let bytes = response.bytes().await?;
let server_message = parse_error_bytes(bytes.as_ref());
@@ -988,7 +887,6 @@ impl SamplingClient {
message,
model_metadata,
retry_after_secs,
should_retry,
});
}
@@ -1111,8 +1009,6 @@ impl SamplingClient {
) -> Result<rs::Response> {
self.apply_response_defaults(&mut request)?;
let x_grok_conv_id = request.x_grok_conv_id.as_deref().unwrap_or_default();
let x_grok_req_id = request.x_grok_req_id.as_deref().unwrap_or_default();
let model_id = request.inner.model.clone().unwrap_or_default();
// The trace field is process-local: it is consumed by upstream
@@ -1123,16 +1019,6 @@ impl SamplingClient {
tracing::debug!("create_response: {:?}", &request);
tracing::debug!("endpoint: {:?}", self.endpoint("responses"));
let grok_headers = GrokRequestHeaders {
conv_id: x_grok_conv_id,
req_id: x_grok_req_id,
model_id: &model_id,
session_id: request.x_grok_session_id.as_deref().unwrap_or_default(),
turn_idx: request.x_grok_turn_idx.as_deref(),
agent_id: request.x_grok_agent_id.as_deref().unwrap_or_default(),
deployment_id: request.x_grok_deployment_id.as_deref(),
user_id: request.x_grok_user_id.as_deref(),
};
let mut request_body = serde_json::to_value(&request.inner).map_err(|e| {
tracing::error!("Failed to serialize responses request: {}", e);
SamplingError::Serialization(e)
@@ -1142,9 +1028,7 @@ impl SamplingClient {
// it in post-serialize. This is the last surviving piece of the
// old raw_output machinery.
kigi_sampling_types::patch_reasoning_text_types(&mut request_body);
let http_request = grok_headers
.apply(self.post(self.endpoint("responses")))
.json(&request_body);
let http_request = self.post(self.endpoint("responses")).json(&request_body);
let response = http_request.send().await.map_err(|e| {
tracing::debug!("HTTP request failed: {}", e);
@@ -1154,7 +1038,6 @@ impl SamplingClient {
let status = response.status();
let model_metadata = extract_model_metadata(response.headers());
let retry_after_secs = extract_retry_after(response.headers());
let should_retry = extract_should_retry(response.headers());
let bytes = response.bytes().await?;
if !status.is_success() {
@@ -1167,8 +1050,7 @@ impl SamplingClient {
)));
}
let req_headers =
self.format_request_headers(x_grok_conv_id, x_grok_req_id, &model_id, false);
let req_headers = self.format_request_headers(false);
let server_message = parse_error_bytes(bytes.as_ref());
let message = self.build_api_error_message(
@@ -1189,7 +1071,6 @@ impl SamplingClient {
message,
model_metadata,
retry_after_secs,
should_retry,
});
}
@@ -1245,8 +1126,6 @@ impl SamplingClient {
// Enable streaming
request.inner.stream = Some(true);
let x_grok_conv_id = request.x_grok_conv_id.as_deref().unwrap_or_default();
let x_grok_req_id = request.x_grok_req_id.as_deref().unwrap_or_default();
let model_id = request.inner.model.clone().unwrap_or_default();
// Drop process-local trace data (see note in `create_response`).
@@ -1258,16 +1137,6 @@ impl SamplingClient {
"Sending responses API stream request"
);
let grok_headers = GrokRequestHeaders {
conv_id: x_grok_conv_id,
req_id: x_grok_req_id,
model_id: &model_id,
session_id: request.x_grok_session_id.as_deref().unwrap_or_default(),
turn_idx: request.x_grok_turn_idx.as_deref(),
agent_id: request.x_grok_agent_id.as_deref().unwrap_or_default(),
deployment_id: request.x_grok_deployment_id.as_deref(),
user_id: request.x_grok_user_id.as_deref(),
};
let extra_raw_tools = std::mem::take(&mut request.extra_raw_tools);
let mut request_body = serde_json::to_value(&request.inner).map_err(|e| {
tracing::error!("Failed to serialize responses request: {}", e);
@@ -1293,8 +1162,8 @@ impl SamplingClient {
.defaults
.doom_loop_recovery
.map(crate::doom_loop::DoomLoopSignalCollector::new);
let mut http_request = grok_headers
.apply(self.post(self.endpoint("responses")))
let mut http_request = self
.post(self.endpoint("responses"))
.header(ACCEPT, HeaderValue::from_static("text/event-stream"));
if doom_loop.is_some() {
// Presence opts in; the server ignores the value.
@@ -1336,9 +1205,7 @@ impl SamplingClient {
}
let model_metadata = extract_model_metadata(response.headers());
let retry_after_secs = extract_retry_after(response.headers());
let should_retry = extract_should_retry(response.headers());
let req_headers =
self.format_request_headers(x_grok_conv_id, x_grok_req_id, &model_id, true);
let req_headers = self.format_request_headers(true);
let resp_headers = Self::format_response_headers(&response);
let bytes = response.bytes().await?;
let server_message = parse_error_bytes(bytes.as_ref());
@@ -1363,7 +1230,6 @@ impl SamplingClient {
message,
model_metadata,
retry_after_secs,
should_retry,
});
}
@@ -1480,8 +1346,6 @@ impl SamplingClient {
) -> Result<messages::MessagesResponse> {
self.apply_message_defaults(&mut request)?;
let x_grok_conv_id = request.x_grok_conv_id.as_deref().unwrap_or_default();
let x_grok_req_id = request.x_grok_req_id.as_deref().unwrap_or_default();
let model_id = request.inner.model.clone();
// Drop process-local trace data.
@@ -1490,19 +1354,7 @@ impl SamplingClient {
tracing::debug!("create_message: {:?}", &request.inner);
tracing::debug!("endpoint: {:?}", self.endpoint("messages"));
let grok_headers = GrokRequestHeaders {
conv_id: x_grok_conv_id,
req_id: x_grok_req_id,
model_id: &model_id,
session_id: request.x_grok_session_id.as_deref().unwrap_or_default(),
turn_idx: request.x_grok_turn_idx.as_deref(),
agent_id: request.x_grok_agent_id.as_deref().unwrap_or_default(),
deployment_id: request.x_grok_deployment_id.as_deref(),
user_id: request.x_grok_user_id.as_deref(),
};
let http_request = grok_headers
.apply(self.post(self.endpoint("messages")))
.json(&request.inner);
let http_request = self.post(self.endpoint("messages")).json(&request.inner);
let response = http_request.send().await.map_err(|e| {
tracing::debug!("HTTP request failed: {}", e);
@@ -1512,7 +1364,6 @@ impl SamplingClient {
let status = response.status();
let model_metadata = extract_model_metadata(response.headers());
let retry_after_secs = extract_retry_after(response.headers());
let should_retry = extract_should_retry(response.headers());
let bytes = response.bytes().await?;
if !status.is_success() {
@@ -1525,8 +1376,7 @@ impl SamplingClient {
)));
}
let req_headers =
self.format_request_headers(x_grok_conv_id, x_grok_req_id, &model_id, false);
let req_headers = self.format_request_headers(false);
let server_message = parse_error_bytes(bytes.as_ref());
let message = self.build_api_error_message(
@@ -1547,7 +1397,6 @@ impl SamplingClient {
message,
model_metadata,
retry_after_secs,
should_retry,
});
}
@@ -1593,8 +1442,6 @@ impl SamplingClient {
// Enable streaming
request.inner.stream = Some(true);
let x_grok_conv_id = request.x_grok_conv_id.as_deref().unwrap_or_default();
let x_grok_req_id = request.x_grok_req_id.as_deref().unwrap_or_default();
let model_id = request.inner.model.clone();
// Drop process-local trace data.
@@ -1606,18 +1453,8 @@ impl SamplingClient {
"Sending Messages API stream request"
);
let grok_headers = GrokRequestHeaders {
conv_id: x_grok_conv_id,
req_id: x_grok_req_id,
model_id: &model_id,
session_id: request.x_grok_session_id.as_deref().unwrap_or_default(),
turn_idx: request.x_grok_turn_idx.as_deref(),
agent_id: request.x_grok_agent_id.as_deref().unwrap_or_default(),
deployment_id: request.x_grok_deployment_id.as_deref(),
user_id: request.x_grok_user_id.as_deref(),
};
let http_request = grok_headers
.apply(self.post(self.endpoint("messages")))
let http_request = self
.post(self.endpoint("messages"))
.header(ACCEPT, HeaderValue::from_static("text/event-stream"))
.json(&request.inner);
@@ -1655,9 +1492,7 @@ impl SamplingClient {
}
let model_metadata = extract_model_metadata(response.headers());
let retry_after_secs = extract_retry_after(response.headers());
let should_retry = extract_should_retry(response.headers());
let req_headers =
self.format_request_headers(x_grok_conv_id, x_grok_req_id, &model_id, true);
let req_headers = self.format_request_headers(true);
let resp_headers = Self::format_response_headers(&response);
let bytes = response.bytes().await?;
let server_message = parse_error_bytes(bytes.as_ref());
@@ -1682,7 +1517,6 @@ impl SamplingClient {
message,
model_metadata,
retry_after_secs,
should_retry,
});
}
@@ -2002,7 +1836,6 @@ impl SamplingClient {
message: info.message,
model_metadata: info.model_metadata,
retry_after_secs: info.retry_after_secs,
should_retry: None,
})
}
}
@@ -2031,10 +1864,6 @@ mod tests {
idle_timeout_secs: None,
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,
@@ -2146,40 +1975,6 @@ mod tests {
assert_eq!(extract_retry_after(&headers), None);
}
#[test]
fn extract_should_retry_true() {
let mut headers = reqwest::header::HeaderMap::new();
headers.insert("x-should-retry", "true".parse().unwrap());
assert_eq!(extract_should_retry(&headers), Some(true));
}
#[test]
fn extract_should_retry_true_case_insensitive() {
let mut headers = reqwest::header::HeaderMap::new();
headers.insert("x-should-retry", "TRUE".parse().unwrap());
assert_eq!(extract_should_retry(&headers), Some(true));
}
#[test]
fn extract_should_retry_false() {
let mut headers = reqwest::header::HeaderMap::new();
headers.insert("x-should-retry", "false".parse().unwrap());
assert_eq!(extract_should_retry(&headers), Some(false));
}
#[test]
fn extract_should_retry_unknown_value_is_none() {
let mut headers = reqwest::header::HeaderMap::new();
headers.insert("x-should-retry", "banana".parse().unwrap());
assert_eq!(extract_should_retry(&headers), None);
}
#[test]
fn extract_should_retry_absent_is_none() {
let headers = reqwest::header::HeaderMap::new();
assert_eq!(extract_should_retry(&headers), None);
}
#[test]
fn new_with_minimal_config_succeeds() {
let client = SamplingClient::new(minimal_config()).expect("client should construct");
@@ -2286,8 +2081,8 @@ mod tests {
version: None,
};
let ua = user_agent_string_for(&origin);
// No slash between product and the grok-shell agent product.
assert!(ua.starts_with("my-client grok-shell/"));
// No slash between product and the kigi agent product.
assert!(ua.starts_with("my-client kigi/"));
}
#[test]
crates/codegen/kigi-sampler/src/config.rs
+5 -9
View File
@@ -71,12 +71,12 @@ pub struct SamplerConfig {
// Reasoning effort
pub reasoning_effort: Option<ReasoningEffort>,
// Client identity
/// Client identity for the User-Agent header (`kigi/{version}` plus an
/// optional origin product). The old xAI proxy's identity headers
/// (`x-grok-client-identifier` / `-client-version` / `-deployment-id` /
/// `-user-id`) are gone — User-Agent and `extra_headers` are the only
/// identity signals on the wire.
pub origin_client: Option<OriginClientInfo>,
pub client_identifier: Option<String>,
pub deployment_id: Option<String>,
pub user_id: Option<String>,
pub client_version: Option<String>,
/// Optional hook invoked at every UNAUTHORIZED (401) response
/// site. The sampler passes the bearer that was actually sent on
@@ -146,10 +146,6 @@ impl Default for SamplerConfig {
idle_timeout_secs: None,
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,
crates/codegen/kigi-sampler/src/events.rs
-3
View File
@@ -292,7 +292,6 @@ mod tests {
message: "boom".into(),
model_metadata: None,
retry_after_secs: None,
should_retry: None,
};
let info = SamplingErrorInfo::from(&err);
assert_eq!(info.kind, SamplingErrorKind::Api);
@@ -307,7 +306,6 @@ mod tests {
message: "slow down".into(),
model_metadata: None,
retry_after_secs: Some(15),
should_retry: None,
};
let info = SamplingErrorInfo::from(&err);
assert_eq!(info.kind, SamplingErrorKind::RateLimited);
@@ -326,7 +324,6 @@ mod tests {
..Default::default()
}),
retry_after_secs: None,
should_retry: None,
};
let info = SamplingErrorInfo::from(&err);
assert_eq!(info.kind, SamplingErrorKind::Api);
crates/codegen/kigi-sampler/src/kimi_compat.rs
+424
View File
@@ -0,0 +1,424 @@
//! Kimi (Moonshot) chat/completions request adaptations.
//!
//! The Kimi endpoints are OpenAI-compatible but deviate in a handful of
//! places (PRD F3 Q1). Every request-side deviation is absorbed HERE, in a
//! single adaptation point applied to the serialized chat/completions body
//! just before it is sent — never as scattered special-cases at call sites.
//! Each adaptation cites the kimi-cli source it was derived from
//! (kimi-cli == the authoritative official client; paths are relative to
//! that repository).
//!
//! Response-side deviations live with the wire types themselves
//! (`kigi_sampling_types::Usage::cached_tokens`,
//! `ChatChunkChoice::usage`) and the L2 stream transform
//! (`stream::chat_completions` synthesizes missing tool-call ids).
//!
//! The `ApiBackend::ChatCompletions` backend is the Kimi dialect: both
//! product channels (subscription OAuth and Moonshot API keys) ride it.
//! Custom providers that need vanilla OpenAI semantics for reasoning use
//! the `Responses` backend, which stays available in model configuration.
use serde_json::Value;
/// Adapt a fully-serialized chat/completions request body to the Kimi
/// dialect, in place. Applied by [`crate::SamplingClient`] to both the
/// streaming and non-streaming chat/completions paths.
pub(crate) fn adapt_chat_completions_body(body: &mut Value) {
adapt_thinking(body);
adapt_messages(body);
adapt_tool_schemas(body);
}
/// Map the OpenAI-style `reasoning_effort` knob onto Kimi's `thinking`
/// request field and drop `reasoning_effort` from the wire.
///
/// kimi-cli controls thinking exclusively through the request body's
/// `thinking: {"type": "enabled" | "disabled"}` field
/// (packages/kosong/src/kosong/chat_provider/kimi.py:214-223 `with_thinking`:
/// `"enabled" if effort != "off" else "disabled"`; wired by
/// src/kimi_cli/llm.py:475-481). When no effort is configured, nothing is
/// sent and the server default applies (llm.py:482 "leave as-is").
fn adapt_thinking(body: &mut Value) {
let Some(obj) = body.as_object_mut() else {
return;
};
let Some(effort) = obj.remove("reasoning_effort") else {
return;
};
let enabled = effort.as_str() != Some("none");
obj.insert(
"thinking".to_owned(),
serde_json::json!({ "type": if enabled { "enabled" } else { "disabled" } }),
);
}
/// Message-level adaptations:
///
/// * Drop `model_id` — a grok-build extension recorded on assistant turns;
/// kimi-cli's message serializer sends no such field
/// (packages/kosong/src/kosong/chat_provider/kimi.py:326-353).
/// * Drop `content` from assistant tool-call messages whose visible content
/// is effectively empty. The Kimi-for-Coding compat layer rejects an
/// empty text content part with 400 "text content is empty"; omitting
/// `content` entirely is always accepted
/// (packages/kosong/src/kosong/chat_provider/kimi.py:339-350, with the
/// "effectively empty" predicate at kimi.py:356-362).
fn adapt_messages(body: &mut Value) {
let Some(messages) = body.get_mut("messages").and_then(Value::as_array_mut) else {
return;
};
for message in messages {
let Some(obj) = message.as_object_mut() else {
continue;
};
obj.remove("model_id");
let is_assistant = obj.get("role").and_then(Value::as_str) == Some("assistant");
let has_tool_calls = obj
.get("tool_calls")
.and_then(Value::as_array)
.is_some_and(|calls| !calls.is_empty());
if is_assistant
&& has_tool_calls
&& obj.get("content").is_some_and(is_effectively_empty_content)
{
obj.remove("content");
}
}
}
/// Port of kimi-cli `_is_effectively_empty_content_parts`
/// (packages/kosong/src/kosong/chat_provider/kimi.py:356-362): a bare
/// whitespace-only string, or a block list whose entries are all
/// whitespace-only text blocks. Any non-text block (e.g. an image) makes
/// the content non-empty.
fn is_effectively_empty_content(content: &Value) -> bool {
match content {
Value::String(s) => s.trim().is_empty(),
Value::Array(blocks) => blocks.iter().all(|block| {
block.get("type").and_then(Value::as_str) == Some("text")
&& block
.get("text")
.and_then(Value::as_str)
.is_some_and(|t| t.trim().is_empty())
}),
Value::Null => true,
_ => false,
}
}
/// Moonshot's schema validator rejects tool parameter schemas whose
/// property schemas omit `type` (e.g. enum-only properties exposed by some
/// MCP servers): HTTP 400 "At path 'properties.X': type is not defined".
/// Fill in an inferred `type` locally so such tools keep working. Port of
/// kimi-cli `ensure_property_types`
/// (packages/kosong/src/kosong/utils/jsonschema.py:88-142, applied per tool
/// at packages/kosong/src/kosong/chat_provider/kimi.py:378-388).
fn adapt_tool_schemas(body: &mut Value) {
let Some(tools) = body.get_mut("tools").and_then(Value::as_array_mut) else {
return;
};
for tool in tools {
if let Some(parameters) = tool.pointer_mut("/function/parameters") {
recurse_schema(parameters);
}
}
}
/// JSON Schema keywords that describe a property's shape without a `type`
/// keyword; nodes carrying one are left alone
/// (kosong/utils/jsonschema.py:15-24 `_COMBINATOR_KEYS`).
const COMBINATOR_KEYS: [&str; 8] = [
"anyOf", "oneOf", "allOf", "not", "if", "then", "else", "$ref",
];
/// Walk property-schema positions under `node` (`properties`, `items`,
/// `additionalProperties`, `anyOf`/`oneOf`/`allOf`); `node` itself is a
/// container and is not normalized (kosong/utils/jsonschema.py:114-142).
fn recurse_schema(node: &mut Value) {
let Some(obj) = node.as_object_mut() else {
return;
};
if let Some(props) = obj.get_mut("properties").and_then(Value::as_object_mut) {
for value in props.values_mut() {
normalize_property(value);
}
}
match obj.get_mut("items") {
Some(items @ Value::Object(_)) => normalize_property(items),
Some(Value::Array(items)) => {
for value in items {
normalize_property(value);
}
}
_ => {}
}
if let Some(additional @ Value::Object(_)) = obj.get_mut("additionalProperties") {
normalize_property(additional);
}
for key in ["anyOf", "oneOf", "allOf"] {
if let Some(branches) = obj.get_mut(key).and_then(Value::as_array_mut) {
for value in branches {
normalize_property(value);
}
}
}
}
/// Ensure a property schema declares `type`, then recurse into it
/// (kosong/utils/jsonschema.py:145-162 `_normalize_property`).
fn normalize_property(node: &mut Value) {
let Some(obj) = node.as_object_mut() else {
return;
};
if !obj.contains_key("type") && !COMBINATOR_KEYS.iter().any(|k| obj.contains_key(*k)) {
let inferred = if let Some(Value::Array(values)) = obj.get("enum") {
if values.is_empty() {
infer_type_from_structure(obj)
} else {
infer_type_from_values(values)
}
} else if let Some(constant) = obj.get("const") {
infer_type_from_values(std::slice::from_ref(constant))
} else {
infer_type_from_structure(obj)
};
obj.insert("type".to_owned(), Value::String(inferred.to_owned()));
}
recurse_schema(node);
}
/// Infer `type` from structural keywords when no enum/const is present;
/// defaults to `"string"` only with no structural hints at all
/// (kosong/utils/jsonschema.py:165-215 `_infer_type_from_structure`).
fn infer_type_from_structure(obj: &serde_json::Map<String, Value>) -> &'static str {
const OBJECT_KEYWORDS: [&str; 7] = [
"properties",
"additionalProperties",
"patternProperties",
"propertyNames",
"required",
"minProperties",
"maxProperties",
];
const ARRAY_KEYWORDS: [&str; 6] = [
"items",
"prefixItems",
"minItems",
"maxItems",
"uniqueItems",
"contains",
];
const STRING_KEYWORDS: [&str; 4] = ["minLength", "maxLength", "pattern", "format"];
const NUMERIC_KEYWORDS: [&str; 5] = [
"minimum",
"maximum",
"multipleOf",
"exclusiveMinimum",
"exclusiveMaximum",
];
if OBJECT_KEYWORDS.iter().any(|k| obj.contains_key(*k)) {
"object"
} else if ARRAY_KEYWORDS.iter().any(|k| obj.contains_key(*k)) {
"array"
} else if STRING_KEYWORDS.iter().any(|k| obj.contains_key(*k)) {
"string"
} else if NUMERIC_KEYWORDS.iter().any(|k| obj.contains_key(*k)) {
"number"
} else {
"string"
}
}
/// Infer a `type` from concrete enum/const values: single JSON type wins,
/// `{integer, number}` collapses to `"number"`, any other mix falls back to
/// `"string"` (kosong/utils/jsonschema.py:218-247 `_infer_type_from_values`).
fn infer_type_from_values(values: &[Value]) -> &'static str {
let mut inferred = std::collections::BTreeSet::new();
for value in values {
let ty = match value {
Value::Bool(_) => "boolean",
Value::Number(n) if n.is_i64() || n.is_u64() => "integer",
Value::Number(_) => "number",
Value::String(_) => "string",
Value::Null => "null",
Value::Object(_) => "object",
Value::Array(_) => "array",
};
inferred.insert(ty);
}
if inferred.len() == 1 {
return inferred.pop_first().expect("non-empty set");
}
if inferred == std::collections::BTreeSet::from(["integer", "number"]) {
return "number";
}
"string"
}
#[cfg(test)]
mod tests {
use super::*;
use serde_json::json;
#[test]
fn reasoning_effort_maps_to_kimi_thinking_field() {
let mut body = json!({ "model": "kimi-for-coding", "reasoning_effort": "high" });
adapt_chat_completions_body(&mut body);
assert_eq!(body.get("reasoning_effort"), None);
assert_eq!(body["thinking"], json!({ "type": "enabled" }));
// kimi.py:218: "off" (our ReasoningEffort::None) → disabled.
let mut body = json!({ "reasoning_effort": "none" });
adapt_chat_completions_body(&mut body);
assert_eq!(body["thinking"], json!({ "type": "disabled" }));
// llm.py:482: unset → leave as-is (no `thinking` at all).
let mut body = json!({ "model": "kimi-for-coding" });
adapt_chat_completions_body(&mut body);
assert_eq!(body.get("thinking"), None);
}
#[test]
fn assistant_tool_call_with_empty_content_drops_content() {
let mut body = json!({
"messages": [
{ "role": "user", "content": "hi" },
{
"role": "assistant",
"content": "",
"model_id": "kimi-for-coding",
"tool_calls": [{ "id": "c1", "type": "function",
"function": { "name": "f", "arguments": "{}" } }]
},
]
});
adapt_chat_completions_body(&mut body);
let assistant = &body["messages"][1];
assert_eq!(assistant.get("content"), None, "empty content dropped");
assert_eq!(assistant.get("model_id"), None, "grok extension dropped");
assert!(assistant.get("tool_calls").is_some());
// The user message keeps its content.
assert_eq!(body["messages"][0]["content"], json!("hi"));
}
#[test]
fn assistant_tool_call_with_real_content_keeps_content() {
let mut body = json!({
"messages": [{
"role": "assistant",
"content": "let me check",
"tool_calls": [{ "id": "c1", "type": "function",
"function": { "name": "f", "arguments": "{}" } }]
}]
});
adapt_chat_completions_body(&mut body);
assert_eq!(body["messages"][0]["content"], json!("let me check"));
}
#[test]
fn assistant_without_tool_calls_keeps_empty_content() {
// Only tool-call turns drop content (kimi.py:339-350 guards on
// `message.tool_calls`); a plain empty assistant turn is left alone.
let mut body = json!({
"messages": [{ "role": "assistant", "content": "" }]
});
adapt_chat_completions_body(&mut body);
assert_eq!(body["messages"][0]["content"], json!(""));
}
#[test]
fn empty_text_block_list_counts_as_empty_content() {
let mut body = json!({
"messages": [{
"role": "assistant",
"content": [{ "type": "text", "text": " " }],
"tool_calls": [{ "id": "c1", "type": "function",
"function": { "name": "f", "arguments": "{}" } }]
}]
});
adapt_chat_completions_body(&mut body);
assert_eq!(body["messages"][0].get("content"), None);
}
#[test]
fn image_block_is_not_empty_content() {
let mut body = json!({
"messages": [{
"role": "assistant",
"content": [{ "type": "image_url", "image_url": { "url": "data:x" } }],
"tool_calls": [{ "id": "c1", "type": "function",
"function": { "name": "f", "arguments": "{}" } }]
}]
});
adapt_chat_completions_body(&mut body);
assert!(body["messages"][0].get("content").is_some());
}
#[test]
fn enum_only_property_gains_inferred_type() {
// The Moonshot validator 400s on `{"enum": [...]}` without `type`
// (kosong/utils/jsonschema.py:91-96).
let mut body = json!({
"tools": [{
"type": "function",
"function": {
"name": "search",
"parameters": {
"type": "object",
"properties": {
"mode": { "enum": ["smart", "full"] },
"count": { "enum": [1, 2, 3] },
"ratio": { "enum": [1, 2.5] },
"nested": {
"type": "object",
"properties": { "inner": { "enum": ["a"] } }
},
"combined": { "anyOf": [{ "type": "string" }] }
}
}
}
}]
});
adapt_chat_completions_body(&mut body);
let props = &body["tools"][0]["function"]["parameters"]["properties"];
assert_eq!(props["mode"]["type"], json!("string"));
assert_eq!(props["count"]["type"], json!("integer"));
assert_eq!(props["ratio"]["type"], json!("number"));
assert_eq!(
props["nested"]["properties"]["inner"]["type"],
json!("string")
);
assert_eq!(
props["combined"].get("type"),
None,
"combinator nodes are left alone"
);
}
#[test]
fn structural_keywords_infer_shape_not_string() {
let mut body = json!({
"tools": [{
"type": "function",
"function": {
"name": "t",
"parameters": {
"type": "object",
"properties": {
"obj": { "properties": { "x": { "type": "string" } } },
"arr": { "items": { "type": "string" } },
"num": { "minimum": 0 },
"free": {}
}
}
}
}]
});
adapt_chat_completions_body(&mut body);
let props = &body["tools"][0]["function"]["parameters"]["properties"];
assert_eq!(props["obj"]["type"], json!("object"));
assert_eq!(props["arr"]["type"], json!("array"));
assert_eq!(props["num"]["type"], json!("number"));
assert_eq!(props["free"]["type"], json!("string"));
}
}
crates/codegen/kigi-sampler/src/lib.rs
+2 -1
View File
@@ -1,4 +1,4 @@
//! kigi-sampler - Actor-based sampling layer for xAI grok.
//! kigi-sampler - Actor-based sampling layer for the Kimi inference APIs.
//!
//! This crate extracts the HTTP streaming + retry logic out of
//! `kigi-shell`'s session actor into a standalone, reusable
@@ -24,6 +24,7 @@ pub mod config;
pub mod doom_loop;
pub mod events;
pub mod handle;
mod kimi_compat;
pub mod metrics;
pub mod retry;
pub mod sampling_log;
crates/codegen/kigi-sampler/src/retry.rs
+7 -80
View File
@@ -24,14 +24,10 @@
//! - `Serialization` (response parsing failure)
//! - `MaxTokensTruncation` (by design)
//!
//! **Server hint** (`x-should-retry` header from CCP):
//! - `false` → Fatal immediately, regardless of status code
//! - `true` / absent → falls through to status-code logic above
//!
//! Today CCP's header mirrors the client's `is_retryable()` logic
//! (4xx except 429 = false, 5xx + 429 = true), so no behavior changes
//! on merge. The header enables future CCP-side refinements (e.g.
//! marking content-caused 500s as non-retryable) without client updates.
//! 429 handling honors the standard `Retry-After` response header when
//! present (delta-seconds; see `client::extract_retry_after`), matching
//! the Kimi/Moonshot API. The old xAI proxy's `x-should-retry` hint
//! header was removed with the proxy.
use std::time::Duration;
@@ -170,22 +166,6 @@ pub fn classify_error(
return RetryDecision::RetryWithImageStrip;
}
// Server explicitly said don't retry (x-should-retry: false).
// Trust the server — it knows if the error is request-content-caused
// (e.g. malformed tool call in conversation history) vs transient.
//
// x-should-retry: true is intentionally NOT handled here — we only
// use the header to suppress retries (false), not to force them
// (true). Forcing retries on non-retryable status codes could
// amplify failures. true falls through to existing status-code logic.
//
// Checked AFTER image-strip guards: image stripping changes the
// request payload, so a server "don't retry" on the original
// request doesn't apply to the stripped request.
if let Some(false) = err.should_retry_header() {
return RetryDecision::Fatal(clone_error(err));
}
// Context-window / size overflow is deterministic — re-sending the same (or
// larger) payload always fails — so never retry it, whatever status the backend
// used (in-stream `ResponseError`→500, HTTP 400/500, OpenAI/Anthropic variants).
@@ -401,13 +381,11 @@ pub(crate) fn clone_error(err: &SamplingError) -> SamplingError {
message,
model_metadata,
retry_after_secs,
should_retry,
} => SamplingError::Api {
status: *status,
message: message.clone(),
model_metadata: model_metadata.clone(),
retry_after_secs: *retry_after_secs,
should_retry: *should_retry,
},
SamplingError::EventStreamError(msg) => SamplingError::EventStreamError(msg.clone()),
SamplingError::StreamError {
@@ -445,7 +423,6 @@ mod tests {
message: message.to_string(),
model_metadata: None,
retry_after_secs: None,
should_retry: None,
}
}
@@ -455,7 +432,6 @@ mod tests {
message: "x".to_string(),
model_metadata: None,
retry_after_secs: Some(retry_after),
should_retry: None,
}
}
@@ -757,31 +733,15 @@ mod tests {
assert!(s.contains("240s"));
}
#[test]
fn should_retry_false_overrides_retryable_status() {
let err = SamplingError::Api {
status: StatusCode::INTERNAL_SERVER_ERROR,
message: "boom".into(),
model_metadata: None,
retry_after_secs: None,
should_retry: Some(false),
};
assert!(matches!(
classify_error(&err, 0, 15, RATE_LIMIT_RETRY_THRESHOLD),
RetryDecision::Fatal(_)
));
}
#[test]
fn context_length_overflow_is_fatal_even_as_500() {
// The backend streams a size overflow as a ResponseError that becomes a 500 with no
// should_retry hint; without the context-length check it would retry the full budget.
// The backend streams a size overflow as a ResponseError that becomes a 500;
// without the context-length check it would retry the full budget.
let err = SamplingError::Api {
status: StatusCode::INTERNAL_SERVER_ERROR,
message: "none: The prompt is too long for this model's context window.".into(),
model_metadata: None,
retry_after_secs: None,
should_retry: None,
};
assert!(matches!(
classify_error(&err, 0, 15, RATE_LIMIT_RETRY_THRESHOLD),
@@ -790,28 +750,12 @@ mod tests {
}
#[test]
fn should_retry_true_falls_through_to_existing_logic() {
fn api_500_first_failure_retries_with_client_rebuild() {
let err = SamplingError::Api {
status: StatusCode::INTERNAL_SERVER_ERROR,
message: "boom".into(),
model_metadata: None,
retry_after_secs: None,
should_retry: Some(true),
};
assert!(matches!(
classify_error(&err, 0, 15, RATE_LIMIT_RETRY_THRESHOLD),
RetryDecision::RetryWithClientRebuild { .. }
));
}
#[test]
fn should_retry_absent_falls_through() {
let err = SamplingError::Api {
status: StatusCode::INTERNAL_SERVER_ERROR,
message: "boom".into(),
model_metadata: None,
retry_after_secs: None,
should_retry: None,
};
assert!(matches!(
classify_error(&err, 0, 15, RATE_LIMIT_RETRY_THRESHOLD),
@@ -836,21 +780,4 @@ mod tests {
}
}
}
#[test]
fn should_retry_false_on_429_is_fatal() {
// Server says don't retry, even though 429 is normally retryable.
// should_retry check runs before rate-limit check.
let err = SamplingError::Api {
status: StatusCode::TOO_MANY_REQUESTS,
message: "rate limited".into(),
model_metadata: None,
retry_after_secs: Some(10),
should_retry: Some(false),
};
assert!(matches!(
classify_error(&err, 0, 15, RATE_LIMIT_RETRY_THRESHOLD),
RetryDecision::Fatal(_)
));
}
}
crates/codegen/kigi-sampler/src/stream/chat_completions.rs
+36 -5
View File
@@ -128,7 +128,16 @@ pub fn stream_chat_completions<'a>(
first_chunk_seen = true;
}
if let Some(u) = chunk.usage.clone() {
// Kimi/Moonshot deviation: usage may ride inside a choice instead
// of (or in addition to) the chunk's top-level `usage`. Same
// fallback as kimi-cli's `extract_usage_from_chunk`
// (packages/kosong/src/kosong/chat_provider/kimi.py:522-533):
// top-level wins, else the first choice carrying one.
let chunk_usage = chunk
.usage
.clone()
.or_else(|| chunk.choices.iter().find_map(|c| c.usage.clone()));
if let Some(u) = chunk_usage {
// Wire cost is cumulative for the response, so last-write-wins.
// Never clobber a known cost with missing/unreported.
let chunk_cost = kigi_sampling_types::reported_cost_ticks(u.cost_in_usd_ticks);
@@ -247,10 +256,27 @@ pub fn stream_chat_completions<'a>(
// ── Build the final response ─────────────────────────────────
let tool_calls: Vec<ToolCall> = tool_call_acc
.into_values()
.map(|(id, name, arguments)| ToolCall {
id: std::sync::Arc::<str>::from(id),
name,
arguments: std::sync::Arc::<str>::from(arguments),
.map(|(id, name, arguments)| {
// Kimi/Moonshot deviation: tool-call deltas may omit `id`.
// Synthesize one so the tool-result round-trip stays keyed,
// exactly like kimi-cli (`id=tool_call.id or str(uuid.uuid4())`,
// packages/kosong/src/kosong/chat_provider/kimi.py:505).
let id = if id.is_empty() {
let synthesized = uuid::Uuid::new_v4().to_string();
tracing::debug!(
tool_name = %name,
synthesized_id = %synthesized,
"tool-call delta carried no id; synthesized one"
);
synthesized
} else {
id
};
ToolCall {
id: std::sync::Arc::<str>::from(id),
name,
arguments: std::sync::Arc::<str>::from(arguments),
}
})
.collect();
@@ -329,6 +355,7 @@ mod tests {
index: i as u32,
delta,
finish_reason: None,
usage: None,
})
.collect(),
usage: None,
@@ -664,6 +691,7 @@ mod tests {
prompt_tokens: 100,
completion_tokens: 50,
total_tokens: 150,
cached_tokens: None,
prompt_tokens_details: None,
completion_tokens_details: None,
cost_in_usd_ticks: None,
@@ -704,6 +732,7 @@ mod tests {
prompt_tokens: 10,
completion_tokens: 5,
total_tokens: 15,
cached_tokens: None,
prompt_tokens_details: None,
completion_tokens_details: None,
cost_in_usd_ticks: wire,
@@ -737,6 +766,7 @@ mod tests {
prompt_tokens: 10,
completion_tokens: 5,
total_tokens: 15,
cached_tokens: None,
prompt_tokens_details: None,
completion_tokens_details: None,
cost_in_usd_ticks: Some(99),
@@ -746,6 +776,7 @@ mod tests {
prompt_tokens: 12,
completion_tokens: 6,
total_tokens: 18,
cached_tokens: None,
prompt_tokens_details: None,
completion_tokens_details: None,
cost_in_usd_ticks: Some(0),
crates/codegen/kigi-sampler/src/stream/collect.rs
+2
View File
@@ -90,6 +90,7 @@ mod tests {
tool_call_id: None,
},
finish_reason: None,
usage: None,
}],
usage: None,
system_fingerprint: None,
@@ -106,6 +107,7 @@ mod tests {
index: 0,
delta: ChatChunkDelta::default(),
finish_reason: Some(FinishReason::Stop),
usage: None,
}],
usage: None,
system_fingerprint: None,
crates/codegen/kigi-sampler/src/stream/messages.rs
-1
View File
@@ -429,7 +429,6 @@ pub fn stream_messages<'a>(
message: error_message,
model_metadata: None,
retry_after_secs: None,
should_retry: None,
};
yield SamplingEvent::Failed {
request_id: request_id.clone(),
crates/codegen/kigi-sampler/src/stream/responses.rs
-3
View File
@@ -299,7 +299,6 @@ pub fn stream_responses<'a>(
message: error_message,
model_metadata: None,
retry_after_secs: None,
should_retry: None,
};
yield SamplingEvent::Failed {
request_id: request_id.clone(),
@@ -316,7 +315,6 @@ pub fn stream_responses<'a>(
message: error_message,
model_metadata: None,
retry_after_secs: None,
should_retry: None,
};
yield SamplingEvent::Failed {
request_id: request_id.clone(),
@@ -419,7 +417,6 @@ pub fn stream_responses<'a>(
.to_string(),
model_metadata: None,
retry_after_secs: None,
should_retry: None,
};
yield SamplingEvent::Failed {
request_id: request_id.clone(),
crates/codegen/kigi-sampler/tests/test_actor.rs
-4
View File
@@ -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,
crates/codegen/kigi-sampler/tests/test_kimi_wire.rs
+564
View File
@@ -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"
);
}