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

docs(comments): rewrite comments across all crates to the guidelines

Browse Source 
Sweep every first-party crate source (1956 .rs files) to the project comment
guidelines: delete redundant restatements, decorative banners, change
narration, and end-of-line comments; keep and tighten the crucial ones
(invariants, bug rationale, SAFETY blocks, ported-source attribution).

No functional code changed. Every edit is proven comment-only against the
prior tree by a comment-stripping lexer (string/char/raw-string aware) plus a
separate doctest-fence check. Where removing a comment made rustfmt or clippy
want to re-lay-out adjacent code, the minimal triggering comment is restored so
code tokens stay byte-identical.

Gates green: cargo fmt --all --check (0 diffs), cargo check and cargo clippy
--workspace --all-targets (0 warnings).

Adds scripts/check_codegen_comment_guidelines.py — the enforcement gate for
these guidelines (flags banners, end-of-line comments, change narration, and
commented-out code).
This commit is contained in:
ZacharyZhang-NY
2026-07-23 16:55:39 -04:00
parent ff0fb56c67
commit a02b555e66
1458 changed files with 10729 additions and 21750 deletions
Download Patch File Download Diff File Expand all files Collapse all files
crates/codegen/kigi-sampler/src/actor/mod.rs
+2 -1
View File
@@ -81,7 +81,8 @@ impl SamplerActor {
cmd = self.cmd_rx.recv() => {
match cmd {
Some(cmd) => self.handle_command(cmd),
None => break, // all handles dropped
// all handles dropped
None => break,
}
}
}
crates/codegen/kigi-sampler/src/actor/request_task.rs
+1 -1
View File
@@ -465,7 +465,7 @@ type ErrorCell = Arc<Mutex<Option<SamplingError>>>;
/// Wrap a raw chunk stream so its first error is captured into a
/// shared cell. The wrapped stream still yields the original
/// `Result<T, SamplingError>` items unchanged so the L2 transform sees
/// `Result<T, SamplingError>` items `unchanged` so the L2 transform sees
/// them and converts them to `SamplingErrorInfo` for events.
fn tee_errors<'a, T: Send + 'a>(
raw: BoxStream<'a, SamplingResult<T>>,
crates/codegen/kigi-sampler/src/actor/state.rs
+9 -18
View File
@@ -1,9 +1,7 @@
//! Actor-internal state.
//!
//! All fields are touched only from the actor task, so no mutex /
//! atomic synchronization is needed -- the actor's command-loop
//! serialization gives us a "single-threaded with shared state"
//! discipline matching the hunk-tracker pattern.
//! Every field is touched only from the actor task, so the command loop
//! serializes all access and no mutex or atomic is needed.
use std::collections::HashMap;
@@ -14,14 +12,12 @@ use crate::types::RequestId;
/// In-flight request bookkeeping.
///
/// `cancel_token` is owned by the actor (cloned into the spawned
/// per-request task). The completion oneshot is moved into the
/// per-request task at spawn time and is therefore not stored here.
/// The completion oneshot is moved into the per-request task at spawn time,
/// so only the cancel token remains reachable from the actor.
pub(crate) struct ActiveRequest {
pub(crate) cancel_token: CancellationToken,
}
/// Actor-owned state.
pub(crate) struct ActorState {
pub(crate) active_requests: HashMap<RequestId, ActiveRequest>,
pub(crate) config: SamplerConfig,
@@ -37,9 +33,8 @@ impl ActorState {
}
}
/// Register a newly-spawned request. Returns the previous entry if
/// the same `request_id` was already in flight (callers should
/// cancel the previous token before overwriting).
/// Returns the displaced entry when `request_id` was already in flight;
/// the caller must cancel that token, nothing else will.
pub(crate) fn register(
&mut self,
request_id: RequestId,
@@ -48,14 +43,12 @@ impl ActorState {
self.active_requests.insert(request_id, active)
}
/// Remove a request from the active set without cancelling its
/// token. Used by the cleanup signal sent from per-request tasks
/// when they exit normally.
/// Drops the entry without cancelling its token, for per-request tasks
/// signalling cleanup after they exit normally.
pub(crate) fn remove(&mut self, request_id: &RequestId) -> Option<ActiveRequest> {
self.active_requests.remove(request_id)
}
/// Cancel and remove an in-flight request.
pub(crate) fn cancel(&mut self, request_id: &RequestId) -> bool {
if let Some(active) = self.active_requests.remove(request_id) {
active.cancel_token.cancel();
@@ -65,8 +58,7 @@ impl ActorState {
}
}
/// Replace the default config. The next request submitted without
/// an override will use this.
/// Applies to every later request that carries no per-request override.
pub(crate) fn update_config(&mut self, config: SamplerConfig) {
self.config = config;
}
@@ -78,7 +70,6 @@ mod tests {
use crate::client::ApiBackend;
use indexmap::IndexMap;
/// Minimal config builder for tests in this module.
fn cfg() -> SamplerConfig {
SamplerConfig {
api_key: None,
crates/codegen/kigi-sampler/src/attribution.rs
+1 -1
View File
@@ -6,7 +6,7 @@
//! the live token from its auth source and the server still rejected
//! it" buckets.
//!
//! `kigi-sampler` is intentionally decoupled from `kigi-shell`
//! `kigi-sampler` is deliberately decoupled from `kigi-shell`
//! (no shell types, no logging crate, no auth-manager dependency). The
//! caller wires an implementation of [`Auth401AttributionCallback`]
//! into [`crate::SamplerConfig::attribution_callback`]; the sampler
crates/codegen/kigi-sampler/src/client.rs
+1 -13
View File
@@ -242,7 +242,7 @@ fn extract_model_metadata(headers: &reqwest::header::HeaderMap) -> Option<Respon
/// `stream_options` fields without modifying the original `ChatCompletionRequest`.
///
/// Uses `#[serde(flatten)]` to inline all fields from the inner request,
/// allowing single-pass serialization instead of the previous two-pass
/// allowing single-pass serialization instead of a two-pass
/// approach (serialize to `Value`, mutate, serialize to bytes).
#[derive(Serialize)]
struct StreamingChatRequest<'a> {
@@ -309,9 +309,7 @@ struct ClientDefaults {
openai_codex: bool,
}
// =============================================================================
// User-Agent helpers
// =============================================================================
#[derive(Clone, Debug, Eq, PartialEq)]
struct PlatformInfo {
@@ -378,9 +376,7 @@ pub fn user_agent_string_for(origin: &OriginClientInfo) -> String {
}
}
// =============================================================================
// SamplingClient
// =============================================================================
impl SamplingClient {
/// Construct a sampling client from a [`SamplerConfig`].
@@ -888,9 +884,7 @@ impl SamplingClient {
Ok(completion)
}
// =========================================================================
// Chat Completions API
// =========================================================================
pub async fn chat_completion(
&self,
@@ -1113,9 +1107,7 @@ impl SamplingClient {
Ok((chunks, model_metadata))
}
// =========================================================================
// Responses API
// =========================================================================
/// Apply default configuration to a Responses API request.
fn apply_response_defaults(&self, request: &mut CreateResponseWrapper) -> Result<()> {
@@ -1479,9 +1471,7 @@ impl SamplingClient {
Ok((events, model_metadata, doom_loop))
}
// =========================================================================
// Anthropic Messages API
// =========================================================================
/// Apply default configuration to a Messages API request.
fn apply_message_defaults(&self, request: &mut MessagesRequestWrapper) -> Result<()> {
@@ -1771,9 +1761,7 @@ impl SamplingClient {
Ok((events, model_metadata))
}
// =========================================================================
// Unified Conversation API
// =========================================================================
/// Apply default configuration to a ConversationRequest.
fn apply_conversation_defaults(&self, request: &mut ConversationRequest) -> Result<()> {
crates/codegen/kigi-sampler/src/commands.rs
+8 -12
View File
@@ -1,8 +1,8 @@
//! Internal actor protocol.
//!
//! `SamplerCommand` is `pub(crate)` because it is the wire between
//! [`SamplerHandle`](crate::handle::SamplerHandle) and the actor task,
//! not a public type. External callers always go through `SamplerHandle`.
//! [`SamplerHandle`](crate::handle::SamplerHandle) and the actor task.
//! External callers always go through `SamplerHandle`.
use tokio::sync::oneshot;
@@ -12,15 +12,12 @@ use crate::config::SamplerConfig;
use crate::metrics::InferenceLatencyStats;
use crate::types::RequestId;
/// Commands sent from a [`SamplerHandle`](crate::handle::SamplerHandle)
/// to the actor task.
///
/// Large payloads (`ConversationRequest`, `SamplerConfig`) are boxed so
/// every command stays cheap to copy through the mpsc channel.
/// Large payloads (`ConversationRequest`, `SamplerConfig`) are boxed so every
/// variant stays cheap to move through the mpsc channel.
pub(crate) enum SamplerCommand {
/// Submit a new sampling request. Fire-and-forget — results come via
/// events. When `completion_tx` is set the per-request task also
/// signals that channel for `submit_and_collect` callers.
/// Fire-and-forget — results arrive as events. When `completion_tx` is set
/// the per-request task also signals that channel, for
/// `submit_and_collect` callers.
Submit {
request_id: RequestId,
request: Box<ConversationRequest>,
@@ -33,10 +30,9 @@ pub(crate) enum SamplerCommand {
/// Cancel an in-flight request.
Cancel { request_id: RequestId },
/// Update the default sampling config (model switch, auth refresh).
/// Sent on a model switch or an auth refresh.
UpdateConfig { config: Box<SamplerConfig> },
/// Query: is a specific request still in flight?
IsActive {
request_id: RequestId,
reply: oneshot::Sender<bool>,
crates/codegen/kigi-sampler/src/config.rs
+34 -55
View File
@@ -1,10 +1,8 @@
//! Sampler configuration types.
//!
//! [`SamplerConfig`] is the per-request configuration handed to the
//! sampler. It deliberately does **not** alias
//! `kigi_sampling_types::SamplingConfig` so that the sampler crate
//! avoids transitive dependencies on shell-specific types
//! (`kigi-tools`, etc.).
//! [`SamplerConfig`] deliberately does **not** alias
//! `kigi_sampling_types::SamplingConfig`, which would drag shell-specific
//! types (`kigi-tools`, etc.) into this crate's dependency graph.
use indexmap::IndexMap;
use kigi_sampling_types::{
@@ -23,27 +21,18 @@ pub enum AuthScheme {
XApiKey,
}
/// All knobs that control a single sampling request.
/// All knobs that control a single sampling request. The session owns one per
/// active model and passes it — or a per-request override — to the actor on
/// every submit.
///
/// The session typically owns one `SamplerConfig` per active model
/// and passes it (or a per-request override) to the actor on every
/// submit.
///
/// # Construction in `kigi-shell`
///
/// `SamplerConfig` is the single source of truth for sampler
/// configuration. The shell builds it directly (see
/// `kigi-shell` builds it by composing chat-state's
/// `kigi_sampling_types::SamplingConfig` with `Credentials`; see
/// `agent::config::sampling_config_for_model` and
/// `session::acp_session::SessionActor::reconstruct_full_config`) by
/// composing chat-state's `kigi_sampling_types::SamplingConfig`
/// with `Credentials` (api key, client version).
///
/// URL-derived request headers are
/// folded into [`Self::extra_headers`] by
/// `agent::config::inject_url_derived_headers` before the
/// `SamplerConfig` is handed to the actor. Auth is selected separately
/// via `auth_scheme`, while `api_backend` controls only the request/response
/// protocol shape.
/// `session::acp_session::SessionActor::reconstruct_full_config`. URL-derived
/// request headers are folded into [`Self::extra_headers`] by
/// `agent::config::inject_url_derived_headers` before the config reaches the
/// actor. Auth is selected separately via `auth_scheme`, while `api_backend`
/// controls only the request/response protocol shape.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct SamplerConfig {
pub api_key: Option<String>,
@@ -91,36 +80,29 @@ pub struct SamplerConfig {
pub stream_tool_calls: bool,
pub idle_timeout_secs: Option<u64>,
// Reasoning effort
pub reasoning_effort: Option<ReasoningEffort>,
/// ChatCompletions body-adaptation dialect (per-platform; BYOK/custom
/// endpoints default to the historical Kimi behavior; lenient default on
/// deserialize so persisted configs from before the field parse).
/// ChatCompletions body-adaptation dialect. BYOK and custom endpoints fall
/// back to Kimi behavior; `serde(default)` keeps persisted configs that
/// lack the key parseable.
#[serde(default)]
pub chat_compat: kigi_sampling_types::ChatCompat,
/// Client identity for the User-Agent header (`kigi/{version}` plus an
/// optional origin product). The old xAI proxy's identity headers
/// (`x-kigi-client-identifier` / `-client-version` / `-deployment-id` /
/// `-user-id`) are gone — User-Agent and `extra_headers` are the only
/// identity signals on the wire.
/// optional origin product). User-Agent and `extra_headers` are the only
/// identity signals this crate puts on the wire.
pub origin_client: Option<OriginClientInfo>,
/// Optional hook invoked at every UNAUTHORIZED (401) response
/// site. The sampler passes the bearer that was actually sent on
/// the wire to the callback; the implementation is free to do
/// whatever it wants with it (typically: join it with a live
/// credential source and emit an attribution event for diagnosis
/// of stale-token vs. server-rejected-live-token 401s). `None`
/// (default) is a no-op -- the 401 arm returns the same
/// `SamplingError::Auth` it always did.
/// Hook invoked at every UNAUTHORIZED (401) response site, receiving the
/// bearer that was actually sent on the wire — typically joined against a
/// live credential source to tell a stale token apart from a live token the
/// server rejected. `None` is a no-op and the 401 arm still yields
/// `SamplingError::Auth`.
///
/// `Arc<dyn Trait>` is not serializable, so the field is skipped
/// in (de)serialization. Round-tripping a config through serde
/// drops the callback; callers that deserialize a `SamplerConfig`
/// from disk must re-attach the callback before passing it to
/// [`crate::SamplingClient::new`] or 401 attribution will be
/// silently disabled for the rebuilt client.
/// `Arc<dyn Trait>` is not serializable, so a config round-tripped through
/// serde comes back without the callback. Callers deserializing a
/// `SamplerConfig` from disk must re-attach it before
/// [`crate::SamplingClient::new`], or 401 attribution is silently disabled
/// for the rebuilt client.
#[serde(skip)]
pub attribution_callback: Option<SharedAttributionCallback>,
@@ -140,11 +122,10 @@ pub struct SamplerConfig {
pub compaction_at_tokens: Option<CompactionAtTokens>,
/// Server-side doom-loop check policy; `None` disables it. When set, the
/// client itself sends the opt-in `x-kigi-doom-loop-check` header on
/// streaming Responses API requests and absorbs the reported trigger
/// events (unlike the environment headers in [`Self::extra_headers`],
/// this header gates the client's own decode behavior, so it lives with
/// the decoder).
/// client sends the opt-in `x-kigi-doom-loop-check` header on streaming
/// Responses API requests and absorbs the reported trigger events. Unlike
/// the environment headers in [`Self::extra_headers`], this header gates
/// the client's own decode behavior, so it lives with the decoder.
#[serde(default)]
pub doom_loop_recovery: Option<DoomLoopRecoveryPolicy>,
@@ -154,8 +135,8 @@ pub struct SamplerConfig {
}
impl Default for SamplerConfig {
/// Empty defaults so callers can use `..Default::default()` and
/// new fields don't ripple through every literal site.
/// Empty defaults so callers can spell `..Default::default()` and a new
/// field does not ripple through every struct literal.
fn default() -> Self {
Self {
api_key: None,
@@ -206,7 +187,6 @@ pub type SharedHeaderInjector = std::sync::Arc<dyn HeaderInjector>;
/// Retry knobs for the sampler's internal transport-error retry loop.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct RetryPolicy {
/// Maximum number of retries before giving up.
pub max_retries: u32,
/// After this many rate-limit (429) retries, escalate to the caller.
/// Lower than `max_retries` because rate-limit waits can be long.
@@ -245,7 +225,6 @@ mod tests {
);
}
/// Configs serialized before the field existed must keep deserializing.
#[test]
fn config_without_doom_loop_recovery_deserializes_to_none() {
let mut stripped = serde_json::to_value(SamplerConfig::default()).unwrap();
crates/codegen/kigi-sampler/src/doom_loop.rs
-8
View File
@@ -35,7 +35,6 @@ struct CollectorState {
}
impl DoomLoopSignalCollector {
/// A fresh, armed collector judging confidence with `policy`.
pub(crate) fn new(policy: DoomLoopRecoveryPolicy) -> Self {
let collector = Self::default();
if let Ok(mut state) = collector.inner.lock() {
@@ -90,7 +89,6 @@ impl DoomLoopSignalCollector {
swallow || named
}
/// Drain the recorded signals; empty when nothing was reported.
pub(crate) fn take(&self) -> Vec<DoomLoopSignal> {
match self.inner.lock() {
Ok(mut state) => std::mem::take(&mut state.signals),
@@ -111,7 +109,6 @@ impl DoomLoopSignalCollector {
}
}
/// Debug-log the first malformed payload per attempt (never per event).
fn log_malformed_once(&self) {
let Ok(mut state) = self.inner.lock() else {
return;
@@ -139,8 +136,6 @@ mod tests {
assert_eq!(signals[0].kind, DoomLoopSignalKind::TailRepetition(4));
}
/// Servers that omit the SSE `event:` name are still handled by the
/// payload `type` check.
#[test]
fn absorb_swallows_check_event_without_sse_name() {
let collector = DoomLoopSignalCollector::default();
@@ -168,7 +163,6 @@ mod tests {
let delta = r#"{"type":"response.output_text.delta","delta":"hi"}"#;
assert!(!collector.absorb("response.output_text.delta", delta));
assert!(collector.take().is_empty());
// Terminal response field is recorded but the event is forwarded.
let terminal = r#"{"type":"response.completed","response":{"id":"r1","doom_loop_check":{"triggers":["low_logprob@response"]}}}"#;
assert!(!collector.absorb("response.completed", terminal));
assert_eq!(collector.take().len(), 1);
@@ -203,8 +197,6 @@ mod tests {
assert!(collector.take().is_empty());
}
/// `abort_triggers` fires only on confident signals, does not drain, and
/// goes quiet once disarmed (the spent-budget attempt must complete).
#[test]
fn abort_triggers_requires_confidence_and_honors_disarm() {
let confident = r#"{"type":"response.doom_loop_check","doom_loop_check":{"triggers":["tail_repetition:8@thinking"]}}"#;
crates/codegen/kigi-sampler/src/events.rs
+1 -1
View File
@@ -143,7 +143,7 @@ pub struct SamplingErrorInfo {
/// Coarse-grained classification of a sampling failure.
///
/// Intentionally narrow — context-window-exceeded does NOT have its own
/// Deliberately narrow — context-window-exceeded does NOT have its own
/// variant because the sampler cannot reliably detect it (it lacks
/// tracked token counts). Context-window errors arrive as
/// `Api { status: 400, .. }` with model metadata; the session inspects
crates/codegen/kigi-sampler/src/kimi_compat.rs
+3 -2
View File
@@ -97,7 +97,8 @@ fn normalize_mistral_tool_call_ids(body: &mut Value) {
let mut h = hash;
while out.len() < LEN {
out.push(digits[(h % 36) as usize] as char);
h = h / 36 + 1; // +1 keeps the stream from collapsing to zeros
// +1 keeps the stream from collapsing to zeros
h = h / 36 + 1;
}
out
}
@@ -588,7 +589,7 @@ mod tests {
);
// Canonical `max` (what the K3 menu token parses to since the
// ReasoningEffort::Max split) passes through unchanged.
// ReasoningEffort::Max split) passes through `unchanged`.
let mut body = json!({ "model": "k3", "reasoning_effort": "max" });
adapt_chat_completions_body(&mut body);
assert_eq!(
crates/codegen/kigi-sampler/src/lib.rs
-10
View File
@@ -1,20 +1,11 @@
//! 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
//! component built on the same actor pattern as `kigi-hunk-tracker`.
//!
//! ## Layered API
//!
//! - **Layer 1**: [`client::SamplingClient`] returns raw chunk streams.
//! - **Layer 2**: [`stream`] transforms raw streams into [`SamplingEvent`]s.
//! - **Layer 3**: [`SamplerHandle`] manages concurrent requests with retry,
//! cancellation, and event-based coordination via the actor.
//!
//! The type skeleton, the pure retry / metrics / client logic, the
//! Layer-2 stream transforms ([`stream_chat_completions`],
//! [`stream_responses`], [`stream_messages`], [`collect_response`]),
//! and the actor with its per-request task tie these layers together.
pub mod actor;
pub mod attribution;
@@ -32,7 +23,6 @@ mod shared_http;
pub mod stream;
pub mod types;
// Public re-exports — the API surface consumers see.
pub use actor::SamplerActor;
pub use attribution::{
Auth401AttributionCallback, SENT_BEARER_PREFIX_LEN, SamplingConsumer, SharedAttributionCallback,
crates/codegen/kigi-sampler/src/metrics.rs
+2 -1
View File
@@ -188,7 +188,8 @@ mod tests {
let chunks: Vec<Instant> = (0..11)
.scan(100u64, |acc, i| {
let t = *acc;
*acc += (i + 1) * 10; // intervals: 10, 20, 30, ...
// intervals: 10, 20, 30, ...
*acc += (i + 1) * 10;
Some(offset(start, t))
})
.collect();
crates/codegen/kigi-sampler/src/retry.rs
+2 -2
View File
@@ -26,8 +26,8 @@
//!
//! 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.
//! the Kimi/Moonshot API. Unlike the old xAI proxy, it sends no
//! `x-should-retry` hint header.
use std::time::Duration;
crates/codegen/kigi-sampler/src/shared_http.rs
+8 -12
View File
@@ -19,11 +19,11 @@ use std::time::Duration;
static SHARED_H2: OnceLock<reqwest::Client> = OnceLock::new();
static SHARED_HTTP1: OnceLock<reqwest::Client> = OnceLock::new();
/// Kill switch: `KIGI_SAMPLER_SHARED_CLIENT=0` (or `false`, any case)
/// restores the old behavior of building a fresh `reqwest::Client` per
/// `SamplingClient`. Resolved once per process: the environment cannot
/// change externally after spawn, and latching keeps the rollback state
/// consistent with the read-once pool knobs.
/// Kill switch: `KIGI_SAMPLER_SHARED_CLIENT=0` (or `false`, any case) makes
/// every `SamplingClient` build its own `reqwest::Client`. Resolved once per
/// process: the environment cannot change externally after spawn, and
/// latching keeps the rollback state consistent with the read-once pool
/// knobs.
fn sharing_disabled() -> bool {
static DISABLED: OnceLock<bool> = OnceLock::new();
*DISABLED.get_or_init(|| {
@@ -56,18 +56,16 @@ fn shared(
Ok(cell.get_or_init(|| built).clone())
}
/// Shared HTTP/2 sampling client (connection pooling + h2 keepalive).
pub(crate) fn client() -> Result<reqwest::Client, reqwest::Error> {
shared(&SHARED_H2, build_http_client, sharing_disabled())
}
/// Shared HTTP/1.1 fallback client. Pool-less by construction, so sharing it
/// is behaviorally identical to building a fresh one.
/// Pool-less by construction, so sharing this client is behaviorally
/// identical to building a fresh one.
pub(crate) fn client_http1() -> Result<reqwest::Client, reqwest::Error> {
shared(&SHARED_HTTP1, build_http_client_http1, sharing_disabled())
}
/// Build a `reqwest::Client` for sampling with HTTP/2 + connection pooling.
/// Env knobs are read once, when the shared client is first built.
fn build_http_client() -> Result<reqwest::Client, reqwest::Error> {
let pool_max_idle: usize = std::env::var("KIGI_POOL_MAX_IDLE")
@@ -88,15 +86,13 @@ fn build_http_client() -> Result<reqwest::Client, reqwest::Error> {
.pool_idle_timeout(Duration::from_secs(pool_idle_timeout_secs))
.connect_timeout(Duration::from_secs(connect_timeout_secs))
.tcp_nodelay(true)
// HTTP/2 keep-alive: ping every 15s, timeout after 5s.
.http2_keep_alive_interval(Duration::from_secs(15))
.http2_keep_alive_timeout(Duration::from_secs(5))
.http2_keep_alive_while_idle(true)
.build()
}
/// Build a `reqwest::Client` constrained to HTTP/1.1 with pooling disabled.
/// Used as a fallback after HTTP/2 transport failures.
/// Fallback used after HTTP/2 transport failures.
fn build_http_client_http1() -> Result<reqwest::Client, reqwest::Error> {
let connect_timeout_secs: u64 = std::env::var("KIGI_CONNECT_TIMEOUT_SECS")
.ok()
crates/codegen/kigi-sampler/src/stream/chat_completions.rs
+3 -2
View File
@@ -96,7 +96,8 @@ pub fn stream_chat_completions<'a>(
loop {
let next = match tokio::time::timeout(idle_timeout, stream.next()).await {
Ok(Some(next)) => next,
Ok(None) => break, // stream ended normally
// stream ended normally
Ok(None) => break,
Err(_elapsed) => {
let err = SamplingError::IdleTimeout {
elapsed_secs: idle_timeout.as_secs(),
@@ -253,7 +254,7 @@ pub fn stream_chat_completions<'a>(
}
}
// ── Build the final response ─────────────────────────────────
// Build the final response
let tool_calls: Vec<ToolCall> = tool_call_acc
.into_values()
.map(|(id, name, arguments)| {
crates/codegen/kigi-sampler/src/stream/collect.rs
-1
View File
@@ -38,7 +38,6 @@ pub async fn collect_response(
response, metrics, ..
} => return Ok((*response, metrics)),
SamplingEvent::Failed { error, .. } => return Err(error),
// Drop intermediate events; this is a buffered collector.
_ => {}
}
}
crates/codegen/kigi-sampler/src/stream/messages.rs
+3 -2
View File
@@ -241,7 +241,8 @@ pub fn stream_messages<'a>(
arguments_delta: None,
};
}
_ => {} // Image / ToolResult are not expected in assistant streams.
// Image / ToolResult are not expected in assistant streams.
_ => {}
},
MessageStreamEvent::ContentBlockDelta { index, delta } => {
@@ -460,7 +461,7 @@ pub fn stream_messages<'a>(
return;
}
// ── Build the final response ─────────────────────────────────
// Build the final response
let model_id = final_model.unwrap_or_default();
// Match the OAI Responses convention: prompt_tokens = full prompt, cached_prompt_tokens = cache hits only.
let total_prompt_tokens = final_input_tokens
crates/codegen/kigi-sampler/src/stream/messages_tests.rs
+1 -1
View File
@@ -537,7 +537,7 @@ fn meaningful_content_classifier_treats_ping_as_keepalive() {
));
}
// ── Token usage: Anthropic Messages API cache-bucket accounting ────────────
// Token usage: Anthropic Messages API cache-bucket accounting
fn message_start_with_cache(
input: u32,
crates/codegen/kigi-sampler/src/stream/mod.rs
+5 -6
View File
@@ -1,12 +1,11 @@
//! Layer-2 stream transforms: turn raw HTTP chunk streams into
//! [`SamplingEvent`](crate::events::SamplingEvent) streams.
//!
//! Each backend has its own transform because the raw chunk types
//! differ; backend dispatch happens in M4's
//! [`actor::request_task`](crate::actor::request_task), which knows
//! the API backend from `SamplerConfig.api_backend` and calls the
//! matching `SamplingClient::conversation_stream*` method before
//! handing the result to the corresponding transform here.
//! Each backend needs its own transform because the raw chunk types differ.
//! Dispatch lives in [`actor::request_task`](crate::actor::request_task),
//! which reads the backend from `SamplerConfig.api_backend`, calls the
//! matching `SamplingClient::conversation_stream*` method, and hands the
//! result to the transform here.
pub mod chat_completions;
pub mod collect;
crates/codegen/kigi-sampler/src/stream/responses.rs
+7 -6
View File
@@ -261,7 +261,7 @@ pub fn stream_responses<'a>(
}
// Continuation chunk for a streaming FunctionCall's args.
// Drop silently if no preceding OutputItemAdded mapped.
// Drop silently if no preceding `OutputItemAdded` mapped.
ResponseStreamEvent::ResponseFunctionCallArgumentsDelta(args_event) => {
let delta = args_event.delta;
if !delta.is_empty()
@@ -323,7 +323,7 @@ pub fn stream_responses<'a>(
return;
}
// ── Backend-hosted tool lifecycle events ────────────
// Backend-hosted tool lifecycle events
// These tools are executed server-side by the agentic
// sampler. We emit progress events so the shell/pager
// can show status to the user.
@@ -406,7 +406,7 @@ pub fn stream_responses<'a>(
}
}
// ── Build the final response ─────────────────────────────────
// Build the final response
let mut response = match final_response {
Some(r) => r,
None => {
@@ -437,7 +437,7 @@ pub fn stream_responses<'a>(
// (`deserialize_response_event`) has already rewritten
// `u.total_tokens` to `context_details.input + output` when
// the backend emits it; on older deployments the wire
// value passes through unchanged.
// value passes through `unchanged`.
let usage = response.usage.as_ref().map(|u| TokenUsage {
prompt_tokens: u.input_tokens,
completion_tokens: u.output_tokens,
@@ -501,7 +501,8 @@ pub fn stream_responses<'a>(
cost_usd_ticks,
message_chunks_emitted: message_chunk_count,
doom_loop_signals,
stop_message: None, // not reported on the Responses API
// not reported on the Responses API
stop_message: None,
};
yield SamplingEvent::Completed {
@@ -848,7 +849,7 @@ mod tests {
#[tokio::test]
async fn function_call_args_delta_without_added_event_is_dropped() {
// ArgumentsDelta with no preceding OutputItemAdded has no
// ArgumentsDelta with no preceding `OutputItemAdded` has no
// output_index → tool_index mapping; drop silently.
let events: Vec<Result<rs::ResponseStreamEvent, SamplingError>> = vec![
Ok(function_call_args_delta_event(7, "{\"oops\":1}")),
crates/codegen/kigi-sampler/src/types.rs
+2 -5
View File
@@ -6,19 +6,16 @@ use serde::{Deserialize, Serialize};
/// Unique identifier for a sampling request.
///
/// Wraps a `String` so callers can pass an externally-assigned ID
/// (e.g., a session-assigned UUID) or generate a fresh random one via
/// [`RequestId::random`].
/// The inner type is `String` rather than a `Uuid` so callers can carry an
/// externally-assigned ID, such as a session-assigned one.
#[derive(Clone, Debug, Hash, Eq, PartialEq, Serialize, Deserialize)]
pub struct RequestId(String);
impl RequestId {
/// Generate a fresh random request ID backed by a UUIDv4.
pub fn random() -> Self {
Self(uuid::Uuid::new_v4().to_string())
}
/// Borrow the underlying string slice.
pub fn as_str(&self) -> &str {
&self.0
}
crates/codegen/kigi-sampler/tests/support/mod.rs
+3 -2
View File
@@ -16,8 +16,9 @@ pub fn test_config(base_url: &str, api_key: &str) -> SamplerConfig {
}
}
/// Drive one POST through the client; the canned `{}` body is not a valid
/// completion, but only the wire-level request matters here.
/// Drive one POST through the client. The response is discarded: the counting
/// server never returns a valid completion, and only the wire-level request
/// matters here.
pub async fn send_one(client: &SamplingClient) {
let request = ConversationRequest {
items: vec![ConversationItem::User(UserItem {
crates/codegen/kigi-sampler/tests/test_actor.rs
-31
View File
@@ -29,9 +29,7 @@ use kigi_sampling_types::{
};
use kigi_test_support::{SseEvent, sse};
// ---------------------------------------------------------------------------
// Mock server harness
// ---------------------------------------------------------------------------
struct MockServer {
addr: SocketAddr,
@@ -64,9 +62,7 @@ impl MockServer {
}
}
// ---------------------------------------------------------------------------
// Config + request helpers
// ---------------------------------------------------------------------------
fn test_config(base_url: String, model: &str) -> SamplerConfig {
SamplerConfig {
@@ -114,9 +110,7 @@ fn user_request(text: &str) -> ConversationRequest {
}
}
// ---------------------------------------------------------------------------
// SSE generators
// ---------------------------------------------------------------------------
/// Render test-helper [`SseEvent`]s (optional `event:` name + `data:`) as
/// axum SSE events for this file's router-based harness.
@@ -148,9 +142,7 @@ fn text_chunk_event(content: &str, finish: bool) -> Event {
Event::default().data(chunk.to_string())
}
// ---------------------------------------------------------------------------
// Actor lifecycle
// ---------------------------------------------------------------------------
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
async fn spawn_then_active_count_zero_then_cancel_unknown_is_noop() {
@@ -163,9 +155,7 @@ async fn spawn_then_active_count_zero_then_cancel_unknown_is_noop() {
assert_eq!(handle.active_count().await, 0);
}
// ---------------------------------------------------------------------------
// Submit + event flow
// ---------------------------------------------------------------------------
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
async fn submit_emits_started_first_token_channel_completed() {
@@ -226,9 +216,7 @@ async fn submit_emits_started_first_token_channel_completed() {
}
}
// ---------------------------------------------------------------------------
// submit_and_collect
// ---------------------------------------------------------------------------
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
async fn submit_and_collect_returns_response() {
@@ -258,9 +246,7 @@ async fn submit_and_collect_returns_response() {
assert_eq!(a.content.as_ref(), "collected response");
}
// ---------------------------------------------------------------------------
// Cancellation
// ---------------------------------------------------------------------------
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
async fn cancel_in_flight_request_terminates_task() {
@@ -313,9 +299,7 @@ async fn cancel_in_flight_request_terminates_task() {
server.shutdown();
}
// ---------------------------------------------------------------------------
// Concurrent requests
// ---------------------------------------------------------------------------
#[tokio::test(flavor = "multi_thread", worker_threads = 4)]
async fn two_concurrent_requests_complete_with_correct_request_ids() {
@@ -371,9 +355,7 @@ async fn two_concurrent_requests_complete_with_correct_request_ids() {
server.shutdown();
}
// ---------------------------------------------------------------------------
// Retry on transient transport error
// ---------------------------------------------------------------------------
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
async fn retries_on_500_then_succeeds() {
@@ -386,7 +368,6 @@ async fn retries_on_500_then_succeeds() {
async move {
let n = counter.fetch_add(1, Ordering::SeqCst);
if n == 0 {
// First attempt: server error.
Err::<Sse<_>, (StatusCode, String)>((
StatusCode::INTERNAL_SERVER_ERROR,
json!({ "error": { "message": "transient" } }).to_string(),
@@ -434,9 +415,7 @@ async fn retries_on_500_then_succeeds() {
);
}
// ---------------------------------------------------------------------------
// Rate limit exhausts threshold
// ---------------------------------------------------------------------------
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
async fn rate_limit_exhausts_at_threshold_and_yields_failed() {
@@ -489,9 +468,7 @@ async fn rate_limit_exhausts_at_threshold_and_yields_failed() {
assert!((1..=3).contains(&hits), "expected 1-3 hits, got {hits}");
}
// ---------------------------------------------------------------------------
// Auth error -> EmitToSession (immediate)
// ---------------------------------------------------------------------------
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
async fn auth_401_emits_failed_immediately_no_retry() {
@@ -542,9 +519,7 @@ async fn auth_401_emits_failed_immediately_no_retry() {
assert_eq!(counter.load(Ordering::SeqCst), 1, "no retries on 401");
}
// ---------------------------------------------------------------------------
// Anthropic Messages API: refusal stop_reason + mid-stream parse failure
// ---------------------------------------------------------------------------
fn messages_config(base_url: String) -> SamplerConfig {
let mut cfg = test_config(base_url, "messages-compatible-model");
@@ -709,9 +684,7 @@ async fn messages_unparseable_event_is_fatal_without_retry() {
assert_eq!(counter.load(Ordering::SeqCst), 1, "exactly one attempt");
}
// ---------------------------------------------------------------------------
// UpdateConfig invalidates cache + applies to subsequent requests
// ---------------------------------------------------------------------------
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
async fn update_config_changes_subsequent_request_model() {
@@ -765,9 +738,7 @@ async fn update_config_changes_subsequent_request_model() {
);
}
// ---------------------------------------------------------------------------
// Responses doom-loop check signals
// ---------------------------------------------------------------------------
fn responses_config(base_url: String, doom_loop: Option<DoomLoopRecoveryPolicy>) -> SamplerConfig {
let mut cfg = test_config(base_url, "test-model");
@@ -881,9 +852,7 @@ async fn responses_confident_doom_loop_signal_resamples_once() {
);
}
// ---------------------------------------------------------------------------
// Helpers for draining the event channel
// ---------------------------------------------------------------------------
/// Drain the event channel until a terminal event (`Completed` or
/// `Failed`) is received, or until `deadline` elapses.
crates/codegen/kigi-sampler/tests/test_kimi_wire.rs
-11
View File
@@ -37,9 +37,7 @@ use kigi_sampling_types::{
ToolResultItem, ToolSpec, UserItem, synthesized_reasoning_item,
};
// ---------------------------------------------------------------------------
// Mock server harness (same shape as test_actor.rs)
// ---------------------------------------------------------------------------
struct MockServer {
addr: SocketAddr,
@@ -139,9 +137,7 @@ async fn drain_until_terminal(
}
}
// ---------------------------------------------------------------------------
// 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() {
@@ -254,9 +250,7 @@ async fn kimi_stream_reasoning_tool_calls_and_choice_usage() {
}
}
// ---------------------------------------------------------------------------
// Request surface: bearer auth, kigi UA, kimi_compat body adaptations
// ---------------------------------------------------------------------------
type Captured = Arc<Mutex<Option<(HeaderMap, Value)>>>;
@@ -415,9 +409,7 @@ async fn request_carries_bearer_kigi_ua_and_kimi_dialect_body() {
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() {
@@ -494,9 +486,7 @@ async fn rate_limit_honors_retry_after_then_succeeds() {
);
}
// ---------------------------------------------------------------------------
// Mid-stream network drop → retry → recovery
// ---------------------------------------------------------------------------
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
async fn mid_stream_drop_recovers_via_retry() {
@@ -509,7 +499,6 @@ async fn mid_stream_drop_recovers_via_retry() {
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(