Publish harness and TUI open-source

initial sync from the monorepo

crates/codegen/xai-grok-sampler/Cargo.toml

@@ -0,0 +1,46 @@
+[package]
+license = "Apache-2.0"
+name = "xai-grok-sampler"
+version = "0.1.0"
+edition.workspace = true
+description = "Actor-based sampling/inference layer for xAI grok (HTTP streaming + retry, no shell coupling)"
+
+[dependencies]
+# Internal
+xai-grok-sampling-types = { path = "../xai-grok-sampling-types" }
+xai-grok-version = { workspace = true }
+
+# External
+async-openai = { workspace = true }
+async-stream = { workspace = true }
+chrono = { workspace = true, features = ["serde"] }
+eventsource-stream = { workspace = true }
+futures-util = { workspace = true }
+indexmap = { workspace = true, features = ["serde"] }
+reqwest = { workspace = true, features = ["stream"] }
+serde = { workspace = true, features = ["derive"] }
+serde_json = { workspace = true }
+tokio = { workspace = true, features = ["rt", "macros", "time", "sync"] }
+tokio-util = { workspace = true }
+tracing = { workspace = true }
+uuid = { workspace = true, features = ["v4"] }
+
+[dev-dependencies]
+# Mock HTTP server for actor / request_task integration tests.
+axum = { workspace = true }
+# Shared SSE generators for the happy-path mock payloads.
+xai-grok-test-support = { workspace = true }
+# `start_paused = true` for deterministic timeout testing; `net` for
+# `TcpListener::bind` in the mock server harness.
+tokio = { workspace = true, features = [
+    "rt",
+    "macros",
+    "time",
+    "sync",
+    "test-util",
+    "net",
+    "rt-multi-thread",
+] }
+
+[lints]
+workspace = true

crates/codegen/xai-grok-sampler/src/actor/mod.rs

@@ -0,0 +1,145 @@
+//! Sampler actor: owns global state, spawns per-request tasks.
+//!
+//! The actor task itself is single-threaded -- it processes one
+//! command at a time -- but it spawns `tokio::spawn` per-request
+//! tasks for the actual streaming work, so multiple requests can be
+//! in flight concurrently.
+
+pub(crate) mod request_task;
+pub(crate) mod state;
+
+use tokio::sync::mpsc;
+use tokio::task::JoinSet;
+use tokio_util::sync::CancellationToken;
+
+use crate::commands::SamplerCommand;
+use crate::config::{RetryPolicy, SamplerConfig};
+use crate::events::SamplingEvent;
+use crate::handle::SamplerHandle;
+use state::{ActiveRequest, ActorState};
+
+use crate::types::RequestId;
+
+/// Sampler actor.
+///
+/// Construct via [`SamplerActor::spawn`]; the returned
+/// [`SamplerHandle`] is the only supported way to interact with it.
+pub struct SamplerActor {
+    cmd_rx: mpsc::UnboundedReceiver<SamplerCommand>,
+    event_tx: mpsc::UnboundedSender<SamplingEvent>,
+    state: ActorState,
+    /// Per-request tasks. The actor's run loop selects on
+    /// `cmd_rx.recv()` and `tasks.join_next()`; when a task finishes
+    /// it returns its `RequestId` so the actor can clean up
+    /// `active_requests`.
+    tasks: JoinSet<RequestId>,
+}
+
+impl SamplerActor {
+    /// Spawn the actor on the current tokio runtime and return a
+    /// handle. The actor stops when the returned handle (and all its
+    /// clones) are dropped.
+    pub fn spawn(
+        config: SamplerConfig,
+        retry_policy: RetryPolicy,
+        event_tx: mpsc::UnboundedSender<SamplingEvent>,
+    ) -> SamplerHandle {
+        let (cmd_tx, cmd_rx) = mpsc::unbounded_channel();
+        let actor = Self {
+            cmd_rx,
+            event_tx,
+            state: ActorState::new(config, retry_policy),
+            tasks: JoinSet::new(),
+        };
+        tokio::spawn(actor.run());
+        SamplerHandle::new(cmd_tx)
+    }
+
+    async fn run(mut self) {
+        loop {
+            tokio::select! {
+                biased;
+                // Prefer cleaning up finished tasks before processing
+                // new commands -- prevents `active_requests` from
+                // staying stale longer than necessary.
+                Some(joined) = self.tasks.join_next(), if !self.tasks.is_empty() => {
+                    match joined {
+                        Ok(request_id) => {
+                            // Task finished normally; remove from
+                            // active set unless the user has already
+                            // cancelled it (Cancel removes it too).
+                            self.state.remove(&request_id);
+                        }
+                        Err(join_err) => {
+                            tracing::warn!(
+                                error = %join_err,
+                                "request task panicked or was aborted"
+                            );
+                        }
+                    }
+                }
+                cmd = self.cmd_rx.recv() => {
+                    match cmd {
+                        Some(cmd) => self.handle_command(cmd),
+                        None => break, // all handles dropped
+                    }
+                }
+            }
+        }
+
+        // Cancel any still-running tasks before exiting so they don't
+        // leak. The cancellation token shutdown is best-effort.
+        for (_, active) in self.state.active_requests.drain() {
+            active.cancel_token.cancel();
+        }
+        self.tasks.shutdown().await;
+    }
+
+    fn handle_command(&mut self, cmd: SamplerCommand) {
+        match cmd {
+            SamplerCommand::Submit {
+                request_id,
+                request,
+                config,
+                completion_tx,
+            } => {
+                let cancel_token = CancellationToken::new();
+                let active = ActiveRequest {
+                    cancel_token: cancel_token.clone(),
+                };
+                if let Some(prev) = self.state.register(request_id.clone(), active) {
+                    // Caller submitted a duplicate id; cancel the
+                    // previous one so we don't leak its task.
+                    prev.cancel_token.cancel();
+                }
+                let effective_config = config
+                    .map(|b| *b)
+                    .unwrap_or_else(|| self.state.config.clone());
+                let event_tx = self.event_tx.clone();
+                let retry_policy = self.state.retry_policy.clone();
+                let request_inner = *request;
+                self.tasks.spawn(request_task::run_request_task(
+                    request_id,
+                    request_inner,
+                    effective_config,
+                    retry_policy,
+                    event_tx,
+                    cancel_token,
+                    completion_tx,
+                ));
+            }
+            SamplerCommand::Cancel { request_id } => {
+                self.state.cancel(&request_id);
+            }
+            SamplerCommand::UpdateConfig { config } => {
+                self.state.update_config(*config);
+            }
+            SamplerCommand::IsActive { request_id, reply } => {
+                let _ = reply.send(self.state.active_requests.contains_key(&request_id));
+            }
+            SamplerCommand::ActiveCount { reply } => {
+                let _ = reply.send(self.state.active_requests.len());
+            }
+        }
+    }
+}

crates/codegen/xai-grok-sampler/src/actor/request_task.rs

@@ -0,0 +1,861 @@
+//! Per-request streaming task.
+//!
+//! Spawned by the actor's `Submit` handler. Owns the retry loop and
+//! consumes a Layer 2 stream from the matching backend transform.
+//! Cancellation is cooperative via `CancellationToken`.
+
+use std::pin::pin;
+use std::sync::{Arc, Mutex};
+use std::time::Duration;
+
+use futures_util::StreamExt;
+use futures_util::stream::BoxStream;
+use tokio::sync::{mpsc, oneshot};
+use tokio_util::sync::CancellationToken;
+use tracing::Instrument;
+
+use xai_grok_sampling_types::{
+    ConversationRequest, ConversationResponse, EmptyResponseContext, SamplingError,
+    error::Result as SamplingResult,
+};
+
+use crate::client::{ApiBackend, SamplingClient};
+use crate::config::{RetryPolicy, SamplerConfig};
+use crate::events::{SamplingErrorInfo, SamplingErrorKind, SamplingEvent};
+use crate::metrics::InferenceLatencyStats;
+use crate::retry::{
+    self as retry_mod, RetryDecision, classify_error, clone_error, resolve_max_retries,
+};
+use crate::stream::{stream_chat_completions, stream_messages, stream_responses};
+use crate::types::RequestId;
+
+/// Default per-chunk idle timeout when neither config nor caller
+/// supplies one. Matches the shell's session-level default
+/// (5 minutes -- long enough for cold-start reasoning, short enough
+/// to detect dead streams before the user gives up).
+const DEFAULT_IDLE_TIMEOUT_SECS: u64 = 300;
+
+/// Result type for the `submit_and_collect` oneshot. Carries the rich
+/// `SamplingError` so callers can inspect retryability, status code,
+/// etc., without losing information through the
+/// `SamplingErrorInfo` round trip.
+pub(crate) type CompletionResult =
+    Result<(ConversationResponse, InferenceLatencyStats), SamplingError>;
+
+/// Outcome of a single attempt within the retry loop.
+enum AttemptOutcome {
+    /// Stream emitted [`SamplingEvent::Completed`] with a non-empty
+    /// response.
+    Completed {
+        response: Box<ConversationResponse>,
+        metrics: InferenceLatencyStats,
+    },
+    /// Stream emitted [`SamplingEvent::Completed`] but the response
+    /// was empty (no text, no tool calls). The retry loop treats this
+    /// as a transient failure (the model returned reasoning-only or
+    /// the stream was truncated). Metrics from the empty attempt are
+    /// discarded; a successful retry produces fresh ones.
+    Empty { context: EmptyResponseContext },
+    /// Stream emitted [`SamplingEvent::Failed`]. The captured raw
+    /// error is what the retry loop classifies; if no rich error was
+    /// captured (e.g. the failure was synthesised inside the L2
+    /// transform), `error` was reconstructed from the
+    /// [`SamplingErrorInfo`].
+    Failed { error: SamplingError },
+    /// `cancel_token` fired mid-attempt. The retry loop bails out
+    /// without further attempts.
+    Cancelled,
+    /// Failed to construct the underlying raw stream (e.g., HTTP
+    /// connect error before any chunks arrive).
+    InitFailed { error: SamplingError },
+}
+
+/// Run a single sampling request to completion (or final failure).
+///
+/// Returns the request id so the actor can clean it up from
+/// `active_requests` via [`tokio::task::JoinSet::join_next`].
+pub(crate) async fn run_request_task(
+    request_id: RequestId,
+    request: ConversationRequest,
+    config: SamplerConfig,
+    retry_policy: RetryPolicy,
+    event_tx: mpsc::UnboundedSender<SamplingEvent>,
+    cancel_token: CancellationToken,
+    completion_tx: Option<oneshot::Sender<CompletionResult>>,
+) -> RequestId {
+    let mut completion_tx = completion_tx;
+    let idle_timeout = Duration::from_secs(
+        config
+            .idle_timeout_secs
+            .unwrap_or(DEFAULT_IDLE_TIMEOUT_SECS),
+    );
+    let max_retries = resolve_max_retries(config.max_retries.or(Some(retry_policy.max_retries)));
+
+    // Build the initial client. Configuration errors here are fatal
+    // (no point retrying with the same broken config).
+    let mut client = match SamplingClient::new(config.clone()) {
+        Ok(c) => c,
+        Err(err) => {
+            emit_failed(&event_tx, &request_id, &err);
+            send_completion(&mut completion_tx, Err(err));
+            return request_id;
+        }
+    };
+
+    let sampling_span = crate::sampling_log::request_span(
+        &request_id,
+        &config.model,
+        &format!("{:?}", client.api_backend()),
+        &config.base_url,
+        &client.auth_info(),
+    );
+    if let Some(eff) = config.reasoning_effort {
+        sampling_span.record("reasoning_effort", eff.as_str());
+    }
+
+    let mut request = request;
+    let mut retry_count: u32 = 0;
+    // Doom-loop recovery keeps its own resample budget, independent of the
+    // transport/empty budget above.
+    let doom_policy = config.doom_loop_recovery;
+    let doom_max_retries = doom_policy.map_or(0, |p| p.max_retries);
+    let mut doom_retry_count: u32 = 0;
+
+    loop {
+        if cancel_token.is_cancelled() {
+            handle_cancellation(&event_tx, &request_id, &mut completion_tx);
+            return request_id;
+        }
+
+        // Once the resample budget is spent, the attempt runs with the abort
+        // disarmed so it can complete and be accepted as-is.
+        let doom_check = doom_policy.filter(|_| doom_retry_count < doom_max_retries);
+        let outcome = run_one_attempt(
+            &client,
+            request.clone(),
+            request_id.clone(),
+            idle_timeout,
+            &event_tx,
+            &cancel_token,
+            doom_check,
+        )
+        .instrument(sampling_span.clone())
+        .await;
+
+        match outcome {
+            AttemptOutcome::Completed {
+                response,
+                mut metrics,
+            } => {
+                metrics.attempts = retry_count + doom_retry_count + 1;
+                if let Some(policy) = doom_policy {
+                    let confident = policy.confident_triggers(&response.doom_loop_signals);
+                    if !confident.is_empty() {
+                        tracing::warn!(
+                            target: crate::sampling_log::TARGET,
+                            triggers = ?confident,
+                            attempt = doom_retry_count + 1,
+                            outcome = "accepted_after_budget",
+                            "doom-loop recovery: resample budget spent; accepting as-is"
+                        );
+                    }
+                }
+                // Surface token usage on the sampling span alongside effort.
+                if let Some(usage) = response.usage.as_ref() {
+                    sampling_span.record("output_tokens", usage.completion_tokens);
+                    sampling_span.record("reasoning_tokens", usage.reasoning_tokens);
+                }
+                // Emit Completed only after the loop succeeds; the L2
+                // stream's terminal event was suppressed by
+                // `run_one_attempt`.
+                let _ = event_tx.send(SamplingEvent::Completed {
+                    request_id: request_id.clone(),
+                    response: response.clone(),
+                    metrics: metrics.clone(),
+                });
+                send_completion(&mut completion_tx, Ok((*response, metrics)));
+                return request_id;
+            }
+            AttemptOutcome::Empty { context } => {
+                tracing::warn!(
+                    target: crate::sampling_log::TARGET,
+                    empty_response = true,
+                    empty_reason = context.reason.as_str(),
+                    had_reasoning = context.had_reasoning,
+                    content_len = context.content_len,
+                    tool_call_count = context.tool_call_count,
+                    completion_tokens = context.completion_tokens.unwrap_or(0),
+                    reasoning_tokens = context.reasoning_tokens.unwrap_or(0),
+                    finish_reason = context.finish_reason_str(),
+                    first_choice_seen = context.first_choice_seen,
+                    model = %context.model,
+                    "empty response from model: {reason} (retrying)",
+                    reason = context.reason,
+                );
+                let err = SamplingError::EmptyResponse { context };
+                if !apply_retry_decision(
+                    &err,
+                    &mut retry_count,
+                    max_retries,
+                    &retry_policy,
+                    &event_tx,
+                    &request_id,
+                    &mut request,
+                    &mut client,
+                    &config,
+                    &mut completion_tx,
+                )
+                .await
+                {
+                    return request_id;
+                }
+            }
+            AttemptOutcome::Failed { error } => {
+                // Doom-loop resamples run on their own budget and never
+                // consult the transport classifier, so no classifier change
+                // can silently debit the transport budget for a doom failure.
+                if let SamplingError::DoomLoopDetected { .. } = &error {
+                    let backoff = retry_mod::doom_loop_backoff(doom_retry_count + 1);
+                    doom_retry_count += 1;
+                    tracing::warn!(
+                        target: crate::sampling_log::TARGET,
+                        reason = %error,
+                        attempt = doom_retry_count,
+                        max_retries = doom_max_retries,
+                        outcome = "resampled",
+                        "doom-loop recovery: discarding the poisoned attempt and resampling"
+                    );
+                    emit_retrying(
+                        &event_tx,
+                        &request_id,
+                        doom_retry_count,
+                        doom_max_retries,
+                        &error,
+                    );
+                    tokio::time::sleep(backoff).await;
+                    continue;
+                }
+                if !apply_retry_decision(
+                    &error,
+                    &mut retry_count,
+                    max_retries,
+                    &retry_policy,
+                    &event_tx,
+                    &request_id,
+                    &mut request,
+                    &mut client,
+                    &config,
+                    &mut completion_tx,
+                )
+                .await
+                {
+                    return request_id;
+                }
+            }
+            AttemptOutcome::Cancelled => {
+                handle_cancellation(&event_tx, &request_id, &mut completion_tx);
+                return request_id;
+            }
+            AttemptOutcome::InitFailed { error } => {
+                if !apply_retry_decision(
+                    &error,
+                    &mut retry_count,
+                    max_retries,
+                    &retry_policy,
+                    &event_tx,
+                    &request_id,
+                    &mut request,
+                    &mut client,
+                    &config,
+                    &mut completion_tx,
+                )
+                .await
+                {
+                    return request_id;
+                }
+            }
+        }
+    }
+}
+
+/// Apply a [`RetryDecision`]. Returns `true` if the loop should
+/// continue, `false` if the request is finished (either fatal or
+/// emit-to-session). Performs the side-effects of the decision:
+/// sleeping, rebuilding the client, stripping images, emitting the
+/// `Retrying` event.
+#[allow(clippy::too_many_arguments)]
+async fn apply_retry_decision(
+    err: &SamplingError,
+    retry_count: &mut u32,
+    max_retries: u32,
+    retry_policy: &RetryPolicy,
+    event_tx: &mpsc::UnboundedSender<SamplingEvent>,
+    request_id: &RequestId,
+    request: &mut ConversationRequest,
+    client: &mut SamplingClient,
+    config: &SamplerConfig,
+    completion_tx: &mut Option<oneshot::Sender<CompletionResult>>,
+) -> bool {
+    let rate_limit_threshold = if retry_policy.rate_limit_retry_threshold == 0 {
+        retry_mod::RATE_LIMIT_RETRY_THRESHOLD
+    } else {
+        retry_policy.rate_limit_retry_threshold
+    };
+    let decision = classify_error(err, *retry_count, max_retries, rate_limit_threshold);
+
+    // Connection-reset / broken-pipe on body upload often means nginx
+    // rejected an oversized payload before responding 413. Strip
+    // images proactively before any retry of those errors so we don't
+    // burn budget re-uploading the same large body.
+    if err.is_likely_body_rejected() {
+        let stripped = request.strip_images();
+        if stripped > 0 {
+            tracing::warn!(
+                stripped,
+                "stripped {stripped} image(s) before retry (likely nginx 413 via connection reset)"
+            );
+        }
+    }
+
+    match decision {
+        RetryDecision::Retry { backoff } => {
+            *retry_count += 1;
+            emit_retrying(event_tx, request_id, *retry_count, max_retries, err);
+            tokio::time::sleep(backoff).await;
+            true
+        }
+        RetryDecision::RetryWithBackoff { backoff, .. } => {
+            *retry_count += 1;
+            emit_retrying(event_tx, request_id, *retry_count, max_retries, err);
+            tokio::time::sleep(backoff).await;
+            true
+        }
+        RetryDecision::RetryWithImageStrip => {
+            let stripped = request.strip_images();
+            if stripped == 0 {
+                // Nothing left to strip; upgrade to fatal.
+                emit_failed(event_tx, request_id, err);
+                send_completion(completion_tx, Err(clone_error(err)));
+                return false;
+            }
+            *retry_count += 1;
+            emit_retrying(event_tx, request_id, *retry_count, max_retries, err);
+            true
+        }
+        RetryDecision::RetryWithClientRebuild { backoff } => {
+            *retry_count += 1;
+            emit_retrying(event_tx, request_id, *retry_count, max_retries, err);
+            tokio::time::sleep(backoff).await;
+
+            // Rebuild client with HTTP/1.1 fallback to escape poisoned
+            // HTTP/2 connection pools.
+            let mut http1_config = config.clone();
+            http1_config.force_http1 = true;
+            match SamplingClient::new(http1_config) {
+                Ok(fresh) => {
+                    *client = fresh;
+                    tracing::info!("rebuilt sampling client with HTTP/1.1 fallback for retry");
+                }
+                Err(rebuild_err) => {
+                    tracing::warn!(
+                        error = %rebuild_err,
+                        "failed to rebuild HTTP/1.1 client for retry; reusing existing client"
+                    );
+                }
+            }
+            true
+        }
+        RetryDecision::EmitToSession(emitted_err) => {
+            emit_failed(event_tx, request_id, &emitted_err);
+            send_completion(completion_tx, Err(emitted_err));
+            false
+        }
+        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".
+            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
+                };
+            if budget_exhausted {
+                let exhausted_span = tracing::info_span!(
+                    "http.retries_exhausted",
+                    total_attempts = next_attempt as i64,
+                    model = %config.model,
+                    error = %err,
+                    status_code = tracing::field::Empty,
+                );
+                let status_code = match err {
+                    SamplingError::Api { status, .. } => Some(status.as_u16()),
+                    SamplingError::Http(e) => e.status().map(|s| s.as_u16()),
+                    _ => None,
+                };
+                if let Some(status) = status_code {
+                    exhausted_span.record("status_code", status as i64);
+                }
+                exhausted_span.in_scope(|| {});
+            }
+            emit_failed(event_tx, request_id, &fatal_err);
+            send_completion(completion_tx, Err(fatal_err));
+            false
+        }
+    }
+}
+
+/// Run a single attempt: build the raw stream, drive it through the
+/// matching L2 transform, and forward all non-terminal events to
+/// `event_tx`. Captures the rich `SamplingError` from the underlying
+/// raw stream so the retry loop can classify it accurately.
+///
+/// `doom_check` is the doom-loop policy while the resample budget lasts;
+/// `None` disarms the mid-stream abort and the terminal confidence check so
+/// the attempt completes and its response can be accepted.
+async fn run_one_attempt(
+    client: &SamplingClient,
+    request: ConversationRequest,
+    request_id: RequestId,
+    idle_timeout: Duration,
+    event_tx: &mpsc::UnboundedSender<SamplingEvent>,
+    cancel_token: &CancellationToken,
+    doom_check: Option<xai_grok_sampling_types::DoomLoopRecoveryPolicy>,
+) -> AttemptOutcome {
+    match client.api_backend() {
+        ApiBackend::ChatCompletions => {
+            let (raw, metadata) = match client.conversation_stream(request).await {
+                Ok(pair) => pair,
+                Err(e) => return AttemptOutcome::InitFailed { error: e },
+            };
+            let (teed, captured) = tee_errors(raw);
+            let l2 = stream_chat_completions(teed, metadata, request_id.clone(), idle_timeout);
+            drive_l2(l2, request_id, event_tx, cancel_token, captured, None).await
+        }
+        ApiBackend::Responses => {
+            let (raw, metadata, doom_loop) =
+                match client.conversation_stream_responses(request).await {
+                    Ok(parts) => parts,
+                    Err(e) => return AttemptOutcome::InitFailed { error: e },
+                };
+            if doom_check.is_none()
+                && let Some(collector) = &doom_loop
+            {
+                collector.disarm_abort();
+            }
+            let (teed, captured) = tee_errors(raw);
+            let l2 = stream_responses(teed, metadata, request_id.clone(), idle_timeout, doom_loop);
+            drive_l2(l2, request_id, event_tx, cancel_token, captured, doom_check).await
+        }
+        ApiBackend::Messages => {
+            let (raw, metadata) = match client.conversation_stream_messages(request).await {
+                Ok(pair) => pair,
+                Err(e) => return AttemptOutcome::InitFailed { error: e },
+            };
+            let (teed, captured) = tee_errors(raw);
+            let l2 = stream_messages(teed, metadata, request_id.clone(), idle_timeout);
+            drive_l2(l2, request_id, event_tx, cancel_token, captured, None).await
+        }
+    }
+}
+
+/// Captured-error cell shared between the tee adapter and the
+/// per-request task.
+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
+/// them and converts them to `SamplingErrorInfo` for events.
+fn tee_errors<'a, T: Send + 'a>(
+    raw: BoxStream<'a, SamplingResult<T>>,
+) -> (BoxStream<'a, SamplingResult<T>>, ErrorCell) {
+    let cell: ErrorCell = Arc::new(Mutex::new(None));
+    let cell_clone = Arc::clone(&cell);
+    let teed = raw
+        .map(move |item| {
+            if let Err(ref e) = item
+                && let Ok(mut guard) = cell_clone.lock()
+                && guard.is_none()
+            {
+                // Capture only the first error -- subsequent errors
+                // on a torn-down stream are usually secondary effects
+                // of the same disconnect.
+                *guard = Some(clone_error(e));
+            }
+            item
+        })
+        .boxed();
+    (teed, cell)
+}
+
+/// Drive an L2 event stream: forward non-terminal events to
+/// `event_tx`, watch `cancel_token`, return `AttemptOutcome` based on
+/// the terminal event (or cancellation). `doom_check`, when set, turns a
+/// completed response carrying confident doom-loop signals into a
+/// retryable failure (belt-and-braces behind the mid-stream abort).
+async fn drive_l2(
+    l2: impl futures_util::Stream<Item = SamplingEvent>,
+    request_id: RequestId,
+    event_tx: &mpsc::UnboundedSender<SamplingEvent>,
+    cancel_token: &CancellationToken,
+    captured: ErrorCell,
+    doom_check: Option<xai_grok_sampling_types::DoomLoopRecoveryPolicy>,
+) -> AttemptOutcome {
+    let mut l2 = pin!(l2);
+    loop {
+        tokio::select! {
+            biased;
+            _ = cancel_token.cancelled() => {
+                return AttemptOutcome::Cancelled;
+            }
+            next = l2.next() => match next {
+                Some(SamplingEvent::Completed { response, metrics, .. }) => {
+                    // Doom outranks the truncation/empty classes: a confident
+                    // loop poisons the attempt whatever else it looks like.
+                    if let Some(policy) = doom_check {
+                        let triggers = policy.confident_triggers(&response.doom_loop_signals);
+                        if !triggers.is_empty() {
+                            return AttemptOutcome::Failed {
+                                error: SamplingError::DoomLoopDetected {
+                                    triggers,
+                                    aborted_at_chunk: None,
+                                },
+                            };
+                        }
+                    }
+                    if response.stop_reason == Some(xai_grok_sampling_types::StopReason::Length) {
+                        return AttemptOutcome::Failed {
+                            error: SamplingError::MaxTokensTruncation,
+                        };
+                    }
+                    // A content-filtered turn (Anthropic refusal, OpenAI
+                    // content_filter stop reason) is legitimately content-less and
+                    // deterministic — resampling it would retry-storm.
+                    let content_filtered = response.stop_reason
+                        == Some(xai_grok_sampling_types::StopReason::ContentFilter);
+                    if !content_filtered && let Some(reason) = response.empty_reason() {
+                        let context = build_empty_context(reason, &response);
+                        return AttemptOutcome::Empty { context };
+                    }
+                    return AttemptOutcome::Completed { response, metrics };
+                }
+                Some(SamplingEvent::Failed { error: info, .. }) => {
+                    let raw = captured
+                        .lock()
+                        .ok()
+                        .and_then(|mut g| g.take());
+                    let error = raw.unwrap_or_else(|| synthesize_from_info(&info));
+                    return AttemptOutcome::Failed { error };
+                }
+                Some(other) => {
+                    let _ = event_tx.send(retag(other, &request_id));
+                }
+                None => {
+                    // L2 streams always terminate with Completed or
+                    // Failed; reaching None means the producer was
+                    // dropped without termination -- treat as a
+                    // synthetic transport error.
+                    return AttemptOutcome::Failed {
+                        error: SamplingError::EventStreamError(
+                            "stream dropped without terminal event".to_string(),
+                        ),
+                    };
+                }
+            }
+        }
+    }
+}
+
+/// Re-tag a forwarded event with the canonical request_id. The L2
+/// transform tags events with the id we passed in, so this is
+/// usually a no-op; keeping the helper makes the data-flow explicit.
+fn retag(event: SamplingEvent, _request_id: &RequestId) -> SamplingEvent {
+    event
+}
+
+/// Reconstruct a [`SamplingError`] from a [`SamplingErrorInfo`] when
+/// the L2 transform fired a synthesised Failed event (idle timeout,
+/// `ResponseFailed`, server error event) and there is no captured raw
+/// error in the cell.
+fn synthesize_from_info(info: &SamplingErrorInfo) -> SamplingError {
+    match info.kind {
+        SamplingErrorKind::IdleTimeout => SamplingError::IdleTimeout {
+            elapsed_secs: info
+                .message
+                .split_whitespace()
+                .find_map(|tok| tok.strip_suffix('s').and_then(|n| n.parse::<u64>().ok()))
+                .unwrap_or(0),
+        },
+        SamplingErrorKind::Auth => SamplingError::Auth(info.message.clone()),
+        // Must stay Serialization: EventStreamError is retryable, and a
+        // response-parse failure is deterministic on retry. `info.message`
+        // is the variant's rendered Display, so rebuild via the constructor
+        // that owns the prefix-stripping.
+        SamplingErrorKind::Serialization => {
+            SamplingError::serialization_from_rendered(&info.message)
+        }
+        SamplingErrorKind::Http => SamplingError::EventStreamError(info.message.clone()),
+        SamplingErrorKind::Api | SamplingErrorKind::RateLimited => {
+            let status = info
+                .status_code
+                .and_then(|c| reqwest::StatusCode::from_u16(c).ok())
+                .unwrap_or(reqwest::StatusCode::INTERNAL_SERVER_ERROR);
+            SamplingError::Api {
+                status,
+                message: info.message.clone(),
+                model_metadata: info.model_metadata.clone(),
+                retry_after_secs: info.retry_after_secs,
+                should_retry: None,
+            }
+        }
+        SamplingErrorKind::EmptyResponse => {
+            if let Some(ctx) = &info.empty_response_context {
+                SamplingError::EmptyResponse {
+                    context: ctx.clone(),
+                }
+            } else {
+                SamplingError::EventStreamError(info.message.clone())
+            }
+        }
+        SamplingErrorKind::MaxTokensTruncation => SamplingError::MaxTokensTruncation,
+        SamplingErrorKind::DoomLoopDetected => SamplingError::DoomLoopDetected {
+            triggers: info.doom_loop_triggers.clone().unwrap_or_default(),
+            aborted_at_chunk: info.doom_loop_aborted_at_chunk,
+        },
+    }
+}
+
+/// Build an [`EmptyResponseContext`] from a completed-but-empty response.
+fn build_empty_context(
+    reason: xai_grok_sampling_types::EmptyReason,
+    response: &ConversationResponse,
+) -> EmptyResponseContext {
+    let had_reasoning = response
+        .reasoning_items()
+        .any(|r| !r.summary.is_empty() || r.content.is_some() || r.encrypted_content.is_some());
+    let (content_len, tool_call_count, model, first_choice_seen) = match response.assistant() {
+        Some(a) => (
+            a.content.len(),
+            a.tool_calls.len(),
+            a.model_id.clone().unwrap_or_default(),
+            // If model_id is set, the L2 saw at least one choice.
+            a.model_id.is_some(),
+        ),
+        None => (0, 0, String::new(), false),
+    };
+
+    let finish_reason = response.stop_reason.map(|sr| sr.as_str().to_owned());
+    let (completion_tokens, reasoning_tokens, prompt_tokens) = response
+        .usage
+        .as_ref()
+        .map(|u| {
+            (
+                Some(u.completion_tokens),
+                Some(u.reasoning_tokens),
+                Some(u.prompt_tokens),
+            )
+        })
+        .unwrap_or((None, None, None));
+
+    EmptyResponseContext {
+        reason,
+        had_reasoning,
+        content_len,
+        tool_call_count,
+        finish_reason,
+        completion_tokens,
+        reasoning_tokens,
+        prompt_tokens,
+        model,
+        first_choice_seen,
+    }
+}
+
+fn emit_failed(
+    event_tx: &mpsc::UnboundedSender<SamplingEvent>,
+    request_id: &RequestId,
+    err: &SamplingError,
+) {
+    let info = SamplingErrorInfo::from(err);
+    let _ = event_tx.send(SamplingEvent::Failed {
+        request_id: request_id.clone(),
+        error: info,
+    });
+}
+
+fn emit_retrying(
+    event_tx: &mpsc::UnboundedSender<SamplingEvent>,
+    request_id: &RequestId,
+    attempt: u32,
+    max_retries: u32,
+    err: &SamplingError,
+) {
+    let info = SamplingErrorInfo::from(err);
+    let _ = event_tx.send(SamplingEvent::Retrying {
+        request_id: request_id.clone(),
+        attempt,
+        max_retries,
+        kind: info.kind,
+        reason: err.to_string(),
+        doom_loop_triggers: info.doom_loop_triggers,
+        doom_loop_aborted_at_chunk: info.doom_loop_aborted_at_chunk,
+    });
+}
+
+fn handle_cancellation(
+    event_tx: &mpsc::UnboundedSender<SamplingEvent>,
+    request_id: &RequestId,
+    completion_tx: &mut Option<oneshot::Sender<CompletionResult>>,
+) {
+    // No status code, no upstream API error -- this is a client-side
+    // termination. Use kind=Api so consumers that switch on kind have
+    // a sensible default; the message clearly identifies it.
+    let info = SamplingErrorInfo {
+        kind: SamplingErrorKind::Api,
+        status_code: None,
+        message: "request cancelled".to_string(),
+        is_retryable: false,
+        retry_after_secs: None,
+        model_metadata: None,
+        empty_response_context: None,
+        doom_loop_triggers: None,
+        doom_loop_aborted_at_chunk: None,
+    };
+    let _ = event_tx.send(SamplingEvent::Failed {
+        request_id: request_id.clone(),
+        error: info,
+    });
+    send_completion(
+        completion_tx,
+        Err(SamplingError::Auth("request cancelled".to_string())),
+    );
+}
+
+fn send_completion(
+    completion_tx: &mut Option<oneshot::Sender<CompletionResult>>,
+    result: CompletionResult,
+) {
+    if let Some(tx) = completion_tx.take() {
+        let _ = tx.send(result);
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+    use futures_util::stream;
+
+    #[test]
+    fn synthesize_idle_timeout_extracts_elapsed_secs() {
+        let info = SamplingErrorInfo {
+            kind: SamplingErrorKind::IdleTimeout,
+            status_code: None,
+            message: "inference idle timeout after 240s with no chunks".to_string(),
+            is_retryable: false,
+            retry_after_secs: None,
+            model_metadata: None,
+            empty_response_context: None,
+            doom_loop_triggers: None,
+            doom_loop_aborted_at_chunk: None,
+        };
+        let err = synthesize_from_info(&info);
+        match err {
+            SamplingError::IdleTimeout { elapsed_secs } => assert_eq!(elapsed_secs, 240),
+            other => panic!("expected IdleTimeout, got {other:?}"),
+        }
+    }
+
+    #[test]
+    fn synthesize_api_500_round_trips() {
+        let info = SamplingErrorInfo {
+            kind: SamplingErrorKind::Api,
+            status_code: Some(500),
+            message: "boom".to_string(),
+            is_retryable: true,
+            retry_after_secs: None,
+            model_metadata: None,
+            empty_response_context: None,
+            doom_loop_triggers: None,
+            doom_loop_aborted_at_chunk: None,
+        };
+        let err = synthesize_from_info(&info);
+        match err {
+            SamplingError::Api {
+                status, message, ..
+            } => {
+                assert_eq!(status.as_u16(), 500);
+                assert_eq!(message, "boom");
+            }
+            other => panic!("expected Api, got {other:?}"),
+        }
+    }
+
+    #[test]
+    fn synthesize_rate_limited_preserves_retry_after() {
+        let info = SamplingErrorInfo {
+            kind: SamplingErrorKind::RateLimited,
+            status_code: Some(429),
+            message: "slow down".to_string(),
+            is_retryable: true,
+            retry_after_secs: Some(7),
+            model_metadata: None,
+            empty_response_context: None,
+            doom_loop_triggers: None,
+            doom_loop_aborted_at_chunk: None,
+        };
+        let err = synthesize_from_info(&info);
+        match err {
+            SamplingError::Api {
+                status,
+                retry_after_secs,
+                ..
+            } => {
+                assert_eq!(status.as_u16(), 429);
+                assert_eq!(retry_after_secs, Some(7));
+            }
+            other => panic!("expected Api(429), got {other:?}"),
+        }
+    }
+
+    #[test]
+    fn synthesize_serialization_stays_serialization() {
+        // Round-trip a REAL error's Display so a Display-template rewording
+        // cannot silently reintroduce double-prefixing.
+        let original = SamplingError::Serialization(
+            serde_json::from_str::<i32>("missing field `delta`").unwrap_err(),
+        );
+        let info = SamplingErrorInfo::from(&original);
+        let err = synthesize_from_info(&info);
+        assert!(
+            matches!(err, SamplingError::Serialization(_)),
+            "expected Serialization, got {err:?}"
+        );
+        assert!(!err.is_retryable());
+        assert_eq!(
+            err.to_string(),
+            info.message,
+            "rebuilt Display must round-trip without double-prefixing"
+        );
+    }
+
+    #[tokio::test]
+    async fn tee_captures_first_error_only() {
+        let items: Vec<SamplingResult<u32>> = vec![
+            Ok(1),
+            Err(SamplingError::EventStreamError("first".into())),
+            Err(SamplingError::EventStreamError("second".into())),
+        ];
+        let raw = stream::iter(items).boxed();
+        let (mut teed, cell) = tee_errors(raw);
+        while teed.next().await.is_some() {}
+        let captured = cell.lock().unwrap().take().expect("error captured");
+        match captured {
+            SamplingError::EventStreamError(msg) => assert_eq!(msg, "first"),
+            other => panic!("expected EventStreamError, got {other:?}"),
+        }
+    }
+}

crates/codegen/xai-grok-sampler/src/actor/state.rs

@@ -0,0 +1,148 @@
+//! 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.
+
+use std::collections::HashMap;
+
+use tokio_util::sync::CancellationToken;
+
+use crate::config::{RetryPolicy, SamplerConfig};
+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.
+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,
+    pub(crate) retry_policy: RetryPolicy,
+}
+
+impl ActorState {
+    pub(crate) fn new(config: SamplerConfig, retry_policy: RetryPolicy) -> Self {
+        Self {
+            active_requests: HashMap::new(),
+            config,
+            retry_policy,
+        }
+    }
+
+    /// 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).
+    pub(crate) fn register(
+        &mut self,
+        request_id: RequestId,
+        active: ActiveRequest,
+    ) -> Option<ActiveRequest> {
+        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.
+    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();
+            true
+        } else {
+            false
+        }
+    }
+
+    /// Replace the default config. The next request submitted without
+    /// an override will use this.
+    pub(crate) fn update_config(&mut self, config: SamplerConfig) {
+        self.config = config;
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+    use crate::client::ApiBackend;
+    use indexmap::IndexMap;
+
+    /// Minimal config builder for tests in this module.
+    fn cfg() -> SamplerConfig {
+        SamplerConfig {
+            api_key: None,
+            base_url: "https://example.test".into(),
+            model: "test-model".into(),
+            max_completion_tokens: None,
+            temperature: None,
+            top_p: None,
+            api_backend: ApiBackend::ChatCompletions,
+            auth_scheme: Default::default(),
+            extra_headers: IndexMap::new(),
+            context_window: 8192,
+            force_http1: false,
+            max_retries: None,
+            stream_tool_calls: false,
+            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,
+            compactions_remaining: None,
+            compaction_at_tokens: None,
+            doom_loop_recovery: None,
+            header_injector: None,
+        }
+    }
+
+    #[test]
+    fn cancel_unknown_request_returns_false() {
+        let mut state = ActorState::new(cfg(), RetryPolicy::default());
+        assert!(!state.cancel(&RequestId::from("unknown")));
+    }
+
+    #[test]
+    fn register_then_cancel_removes() {
+        let mut state = ActorState::new(cfg(), RetryPolicy::default());
+        let id = RequestId::from("req-1");
+        state.register(
+            id.clone(),
+            ActiveRequest {
+                cancel_token: CancellationToken::new(),
+            },
+        );
+        assert_eq!(state.active_requests.len(), 1);
+        assert!(state.cancel(&id));
+        assert_eq!(state.active_requests.len(), 0);
+    }
+
+    #[test]
+    fn register_returns_previous_when_same_id() {
+        let mut state = ActorState::new(cfg(), RetryPolicy::default());
+        let id = RequestId::from("req-1");
+        let first = ActiveRequest {
+            cancel_token: CancellationToken::new(),
+        };
+        let second = ActiveRequest {
+            cancel_token: CancellationToken::new(),
+        };
+        assert!(state.register(id.clone(), first).is_none());
+        assert!(state.register(id.clone(), second).is_some());
+    }
+}

crates/codegen/xai-grok-sampler/src/attribution.rs

@@ -0,0 +1,120 @@
+//! 401 attribution callback hook for the sampling client.
+//!
+//! Every 401 response site can optionally emit an attribution event so
+//! a downstream observer can split production 401s into "client sent a
+//! stale snapshot bearer that the server rejected" vs. "client sent
+//! the live token from its auth source and the server still rejected
+//! it" buckets.
+//!
+//! `xai-grok-sampler` is intentionally decoupled from `xai-grok-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
+//! invokes the callback at each UNAUTHORIZED arm with the bearer that
+//! was actually sent on the wire. The implementation is free to join
+//! the bearer with whatever live credential source it owns and emit
+//! the attribution however it wants.
+//!
+//! When the callback is `None` (the default), the 401 sites are silent
+//! and return the same `SamplingError::Auth` they would otherwise.
+
+use std::sync::Arc;
+
+/// A logical 401-emitting site inside the sampling client. The string
+/// identifier ends up in the consumer field of the attribution event
+/// so downstream queries can break down 401s by API path.
+///
+/// # Scope: sampler endpoints only
+///
+/// This enum enumerates the six HTTP endpoints owned by
+/// `SamplingClient` (chat completions, responses, messages -- each in
+/// streaming and non-streaming form). It does *not* cover image
+/// generation, video generation, web search, or embedding -- those
+/// tools live in `xai-grok-tools`
+/// (`crates/codegen/xai-grok-tools/src/implementations/`), have their
+/// own HTTP clients that do not flow through `SamplingClient`, and
+/// hook into the `xai_grok_tools::ApiKeyProvider` trait rather than
+/// this enum.
+#[derive(Debug, Clone, Copy, PartialEq, Eq)]
+pub enum SamplingConsumer {
+    /// `chat_completion_stream`: OpenAI-compatible streaming OpenAI Chat Completions API.
+    ChatCompletionsStream,
+    /// `chat_completion`: OpenAI-compatible non-streaming OpenAI Chat Completions API.
+    ChatCompletions,
+    /// `create_response_stream`: Responses API streaming.
+    ResponsesStream,
+    /// `create_response`: Responses API non-streaming.
+    Responses,
+    /// `messages_stream`: Anthropic Messages API streaming.
+    MessagesStream,
+    /// `messages`: Anthropic Messages API non-streaming.
+    Messages,
+}
+
+impl SamplingConsumer {
+    /// Stable string identifier for this emit site. Callbacks
+    /// typically combine this with a fixed prefix (e.g. the client
+    /// type) when building the consumer field of the attribution
+    /// event.
+    pub fn as_endpoint(self) -> &'static str {
+        match self {
+            Self::ChatCompletionsStream => "chat_completions_stream",
+            Self::ChatCompletions => "chat_completions",
+            Self::ResponsesStream => "responses_stream",
+            Self::Responses => "responses",
+            Self::MessagesStream => "messages_stream",
+            Self::Messages => "messages",
+        }
+    }
+}
+
+/// Maximum prefix length the sampler shares with attribution
+/// callbacks across the crate boundary. Mirrors
+/// `xai_grok_shell::auth::token_suffix` (which truncates to 12 chars
+/// before any sink) so the two crates stay in lock-step on the
+/// "bearers leaving the sampler are 12-char prefixes only" invariant.
+///
+/// The cross-crate boundary is the only place this constant is
+/// load-bearing -- changing it requires updating `token_suffix` in
+/// `xai-grok-shell/src/auth/manager.rs` to match, otherwise the
+/// shell's local-log payload and the sampler's callback argument
+/// will disagree on prefix length.
+pub const SENT_BEARER_PREFIX_LEN: usize = 12;
+/// Hook invoked by [`crate::SamplingClient`] at every 401 response site.
+///
+/// Implementations are responsible for joining `sent_bearer_prefix`
+/// with whatever live credential source they own (e.g. an auth
+/// manager holding the most-recently-refreshed token) and emitting
+/// whatever attribution event makes sense for their observability
+/// stack.
+///
+/// Implementations must be cheap to invoke and must not block. They
+/// run inside the request's response-handling path and any latency
+/// they add is paid by the user-visible 401 error path.
+//
+// The `Debug` bound is a structural requirement: [`crate::SamplerConfig`]
+// derives `Debug` and carries an `Option<Arc<dyn Auth401AttributionCallback>>`
+// field, which only compiles when the trait is `Debug`. Do not remove
+// the bound when factoring this trait out -- it will break
+// `derive(Debug)` on `SamplerConfig`.
+pub trait Auth401AttributionCallback: Send + Sync + std::fmt::Debug {
+    /// Record a 401 attribution event for one logical 401 response.
+    ///
+    /// `sent_bearer_prefix` is the **first
+    /// [`SENT_BEARER_PREFIX_LEN`] characters** of the bearer that
+    /// was actually sent on the wire. The sampler extracts the
+    /// bearer from the `Authorization` header (or `x-api-key` for
+    /// Anthropic Messages API backends) and truncates it to the prefix
+    /// length **before crossing this trait boundary** -- the full
+    /// bearer never leaves [`crate::SamplingClient`]. This is the
+    /// scrub-at-the-boundary invariant: even a misbehaving callback
+    /// implementation that logs `sent_bearer_prefix` directly leaks
+    /// only the prefix, never the full credential.
+    ///
+    /// `None` indicates the request had no bearer header at all
+    /// (distinct from "had a bearer that turned out to be stale").
+    fn record_401(&self, consumer: SamplingConsumer, sent_bearer_prefix: Option<&str>);
+}
+
+/// Shared, cheap-to-clone alias for the attribution callback.
+pub type SharedAttributionCallback = Arc<dyn Auth401AttributionCallback>;

crates/codegen/xai-grok-sampler/src/commands.rs

@@ -0,0 +1,47 @@
+//! 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`.
+
+use tokio::sync::oneshot;
+
+use xai_grok_sampling_types::{ConversationRequest, ConversationResponse, SamplingError};
+
+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.
+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.
+    Submit {
+        request_id: RequestId,
+        request: Box<ConversationRequest>,
+        config: Option<Box<SamplerConfig>>,
+        completion_tx: Option<
+            oneshot::Sender<Result<(ConversationResponse, InferenceLatencyStats), SamplingError>>,
+        >,
+    },
+
+    /// Cancel an in-flight request.
+    Cancel { request_id: RequestId },
+
+    /// Update the default sampling config (model switch, auth refresh).
+    UpdateConfig { config: Box<SamplerConfig> },
+
+    /// Query: is a specific request still in flight?
+    IsActive {
+        request_id: RequestId,
+        reply: oneshot::Sender<bool>,
+    },
+
+    /// Query: how many requests are in flight?
+    ActiveCount { reply: oneshot::Sender<usize> },
+}

crates/codegen/xai-grok-sampler/src/config.rs

@@ -0,0 +1,246 @@
+//! Sampler configuration types.
+//!
+//! [`SamplerConfig`] is the per-request configuration handed to the
+//! sampler. It deliberately does **not** alias
+//! `xai_grok_sampling_types::SamplingConfig` so that the sampler crate
+//! avoids transitive dependencies on shell-specific types
+//! (`xai-grok-tools`, etc.).
+
+use indexmap::IndexMap;
+use serde::{Deserialize, Serialize};
+use xai_grok_sampling_types::{
+    ApiBackend, CompactionAtTokens, CompactionsRemaining, DoomLoopRecoveryPolicy, ReasoningEffort,
+};
+
+use crate::attribution::SharedAttributionCallback;
+use crate::retry::{DEFAULT_MAX_RETRIES, RATE_LIMIT_RETRY_THRESHOLD};
+
+#[derive(Debug, Clone, Copy, Serialize, Deserialize, PartialEq, Eq, Default)]
+#[serde(rename_all = "snake_case")]
+pub enum AuthScheme {
+    #[default]
+    Bearer,
+    XApiKey,
+}
+
+/// All knobs that control a single sampling request.
+///
+/// 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 `xai-grok-shell`
+///
+/// `SamplerConfig` is the single source of truth for sampler
+/// configuration. The shell builds it directly (see
+/// `agent::config::resolve_model_to_sampling_config` and
+/// `session::acp_session::SessionActor::reconstruct_full_config`) by
+/// composing chat-state's `xai_grok_sampling_types::SamplingConfig`
+/// with `Credentials` (api key, client version).
+///
+/// URL-derived request headers (e.g. `X-XAI-Token-Auth` for the
+/// cli-chat-proxy) 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.
+#[derive(Debug, Clone, Serialize, Deserialize)]
+pub struct SamplerConfig {
+    pub api_key: Option<String>,
+    pub base_url: String,
+    pub model: String,
+    pub max_completion_tokens: Option<u32>,
+    pub temperature: Option<f32>,
+    pub top_p: Option<f32>,
+    pub api_backend: ApiBackend,
+    #[serde(default)]
+    pub auth_scheme: AuthScheme,
+    /// Extra request headers applied verbatim. The sampler never inspects
+    /// the URL to derive headers; callers (the session) inject proxy auth
+    /// and other access headers here before constructing the config.
+    pub extra_headers: IndexMap<String, String>,
+    /// Total context window size in tokens. The sampler does not enforce
+    /// it; it is informational metadata used by the session for compaction
+    /// decisions.
+    pub context_window: u64,
+    pub force_http1: bool,
+    pub max_retries: Option<u32>,
+    pub stream_tool_calls: bool,
+    pub idle_timeout_secs: Option<u64>,
+
+    // Reasoning effort
+    pub reasoning_effort: Option<ReasoningEffort>,
+
+    // Client identity
+    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
+    /// 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.
+    ///
+    /// `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.
+    #[serde(skip)]
+    pub attribution_callback: Option<SharedAttributionCallback>,
+
+    /// Live bearer resolve per request. `None` uses construction-time `api_key`.
+    #[serde(skip)]
+    pub bearer_resolver: Option<SharedBearerResolver>,
+
+    #[serde(default)]
+    pub supports_backend_search: bool,
+
+    /// Per-model config for the `x-compactions-remaining` header; `None` disables it.
+    #[serde(default)]
+    pub compactions_remaining: Option<CompactionsRemaining>,
+
+    /// Per-model config for the `x-compaction-at` header; `None` disables it.
+    #[serde(default)]
+    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-grok-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>,
+
+    /// Per-request header injector (e.g. OTel traceparent). Called in `post()`.
+    #[serde(skip)]
+    pub header_injector: Option<SharedHeaderInjector>,
+}
+
+impl Default for SamplerConfig {
+    /// Empty defaults so callers can use `..Default::default()` and
+    /// new fields don't ripple through every literal site.
+    fn default() -> Self {
+        Self {
+            api_key: None,
+            base_url: String::new(),
+            model: String::new(),
+            max_completion_tokens: None,
+            temperature: None,
+            top_p: None,
+            api_backend: ApiBackend::default(),
+            auth_scheme: AuthScheme::default(),
+            extra_headers: IndexMap::new(),
+            context_window: 0,
+            force_http1: false,
+            max_retries: None,
+            stream_tool_calls: false,
+            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,
+            compactions_remaining: None,
+            compaction_at_tokens: None,
+            doom_loop_recovery: None,
+            header_injector: None,
+        }
+    }
+}
+
+/// Cheap sync read of the current bearer for [`SamplerConfig::bearer_resolver`].
+pub trait BearerResolver: Send + Sync + std::fmt::Debug {
+    fn current_bearer(&self) -> Option<String>;
+}
+
+pub type SharedBearerResolver = std::sync::Arc<dyn BearerResolver>;
+
+/// Per-request header injection (e.g. OTel `traceparent`).
+pub trait HeaderInjector: Send + Sync + std::fmt::Debug {
+    fn inject(&self, headers: &mut reqwest::header::HeaderMap);
+}
+
+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.
+    pub rate_limit_retry_threshold: u32,
+}
+
+impl Default for RetryPolicy {
+    fn default() -> Self {
+        Self {
+            max_retries: DEFAULT_MAX_RETRIES,
+            rate_limit_retry_threshold: RATE_LIMIT_RETRY_THRESHOLD,
+        }
+    }
+}
+
+/// Identity of the client that originated the request, used for
+/// User-Agent rendering. The shell layer composes this with platform
+/// info into a final UA string.
+#[derive(Clone, Debug, Eq, PartialEq, Serialize, Deserialize)]
+pub struct OriginClientInfo {
+    pub product: String,
+    pub version: Option<String>,
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    #[test]
+    fn retry_policy_defaults() {
+        let policy = RetryPolicy::default();
+        assert_eq!(policy.max_retries, DEFAULT_MAX_RETRIES);
+        assert_eq!(
+            policy.rate_limit_retry_threshold,
+            RATE_LIMIT_RETRY_THRESHOLD
+        );
+    }
+
+    /// 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();
+        stripped
+            .as_object_mut()
+            .unwrap()
+            .remove("doom_loop_recovery");
+        let config: SamplerConfig = serde_json::from_value(stripped).unwrap();
+        assert!(config.doom_loop_recovery.is_none());
+
+        let with_policy = SamplerConfig {
+            doom_loop_recovery: Some(DoomLoopRecoveryPolicy {
+                max_threshold: 8,
+                max_retries: 2,
+            }),
+            ..Default::default()
+        };
+        let round_tripped: SamplerConfig =
+            serde_json::from_value(serde_json::to_value(&with_policy).unwrap()).unwrap();
+        assert_eq!(
+            round_tripped.doom_loop_recovery,
+            with_policy.doom_loop_recovery
+        );
+    }
+}

crates/codegen/xai-grok-sampler/src/doom_loop.rs

@@ -0,0 +1,229 @@
+//! Per-request transport for server-reported doom-loop signals.
+//!
+//! The wire shapes and tolerant parsers live in
+//! [`xai_grok_sampling_types::doom_loop`]; this module only moves the parsed
+//! signals across the layer boundary: the Layer-1 SSE decoder in
+//! [`crate::client`] records them as raw payloads arrive, and the Layer-2
+//! transform in [`crate::stream::responses`] drains them into the final
+//! `ConversationResponse`.
+
+use std::sync::{Arc, Mutex};
+
+use xai_grok_sampling_types::doom_loop::{
+    DOOM_LOOP_CHECK_EVENT_TYPE, DoomLoopPeek, DoomLoopRecoveryPolicy, DoomLoopSignal,
+    peek_doom_loop,
+};
+
+/// Cheap-to-clone accumulator shared between the SSE decode closure and the
+/// stream transform of one request attempt. Created fresh per attempt so
+/// signals from a failed attempt can never leak into the next one. Carries
+/// the policy so the stream transform can judge confidence for the
+/// mid-stream abort; the retry loop disarms the abort once the recovery
+/// budget is spent so the final attempt completes and can be accepted.
+#[derive(Clone, Debug, Default)]
+pub struct DoomLoopSignalCollector {
+    inner: Arc<Mutex<CollectorState>>,
+}
+
+#[derive(Debug, Default)]
+struct CollectorState {
+    signals: Vec<DoomLoopSignal>,
+    malformed_logged: bool,
+    policy: DoomLoopRecoveryPolicy,
+    // Inverted so `derive(Default)` starts attempts armed.
+    abort_disarmed: bool,
+}
+
+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() {
+            state.policy = policy;
+        }
+        collector
+    }
+
+    /// Stop the mid-stream abort for this attempt; signals keep recording.
+    pub(crate) fn disarm_abort(&self) {
+        if let Ok(mut state) = self.inner.lock() {
+            state.abort_disarmed = true;
+        }
+    }
+
+    /// While armed: the raw labels of the confident signals recorded so far
+    /// (non-draining), or `None` when there is nothing to act on.
+    pub(crate) fn abort_triggers(&self) -> Option<Vec<String>> {
+        let state = self.inner.lock().ok()?;
+        if state.abort_disarmed {
+            return None;
+        }
+        let confident = state.policy.confident_triggers(&state.signals);
+        (!confident.is_empty()).then_some(confident)
+    }
+
+    /// Inspect a raw SSE frame. Returns `true` when the frame is the
+    /// non-standard `response.doom_loop_check` event — by its SSE `event:`
+    /// name or its payload `type` — which the caller must swallow;
+    /// forwarding it would fail typed deserialization. Reported triggers
+    /// (mid-stream or on the terminal response object) are recorded,
+    /// deduplicated by raw label. Never fails.
+    pub(crate) fn absorb(&self, event_name: &str, data: &str) -> bool {
+        // The name check keeps a check event with an unparseable payload
+        // from ever reaching the typed parser.
+        let named = event_name == DOOM_LOOP_CHECK_EVENT_TYPE;
+        let (signals, swallow) = match peek_doom_loop(data) {
+            DoomLoopPeek::CheckEvent(signals) => (signals, true),
+            DoomLoopPeek::ResponseField(signals) => (signals, false),
+            DoomLoopPeek::None => {
+                if named {
+                    self.log_malformed_once();
+                }
+                return named;
+            }
+        };
+        if signals.is_empty() {
+            self.log_malformed_once();
+        } else {
+            self.record(signals);
+        }
+        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),
+            Err(_) => Vec::new(),
+        }
+    }
+
+    fn record(&self, signals: Vec<DoomLoopSignal>) {
+        let Ok(mut state) = self.inner.lock() else {
+            return;
+        };
+        // Cumulative sets are re-sent as they grow; the raw label is the
+        // stable identity. Linear scan is fine for these tiny sets.
+        for signal in signals {
+            if !state.signals.iter().any(|s| s.raw == signal.raw) {
+                state.signals.push(signal);
+            }
+        }
+    }
+
+    /// 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;
+        };
+        if !state.malformed_logged {
+            state.malformed_logged = true;
+            tracing::debug!("doom-loop check payload malformed or empty; ignoring");
+        }
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+    use xai_grok_sampling_types::doom_loop::{
+        DoomLoopSignalKind, SAMPLE_CHECK_EVENT_DATA, SAMPLE_CHECK_EVENT_DATA_CUMULATIVE,
+    };
+
+    #[test]
+    fn absorb_swallows_check_event_and_records_signals() {
+        let collector = DoomLoopSignalCollector::default();
+        assert!(collector.absorb(DOOM_LOOP_CHECK_EVENT_TYPE, SAMPLE_CHECK_EVENT_DATA));
+        let signals = collector.take();
+        assert_eq!(signals.len(), 1);
+        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();
+        assert!(collector.absorb("message", SAMPLE_CHECK_EVENT_DATA));
+        assert_eq!(collector.take().len(), 1);
+    }
+
+    #[test]
+    fn absorb_dedupes_cumulative_sets_by_raw_label() {
+        let collector = DoomLoopSignalCollector::default();
+        assert!(collector.absorb(DOOM_LOOP_CHECK_EVENT_TYPE, SAMPLE_CHECK_EVENT_DATA));
+        assert!(collector.absorb(
+            DOOM_LOOP_CHECK_EVENT_TYPE,
+            SAMPLE_CHECK_EVENT_DATA_CUMULATIVE
+        ));
+        let signals = collector.take();
+        assert_eq!(signals.len(), 2);
+        assert_eq!(signals[0].raw, "tail_repetition:4@response");
+        assert_eq!(signals[1].raw, "tail_repetition:2@response");
+    }
+
+    #[test]
+    fn absorb_forwards_ordinary_and_terminal_payloads() {
+        let collector = DoomLoopSignalCollector::default();
+        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);
+    }
+
+    #[test]
+    fn malformed_check_event_swallowed_without_signals() {
+        let collector = DoomLoopSignalCollector::default();
+        let wrong_type =
+            r#"{"type":"response.doom_loop_check","doom_loop_check":{"triggers":"nope"}}"#;
+        assert!(collector.absorb(DOOM_LOOP_CHECK_EVENT_TYPE, wrong_type));
+        assert!(collector.absorb("message", r#"{"type":"response.doom_loop_check"}"#));
+        assert!(collector.take().is_empty());
+    }
+
+    /// A frame with the check event's SSE name but an unparseable payload
+    /// (non-JSON, or JSON without the `type` tag) must still be swallowed —
+    /// forwarding it would fail the typed parse and the whole attempt.
+    #[test]
+    fn named_event_with_garbage_payload_still_swallowed() {
+        let collector = DoomLoopSignalCollector::default();
+        assert!(collector.absorb(DOOM_LOOP_CHECK_EVENT_TYPE, "not json at all"));
+        assert!(collector.absorb(DOOM_LOOP_CHECK_EVENT_TYPE, r#"{"no_type_tag":true}"#));
+        assert!(collector.take().is_empty());
+    }
+
+    #[test]
+    fn take_drains_once() {
+        let collector = DoomLoopSignalCollector::default();
+        collector.absorb(DOOM_LOOP_CHECK_EVENT_TYPE, SAMPLE_CHECK_EVENT_DATA);
+        assert!(!collector.take().is_empty());
+        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"]}}"#;
+
+        let collector = DoomLoopSignalCollector::new(DoomLoopRecoveryPolicy::default());
+        // Non-confident channel: recorded but not actionable.
+        assert!(collector.absorb(DOOM_LOOP_CHECK_EVENT_TYPE, SAMPLE_CHECK_EVENT_DATA));
+        assert!(collector.abort_triggers().is_none());
+
+        assert!(collector.absorb(DOOM_LOOP_CHECK_EVENT_TYPE, confident));
+        assert_eq!(
+            collector.abort_triggers(),
+            Some(vec!["tail_repetition:8@thinking".to_string()])
+        );
+        // Non-draining: probing twice and taking afterwards both work.
+        assert!(collector.abort_triggers().is_some());
+
+        collector.disarm_abort();
+        assert!(collector.abort_triggers().is_none());
+        assert_eq!(collector.take().len(), 2, "recording survives the disarm");
+    }
+}

crates/codegen/xai-grok-sampler/src/events.rs

@@ -0,0 +1,367 @@
+//! Outbound events emitted by the sampler.
+
+use serde::{Deserialize, Serialize};
+
+use xai_grok_sampling_types::{
+    ConversationResponse, EmptyResponseContext, ResponseModelMetadata, SamplingError,
+};
+
+use crate::metrics::InferenceLatencyStats;
+use crate::types::RequestId;
+
+/// Which content channel a token belongs to.
+///
+/// Extensible — adding a new channel (e.g., `Planning`) only requires a
+/// new variant here, not new [`SamplingEvent`] variants. Mirrors the
+/// agentic-sampler's `AgentChannel` pattern.
+#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
+pub enum SamplingChannel {
+    Text,
+    Reasoning,
+}
+
+/// Events emitted by the sampler for a single in-flight request.
+///
+/// Sent on the shared event channel that callers subscribe to. The
+/// session translates these into ACP notifications.
+#[derive(Debug, Clone)]
+pub enum SamplingEvent {
+    /// HTTP stream established, headers read. Emitted before any content.
+    StreamStarted {
+        request_id: RequestId,
+        timestamp_ms: i64,
+    },
+
+    /// First content token received for a request.
+    FirstToken { request_id: RequestId },
+
+    /// Content token in a named channel (text or reasoning).
+    ChannelToken {
+        request_id: RequestId,
+        channel: SamplingChannel,
+        text: String,
+        chunk_index: u64,
+    },
+
+    /// Streaming delta carrying a fragment of a tool call.
+    ///
+    /// Emitted by the L2 transforms (Chat Completions, Responses, Messages)
+    /// per-chunk as the model streams tool-call arguments. Any single
+    /// `arguments_delta` is NOT necessarily valid JSON in isolation.
+    ToolCallDelta {
+        request_id: RequestId,
+        tool_index: u32,
+        id: Option<String>,
+        name: Option<String>,
+        arguments_delta: Option<String>,
+    },
+
+    /// Streaming completed successfully.
+    Completed {
+        request_id: RequestId,
+        response: Box<ConversationResponse>,
+        metrics: InferenceLatencyStats,
+    },
+
+    /// Request is being retried.
+    Retrying {
+        request_id: RequestId,
+        attempt: u32,
+        max_retries: u32,
+        /// Typed retry class so consumers never have to sniff `reason`
+        /// (e.g. the shell's doom-loop recovery counter).
+        kind: SamplingErrorKind,
+        reason: String,
+        /// Doom-loop telemetry payload when `kind == DoomLoopDetected`:
+        /// raw trigger labels + the chunk index the mid-stream abort fired
+        /// at (`None` for terminal-response detections). Labels only.
+        doom_loop_triggers: Option<Vec<String>>,
+        doom_loop_aborted_at_chunk: Option<u64>,
+    },
+
+    /// Request failed (after exhausting retries or non-retryable error).
+    Failed {
+        request_id: RequestId,
+        error: SamplingErrorInfo,
+    },
+
+    /// Model metadata received from response headers.
+    ModelMetadata {
+        request_id: RequestId,
+        metadata: ResponseModelMetadata,
+    },
+
+    /// A backend-hosted tool call has started execution on the server
+    /// (e.g., web search is in progress). The client does NOT execute
+    /// these — the backend's agentic sampler handles them.
+    BackendToolCallStarted {
+        request_id: RequestId,
+        call_id: String,
+        name: String,
+    },
+
+    /// A backend-hosted tool call has completed execution on the server.
+    BackendToolCallCompleted {
+        request_id: RequestId,
+        call_id: String,
+        name: String,
+        /// Structured result data from the backend tool (tool-specific).
+        /// For web search: `{"query": "...", "sources": [{"url": "..."}, ...]}`
+        result: Option<serde_json::Value>,
+    },
+}
+
+/// Serializable mirror of [`SamplingError`].
+///
+/// The rich `SamplingError` carries non-serializable inner values
+/// (`reqwest::Error`, `serde_json::Error`) so it cannot cross a network
+/// boundary. `SamplingErrorInfo` extracts the bits that downstream
+/// consumers (UIs, gRPC adapters) actually need.
+#[derive(Debug, Clone, Serialize, Deserialize)]
+pub struct SamplingErrorInfo {
+    pub kind: SamplingErrorKind,
+    pub status_code: Option<u16>,
+    pub message: String,
+    pub is_retryable: bool,
+    pub retry_after_secs: Option<u64>,
+    pub model_metadata: Option<ResponseModelMetadata>,
+    /// Present only when `kind == EmptyResponse`. Carries the structured
+    /// context from the L2 stream so downstream consumers can distinguish
+    /// reasoning-only completions from transport failures.
+    #[serde(default, skip_serializing_if = "Option::is_none")]
+    pub empty_response_context: Option<EmptyResponseContext>,
+    /// Present only when `kind == DoomLoopDetected`. Raw trigger labels
+    /// (never generation content) so the retry loop can reconstruct the
+    /// rich error from a synthesized L2 failure.
+    #[serde(default, skip_serializing_if = "Option::is_none")]
+    pub doom_loop_triggers: Option<Vec<String>>,
+    /// Stream chunk index the mid-stream doom-loop abort fired at.
+    /// Telemetry only; `None` for terminal-response detections.
+    #[serde(default, skip_serializing_if = "Option::is_none")]
+    pub doom_loop_aborted_at_chunk: Option<u64>,
+}
+
+/// Coarse-grained classification of a sampling failure.
+///
+/// Intentionally 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
+/// the metadata and decides whether to compact.
+#[derive(Debug, Clone, Copy, Serialize, Deserialize, PartialEq, Eq)]
+pub enum SamplingErrorKind {
+    Auth,
+    Http,
+    Api,
+    Serialization,
+    IdleTimeout,
+    RateLimited,
+    EmptyResponse,
+    MaxTokensTruncation,
+    DoomLoopDetected,
+}
+
+impl SamplingErrorKind {
+    /// Stable, lowercase string form suitable for telemetry tags
+    /// (e.g., analytics `error_type` columns and signals histograms).
+    /// Mirrors the strings used in the shell's
+    /// `stream_conversation_with_retries` error classifier so tags stay
+    /// consistent across surfaces.
+    pub fn as_str(self) -> &'static str {
+        match self {
+            SamplingErrorKind::Auth => "auth",
+            SamplingErrorKind::Http => "http",
+            SamplingErrorKind::Api => "api",
+            SamplingErrorKind::Serialization => "serialization",
+            SamplingErrorKind::IdleTimeout => "idle_timeout",
+            SamplingErrorKind::RateLimited => "rate_limited",
+            SamplingErrorKind::EmptyResponse => "empty_response",
+            SamplingErrorKind::MaxTokensTruncation => "max_tokens_truncation",
+            SamplingErrorKind::DoomLoopDetected => "doom_loop_detected",
+        }
+    }
+}
+
+impl From<&SamplingError> for SamplingErrorInfo {
+    fn from(err: &SamplingError) -> Self {
+        let is_retryable = err.is_retryable();
+        let message = err.to_string();
+
+        let (kind, status_code, retry_after_secs, model_metadata) = match err {
+            SamplingError::Auth(_) => (SamplingErrorKind::Auth, None, None, None),
+            SamplingError::InvalidConfiguration(_) => (SamplingErrorKind::Api, None, None, None),
+            SamplingError::Http(_) => (SamplingErrorKind::Http, None, None, None),
+            SamplingError::Serialization(_) => (SamplingErrorKind::Serialization, None, None, None),
+            SamplingError::Api {
+                status,
+                model_metadata,
+                retry_after_secs,
+                ..
+            } => {
+                let kind = if err.is_rate_limited() {
+                    SamplingErrorKind::RateLimited
+                } else {
+                    SamplingErrorKind::Api
+                };
+                (
+                    kind,
+                    Some(status.as_u16()),
+                    *retry_after_secs,
+                    model_metadata.clone(),
+                )
+            }
+            SamplingError::EventStreamError(_) => (SamplingErrorKind::Http, None, None, None),
+            SamplingError::StreamError { .. } => (SamplingErrorKind::Api, None, None, None),
+            SamplingError::IdleTimeout { .. } => (SamplingErrorKind::IdleTimeout, None, None, None),
+            SamplingError::EmptyResponse { .. } => {
+                (SamplingErrorKind::EmptyResponse, None, None, None)
+            }
+            SamplingError::MaxTokensTruncation => {
+                (SamplingErrorKind::MaxTokensTruncation, None, None, None)
+            }
+            SamplingError::DoomLoopDetected { .. } => {
+                (SamplingErrorKind::DoomLoopDetected, None, None, None)
+            }
+        };
+
+        let empty_response_context = match err {
+            SamplingError::EmptyResponse { context } => Some(context.clone()),
+            _ => None,
+        };
+        let (doom_loop_triggers, doom_loop_aborted_at_chunk) = match err {
+            SamplingError::DoomLoopDetected {
+                triggers,
+                aborted_at_chunk,
+            } => (Some(triggers.clone()), *aborted_at_chunk),
+            _ => (None, None),
+        };
+
+        Self {
+            kind,
+            status_code,
+            message,
+            is_retryable,
+            retry_after_secs,
+            model_metadata,
+            empty_response_context,
+            doom_loop_triggers,
+            doom_loop_aborted_at_chunk,
+        }
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+    use reqwest::StatusCode;
+
+    #[test]
+    fn auth_variant_classified_as_auth() {
+        let err = SamplingError::Auth("bad token".into());
+        let info = SamplingErrorInfo::from(&err);
+        assert_eq!(info.kind, SamplingErrorKind::Auth);
+        assert_eq!(info.status_code, None);
+        assert!(!info.is_retryable);
+        assert_eq!(info.retry_after_secs, None);
+        assert!(info.model_metadata.is_none());
+        assert!(info.message.contains("bad token"));
+    }
+
+    #[test]
+    fn invalid_configuration_classified_as_api() {
+        let err = SamplingError::InvalidConfiguration("missing model");
+        let info = SamplingErrorInfo::from(&err);
+        assert_eq!(info.kind, SamplingErrorKind::Api);
+        assert_eq!(info.status_code, None);
+        assert!(!info.is_retryable);
+    }
+
+    #[test]
+    fn serialization_variant_classified_as_serialization() {
+        let json_err = serde_json::from_str::<i32>("not a number").unwrap_err();
+        let err: SamplingError = json_err.into();
+        let info = SamplingErrorInfo::from(&err);
+        assert_eq!(info.kind, SamplingErrorKind::Serialization);
+        assert!(!info.is_retryable);
+    }
+
+    #[test]
+    fn api_500_classified_as_api_and_retryable() {
+        let err = SamplingError::Api {
+            status: StatusCode::INTERNAL_SERVER_ERROR,
+            message: "boom".into(),
+            model_metadata: None,
+            retry_after_secs: None,
+            should_retry: None,
+        };
+        let info = SamplingErrorInfo::from(&err);
+        assert_eq!(info.kind, SamplingErrorKind::Api);
+        assert_eq!(info.status_code, Some(500));
+        assert!(info.is_retryable, "5xx should be retryable");
+    }
+
+    #[test]
+    fn api_429_classified_as_rate_limited_and_extracts_retry_after() {
+        let err = SamplingError::Api {
+            status: StatusCode::TOO_MANY_REQUESTS,
+            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);
+        assert_eq!(info.status_code, Some(429));
+        assert_eq!(info.retry_after_secs, Some(15));
+        assert!(info.is_retryable, "429 should be retryable");
+    }
+
+    #[test]
+    fn api_400_classified_as_api_and_not_retryable() {
+        let err = SamplingError::Api {
+            status: StatusCode::BAD_REQUEST,
+            message: "context window exceeded".into(),
+            model_metadata: Some(ResponseModelMetadata {
+                context_window: Some(8000),
+                ..Default::default()
+            }),
+            retry_after_secs: None,
+            should_retry: None,
+        };
+        let info = SamplingErrorInfo::from(&err);
+        assert_eq!(info.kind, SamplingErrorKind::Api);
+        assert_eq!(info.status_code, Some(400));
+        assert!(!info.is_retryable, "4xx (non-429) should not be retryable");
+        let metadata = info.model_metadata.expect("metadata preserved");
+        assert_eq!(metadata.context_window, Some(8000));
+    }
+
+    #[test]
+    fn event_stream_error_classified_as_http_and_retryable() {
+        let err = SamplingError::EventStreamError("conn reset".into());
+        let info = SamplingErrorInfo::from(&err);
+        assert_eq!(info.kind, SamplingErrorKind::Http);
+        assert!(info.is_retryable);
+    }
+
+    #[test]
+    fn stream_error_classified_as_api_and_retryable() {
+        let err = SamplingError::StreamError {
+            error_type: "server_error".into(),
+            message: "transient".into(),
+        };
+        let info = SamplingErrorInfo::from(&err);
+        assert_eq!(info.kind, SamplingErrorKind::Api);
+        assert_eq!(info.status_code, None);
+        assert!(info.is_retryable, "stream errors should be retryable");
+    }
+
+    #[test]
+    fn idle_timeout_classified_as_idle_timeout_and_not_retryable() {
+        let err = SamplingError::IdleTimeout { elapsed_secs: 300 };
+        let info = SamplingErrorInfo::from(&err);
+        assert_eq!(info.kind, SamplingErrorKind::IdleTimeout);
+        assert!(!info.is_retryable);
+        assert!(info.message.contains("300s"));
+    }
+}

crates/codegen/xai-grok-sampler/src/handle.rs

@@ -0,0 +1,157 @@
+//! Public handle for talking to the sampler actor.
+
+use tokio::sync::{mpsc, oneshot};
+
+use xai_grok_sampling_types::{ConversationRequest, ConversationResponse, SamplingError};
+
+use crate::commands::SamplerCommand;
+use crate::config::SamplerConfig;
+use crate::metrics::InferenceLatencyStats;
+use crate::types::RequestId;
+
+/// Cheaply-cloneable handle to the sampler actor.
+///
+/// Internally just an `mpsc::UnboundedSender<SamplerCommand>`. All
+/// methods are non-blocking (fire-and-forget) except for the
+/// `*_async` queries which return a future awaiting an
+/// `oneshot::Receiver`.
+#[derive(Clone)]
+pub struct SamplerHandle {
+    cmd_tx: mpsc::UnboundedSender<SamplerCommand>,
+}
+
+impl SamplerHandle {
+    /// Construct a handle from a command sender. `pub(crate)` because
+    /// only [`SamplerActor::spawn`](crate::actor::SamplerActor::spawn)
+    /// produces one of these.
+    pub(crate) fn new(cmd_tx: mpsc::UnboundedSender<SamplerCommand>) -> Self {
+        Self { cmd_tx }
+    }
+
+    /// Create a no-op handle that discards all commands.
+    ///
+    /// Useful for tests and callers that need a `SamplerHandle` field
+    /// before the actor is wired up. Mirrors
+    /// [`HunkTrackerHandle::noop`](https://docs.rs/xai-hunk-tracker).
+    pub fn noop() -> Self {
+        let (cmd_tx, _cmd_rx) = mpsc::unbounded_channel();
+        // Receiver is dropped immediately; sends will fail but every
+        // send-site uses `let _ = ...` so that is fine.
+        Self { cmd_tx }
+    }
+
+    /// Submit a sampling request. Fire-and-forget -- results arrive
+    /// via the shared event channel.
+    pub fn submit(&self, request_id: RequestId, request: ConversationRequest) {
+        let _ = self.cmd_tx.send(SamplerCommand::Submit {
+            request_id,
+            request: Box::new(request),
+            config: None,
+            completion_tx: None,
+        });
+    }
+
+    /// Submit a sampling request with an explicit per-request config
+    /// override (e.g., a different model than the actor's default).
+    pub fn submit_with_config(
+        &self,
+        request_id: RequestId,
+        request: ConversationRequest,
+        config: SamplerConfig,
+    ) {
+        let _ = self.cmd_tx.send(SamplerCommand::Submit {
+            request_id,
+            request: Box::new(request),
+            config: Some(Box::new(config)),
+            completion_tx: None,
+        });
+    }
+
+    /// Cancel an in-flight request. No-op if the request id is
+    /// unknown (already finished or never submitted).
+    pub fn cancel(&self, request_id: RequestId) {
+        let _ = self.cmd_tx.send(SamplerCommand::Cancel { request_id });
+    }
+
+    /// Update the default sampling config (e.g., after model switch
+    /// or auth refresh). The next request submitted without an
+    /// override will use it.
+    pub fn update_config(&self, config: SamplerConfig) {
+        let _ = self.cmd_tx.send(SamplerCommand::UpdateConfig {
+            config: Box::new(config),
+        });
+    }
+
+    /// Query whether a request is still in flight. Returns `false`
+    /// for unknown / finished / cancelled ids.
+    pub async fn is_active(&self, request_id: RequestId) -> bool {
+        let (reply_tx, reply_rx) = oneshot::channel();
+        let _ = self.cmd_tx.send(SamplerCommand::IsActive {
+            request_id,
+            reply: reply_tx,
+        });
+        reply_rx.await.unwrap_or(false)
+    }
+
+    /// Query the number of in-flight requests. Returns 0 if the
+    /// actor has been shut down.
+    pub async fn active_count(&self) -> usize {
+        let (reply_tx, reply_rx) = oneshot::channel();
+        let _ = self
+            .cmd_tx
+            .send(SamplerCommand::ActiveCount { reply: reply_tx });
+        reply_rx.await.unwrap_or(0)
+    }
+
+    /// Submit a request and await its completion. Events still flow
+    /// to the shared channel for live UI updates -- this method just
+    /// additionally awaits the per-request completion oneshot so the
+    /// caller gets a clean `Result` without filtering events.
+    ///
+    /// Used by sequential callers like compaction / summary /
+    /// `/btw` side questions.
+    pub async fn submit_and_collect(
+        &self,
+        request_id: RequestId,
+        request: ConversationRequest,
+    ) -> Result<(ConversationResponse, InferenceLatencyStats), SamplingError> {
+        // RAII guard: when this future is dropped (cancel, panic, or normal return),
+        // tell the sampler actor to cancel the in-flight request_id. No-op if the
+        // actor already finished and removed it from its active set.
+        struct CancelOnDrop {
+            cmd_tx: mpsc::UnboundedSender<SamplerCommand>,
+            request_id: RequestId,
+        }
+        impl Drop for CancelOnDrop {
+            fn drop(&mut self) {
+                // fire-and-forget the send.
+                let _ = self.cmd_tx.send(SamplerCommand::Cancel {
+                    request_id: self.request_id.clone(),
+                });
+            }
+        }
+
+        let (completion_tx, completion_rx) = oneshot::channel();
+        let cancel_id = request_id.clone();
+
+        // Only arm the guard if Submit actually reached the actor.
+        let _guard = self
+            .cmd_tx
+            .send(SamplerCommand::Submit {
+                request_id,
+                request: Box::new(request),
+                config: None,
+                completion_tx: Some(completion_tx),
+            })
+            .ok()
+            .map(|_| CancelOnDrop {
+                cmd_tx: self.cmd_tx.clone(),
+                request_id: cancel_id,
+            });
+        completion_rx.await.unwrap_or_else(|_| {
+            Err(SamplingError::Auth(
+                "sampler actor dropped before completion".to_string(),
+            ))
+        })
+    }
+}

crates/codegen/xai-grok-sampler/src/lib.rs

@@ -0,0 +1,54 @@
+//! xai-grok-sampler - Actor-based sampling layer for xAI grok.
+//!
+//! This crate extracts the HTTP streaming + retry logic out of
+//! `xai-grok-shell`'s session actor into a standalone, reusable
+//! component built on the same actor pattern as `xai-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;
+pub mod client;
+pub mod commands;
+pub mod config;
+pub mod doom_loop;
+pub mod events;
+pub mod handle;
+pub mod metrics;
+pub mod retry;
+pub mod sampling_log;
+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,
+};
+pub use client::{ApiBackend, SamplingClient, user_agent_string_for};
+pub use config::{
+    AuthScheme, BearerResolver, HeaderInjector, OriginClientInfo, RetryPolicy, SamplerConfig,
+    SharedBearerResolver, SharedHeaderInjector,
+};
+pub use doom_loop::DoomLoopSignalCollector;
+pub use events::{SamplingChannel, SamplingErrorInfo, SamplingErrorKind, SamplingEvent};
+pub use handle::SamplerHandle;
+pub use metrics::{InferenceLatencyStats, compute_percentiles};
+pub use retry::{
+    DEFAULT_MAX_RETRIES, RATE_LIMIT_RETRY_THRESHOLD, RetryDecision, classify_error,
+    format_sampling_error, resolve_max_retries, retry_backoff_with_jitter,
+};
+pub use sampling_log::AuthInfo;
+pub use stream::{collect_response, stream_chat_completions, stream_messages, stream_responses};
+pub use types::RequestId;

crates/codegen/xai-grok-sampler/src/metrics.rs

@@ -0,0 +1,246 @@
+//! Per-response inference latency metrics.
+//!
+//! Captures token-level timing from streaming inference responses:
+//! TTFB, TTLB, and inter-token latency (ITL) statistics.
+
+use std::time::Instant;
+
+use serde::{Deserialize, Serialize};
+
+/// Compute percentiles from sorted intervals.
+///
+/// Returns (p50, p99, max, mean, sum) from a slice of sorted values.
+/// Panics if `sorted` is empty.
+pub fn compute_percentiles(sorted: &[u64]) -> (u64, u64, u64, u64, u64) {
+    let len = sorted.len();
+    assert!(len > 0, "Cannot compute percentiles from empty slice");
+
+    let p50 = sorted[len / 2];
+    let p99_idx = ((len as f64 * 0.99).ceil() as usize)
+        .saturating_sub(1)
+        .min(len - 1);
+    let p99 = sorted[p99_idx];
+    let max = sorted[len - 1];
+    let sum: u64 = sorted.iter().sum();
+    let mean = sum / len as u64;
+
+    (p50, p99, max, mean, sum)
+}
+
+/// Per-response inference latency metrics computed from chunk timestamps.
+#[derive(Debug, Clone, Default, Serialize, Deserialize)]
+pub struct InferenceLatencyStats {
+    /// Time to first content token (ms)
+    pub time_to_first_token_ms: Option<u64>,
+    /// Time to last byte / stream end (ms). Measured at stream exhaustion,
+    /// not at the last content chunk, so it includes trailing metadata chunks.
+    pub time_to_last_byte_ms: u64,
+    /// Number of content chunks received
+    pub chunk_count: u32,
+    /// Inter-token latency intervals (raw data for session aggregation)
+    pub itl_intervals_ms: Vec<u64>,
+    /// Inter-token latency: median (ms)
+    pub itl_p50_ms: Option<u64>,
+    /// Inter-token latency: 99th percentile (ms)
+    pub itl_p99_ms: Option<u64>,
+    /// Inter-token latency: maximum (ms)
+    pub itl_max_ms: Option<u64>,
+    /// Inter-token latency: mean (ms)
+    pub itl_mean_ms: Option<u64>,
+    /// Total request attempts (`1` = no retries); set by the retry loop on success.
+    pub attempts: u32,
+}
+
+impl InferenceLatencyStats {
+    /// Record the computed stats as fields on a tracing span.
+    pub fn record_on_span(&self, span: &tracing::Span) {
+        if let Some(ttfb) = self.time_to_first_token_ms {
+            span.record("ttfb_ms", ttfb);
+        }
+        span.record("ttlb_ms", self.time_to_last_byte_ms);
+        span.record("chunk_count", self.chunk_count);
+        if let Some(p50) = self.itl_p50_ms {
+            span.record("itl_p50_ms", p50);
+        }
+        if let Some(p99) = self.itl_p99_ms {
+            span.record("itl_p99_ms", p99);
+        }
+    }
+
+    /// Compute latency stats from chunk timestamps.
+    ///
+    /// # Arguments
+    /// * `stream_start` - `Instant::now()` captured before initiating the stream.
+    /// * `chunk_timestamps` - `Instant` recorded on each content-bearing chunk.
+    /// * `stream_end` - `Instant::now()` captured after the stream is fully exhausted
+    ///   (after trailing metadata/`[DONE]` chunks). Used for TTLB.
+    pub fn from_timestamps(
+        stream_start: Instant,
+        chunk_timestamps: &[Instant],
+        stream_end: Instant,
+    ) -> Self {
+        let ttlb = stream_end.duration_since(stream_start).as_millis() as u64;
+
+        if chunk_timestamps.is_empty() {
+            return Self {
+                time_to_last_byte_ms: ttlb,
+                ..Default::default()
+            };
+        }
+
+        let ttfb = chunk_timestamps[0].duration_since(stream_start);
+
+        // Compute inter-token intervals
+        let intervals: Vec<u64> = chunk_timestamps
+            .windows(2)
+            .map(|w| w[1].duration_since(w[0]).as_millis() as u64)
+            .collect();
+
+        let (itl_p50, itl_p99, itl_max, itl_mean) = if intervals.is_empty() {
+            (None, None, None, None)
+        } else {
+            let mut sorted = intervals.clone();
+            sorted.sort_unstable();
+            let (p50, p99, max, mean, _sum) = compute_percentiles(&sorted);
+            (Some(p50), Some(p99), Some(max), Some(mean))
+        };
+
+        Self {
+            time_to_first_token_ms: Some(ttfb.as_millis() as u64),
+            time_to_last_byte_ms: ttlb,
+            chunk_count: u32::try_from(chunk_timestamps.len()).unwrap_or(u32::MAX),
+            itl_intervals_ms: intervals,
+            itl_p50_ms: itl_p50,
+            itl_p99_ms: itl_p99,
+            itl_max_ms: itl_max,
+            itl_mean_ms: itl_mean,
+            attempts: 0,
+        }
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+    use std::time::Duration;
+
+    /// Helper: create an Instant offset from a base by a given duration.
+    fn offset(base: Instant, ms: u64) -> Instant {
+        base + Duration::from_millis(ms)
+    }
+
+    #[test]
+    fn test_empty_timestamps() {
+        let start = Instant::now();
+        let end = start + Duration::from_millis(500);
+
+        let stats = InferenceLatencyStats::from_timestamps(start, &[], end);
+
+        assert_eq!(stats.time_to_first_token_ms, None);
+        assert_eq!(stats.time_to_last_byte_ms, 500);
+        assert_eq!(stats.chunk_count, 0);
+        assert_eq!(stats.itl_p50_ms, None);
+        assert_eq!(stats.itl_p99_ms, None);
+        assert_eq!(stats.itl_max_ms, None);
+        assert_eq!(stats.itl_mean_ms, None);
+    }
+
+    #[test]
+    fn test_single_chunk() {
+        let start = Instant::now();
+        let chunks = vec![offset(start, 100)];
+        let end = offset(start, 200);
+
+        let stats = InferenceLatencyStats::from_timestamps(start, &chunks, end);
+
+        assert_eq!(stats.time_to_first_token_ms, Some(100));
+        assert_eq!(stats.time_to_last_byte_ms, 200);
+        assert_eq!(stats.chunk_count, 1);
+        // Single chunk => no intervals => no ITL stats
+        assert_eq!(stats.itl_p50_ms, None);
+        assert_eq!(stats.itl_p99_ms, None);
+        assert_eq!(stats.itl_max_ms, None);
+        assert_eq!(stats.itl_mean_ms, None);
+    }
+
+    #[test]
+    fn test_two_chunks() {
+        let start = Instant::now();
+        let chunks = vec![offset(start, 100), offset(start, 150)];
+        let end = offset(start, 200);
+
+        let stats = InferenceLatencyStats::from_timestamps(start, &chunks, end);
+
+        assert_eq!(stats.time_to_first_token_ms, Some(100));
+        assert_eq!(stats.time_to_last_byte_ms, 200);
+        assert_eq!(stats.chunk_count, 2);
+        // One interval of 50ms => p50=p99=max=mean=50
+        assert_eq!(stats.itl_p50_ms, Some(50));
+        assert_eq!(stats.itl_p99_ms, Some(50));
+        assert_eq!(stats.itl_max_ms, Some(50));
+        assert_eq!(stats.itl_mean_ms, Some(50));
+    }
+
+    #[test]
+    fn test_many_chunks() {
+        let start = Instant::now();
+        // 11 chunks: intervals are [10, 20, 30, 40, 50, 60, 70, 80, 90, 100]
+        let chunks: Vec<Instant> = (0..11)
+            .scan(100u64, |acc, i| {
+                let t = *acc;
+                *acc += (i + 1) * 10; // intervals: 10, 20, 30, ...
+                Some(offset(start, t))
+            })
+            .collect();
+        let end = offset(start, 1000);
+
+        let stats = InferenceLatencyStats::from_timestamps(start, &chunks, end);
+
+        assert_eq!(stats.time_to_first_token_ms, Some(100));
+        assert_eq!(stats.time_to_last_byte_ms, 1000);
+        assert_eq!(stats.chunk_count, 11);
+
+        // 10 intervals: [10, 20, 30, 40, 50, 60, 70, 80, 90, 100]
+        // sorted: same
+        // p50: intervals[5] = 60
+        assert_eq!(stats.itl_p50_ms, Some(60));
+        // p99_idx: ceil(10 * 0.99) - 1 = ceil(9.9) - 1 = 10 - 1 = 9, min(9, 9) = 9
+        // intervals[9] = 100
+        assert_eq!(stats.itl_p99_ms, Some(100));
+        assert_eq!(stats.itl_max_ms, Some(100));
+        // mean: (10+20+30+40+50+60+70+80+90+100) / 10 = 550 / 10 = 55
+        assert_eq!(stats.itl_mean_ms, Some(55));
+    }
+
+    #[test]
+    fn test_p99_does_not_overflow() {
+        let start = Instant::now();
+        // 101 chunks => 100 intervals (indices 0..99)
+        let chunks: Vec<Instant> = (0..101).map(|i| offset(start, 100 + i * 10)).collect();
+        let end = offset(start, 2000);
+
+        let stats = InferenceLatencyStats::from_timestamps(start, &chunks, end);
+
+        assert_eq!(stats.chunk_count, 101);
+        // 100 intervals, all 10ms
+        // p99_idx: ceil(100 * 0.99) - 1 = 100 - 1 = 99, min(99, 99) = 99 -> in bounds
+        assert_eq!(stats.itl_p99_ms, Some(10));
+        assert_eq!(stats.itl_max_ms, Some(10));
+        assert_eq!(stats.itl_p50_ms, Some(10));
+        assert_eq!(stats.itl_mean_ms, Some(10));
+    }
+
+    #[test]
+    fn test_ttlb_uses_stream_end_not_last_chunk() {
+        let start = Instant::now();
+        let chunks = vec![offset(start, 100), offset(start, 200)];
+        // stream_end is 500ms after start, well past the last chunk at 200ms
+        let end = offset(start, 500);
+
+        let stats = InferenceLatencyStats::from_timestamps(start, &chunks, end);
+
+        // TTLB should be 500 (from stream_end), not 200 (from last chunk)
+        assert_eq!(stats.time_to_last_byte_ms, 500);
+        assert_eq!(stats.time_to_first_token_ms, Some(100));
+    }
+}

crates/codegen/xai-grok-sampler/src/retry.rs

@@ -0,0 +1,856 @@
+//! Retry classification, backoff, and decision-making.
+//!
+//! Pure logic only: no I/O, no notifications, no logging side-effects.
+//! The actor (M4) wraps this with the actual retry loop.
+//!
+//! # Retry behavior summary
+//!
+//! **Retried** (up to [`DEFAULT_MAX_RETRIES`] = 15, ~6 min with 30s backoff cap):
+//! - 500, 502, 503, 504, 520 (server errors)
+//! - Connection errors (timeout, refused, reset)
+//! - `EventStreamError` / `StreamError` (mid-stream failures)
+//! - `EmptyResponse` (model returned no content/tool calls)
+//!
+//! **Retried with lower cap** ([`RATE_LIMIT_RETRY_THRESHOLD`] = 2):
+//! - 429 (rate limited) — avoids burning long waits
+//!
+//! **Special handling** (not counted against retry budget):
+//! - 413 / image processing errors → strip images and retry once
+//!
+//! **Not retried** (Fatal immediately):
+//! - 400, 401, 403, 404, 408, 422 (client errors)
+//! - `Auth` / `InvalidConfiguration` (credential/config issues)
+//! - `IdleTimeout` (model stuck, retry would stall again)
+//! - `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.
+
+use std::time::Duration;
+
+use xai_grok_sampling_types::SamplingError;
+
+/// After this many rate-limit (429) retries, escalate to the caller
+/// instead of waiting again. Rate-limit waits can be long and there is
+/// no point burning a long backoff just to be rate-limited again.
+pub const RATE_LIMIT_RETRY_THRESHOLD: u32 = 2;
+
+/// Default max retries when no env or model override is set.
+/// With 30s backoff cap this gives ~6 min of retry budget:
+/// retries 1-4 are exponential (2s+4s+8s+16s ≈ 30s), retries
+/// 5-15 are flat at ~30s each (≈ 5.5 min).
+pub const DEFAULT_MAX_RETRIES: u32 = 15;
+
+/// Resolve max API retries from an optional env override, model config,
+/// or default ([`DEFAULT_MAX_RETRIES`]).
+pub(crate) fn resolve_max_retries_with_env(
+    env_override: Option<&str>,
+    model_max_retries: Option<u32>,
+) -> u32 {
+    env_override
+        .and_then(|value| value.parse::<u32>().ok())
+        .or(model_max_retries)
+        .unwrap_or(DEFAULT_MAX_RETRIES)
+}
+
+/// Resolve max API retries: `GROK_MAX_RETRIES` env > model config > default ([`DEFAULT_MAX_RETRIES`]).
+pub fn resolve_max_retries(model_max_retries: Option<u32>) -> u32 {
+    let env_override = std::env::var("GROK_MAX_RETRIES").ok();
+    resolve_max_retries_with_env(env_override.as_deref(), model_max_retries)
+}
+
+/// Backoff for doom-loop resamples: near-immediate with a small jitter.
+/// Loops are stochastic at sampling temperature, so a fresh sample is the
+/// remedy — waiting buys nothing beyond de-syncing concurrent resamples.
+pub fn doom_loop_backoff(retry_count: u32) -> Duration {
+    use std::hash::{Hash, Hasher};
+    use std::sync::atomic::{AtomicU64, Ordering};
+
+    static JITTER_SEQ: AtomicU64 = AtomicU64::new(0);
+
+    let mut hasher = std::hash::DefaultHasher::new();
+    JITTER_SEQ.fetch_add(1, Ordering::Relaxed).hash(&mut hasher);
+    retry_count.hash(&mut hasher);
+    Duration::from_millis(hasher.finish() % 251)
+}
+
+/// Exponential backoff (2s, 4s, 8s, ..., capped 30s) with +/-20% jitter
+/// to prevent thundering-herd retry storms.
+pub fn retry_backoff_with_jitter(retry_count: u32) -> Duration {
+    use std::hash::{Hash, Hasher};
+    use std::sync::atomic::{AtomicU64, Ordering};
+
+    static JITTER_SEQ: AtomicU64 = AtomicU64::new(0);
+
+    let shift = retry_count.saturating_sub(1);
+    let base_ms = 2000u64.checked_shl(shift).unwrap_or(u64::MAX).min(30_000);
+    let jitter_range = base_ms / 5;
+    let mut hasher = std::hash::DefaultHasher::new();
+    JITTER_SEQ.fetch_add(1, Ordering::Relaxed).hash(&mut hasher);
+    std::thread::current().id().hash(&mut hasher);
+    let jitter = hasher.finish() % (jitter_range * 2 + 1);
+    Duration::from_millis(base_ms - jitter_range + jitter)
+}
+
+/// What the actor should do next given a sampling error and retry context.
+///
+/// Pure data: callers (the actor's per-request task) are responsible for
+/// performing the actual sleep, image strip, client rebuild, or emit.
+#[derive(Debug)]
+pub enum RetryDecision {
+    /// Retry with exponential backoff (transport errors, 5xx,
+    /// empty responses).
+    Retry { backoff: Duration },
+
+    /// Retry honoring the server's `Retry-After` header (429 rate
+    /// limits). `is_rate_limited` distinguishes 429s from generic
+    /// retry-with-backoff cases for telemetry.
+    RetryWithBackoff {
+        backoff: Duration,
+        is_rate_limited: bool,
+    },
+
+    /// Retry after stripping inline images from the request (413
+    /// Payload Too Large or image processing rejection).
+    RetryWithImageStrip,
+
+    /// Retry after rebuilding the HTTP client with HTTP/1.1 (transport
+    /// error, first retry only).
+    RetryWithClientRebuild { backoff: Duration },
+
+    /// Emit the error to the session and let it decide what to do
+    /// (auth refresh, encrypted-content mismatch).
+    EmitToSession(SamplingError),
+
+    /// Fatal: no further retries possible. Surface to the caller as the
+    /// final outcome of the sampling request.
+    Fatal(SamplingError),
+}
+
+/// Classify a sampling error into a [`RetryDecision`].
+///
+/// `retry_count` is the number of retries already performed (0 on first
+/// failure). `max_retries` is the total budget. `rate_limit_threshold`
+/// caps consecutive 429 retries (see [`RATE_LIMIT_RETRY_THRESHOLD`]).
+///
+/// The function is pure: it does not sleep, log, or perform I/O.
+pub fn classify_error(
+    err: &SamplingError,
+    retry_count: u32,
+    max_retries: u32,
+    rate_limit_threshold: u32,
+) -> RetryDecision {
+    // Auth and encrypted-content errors are session-owned. The sampler
+    // surfaces the raw error and lets the session refresh credentials
+    // or show a friendly message.
+    if err.is_auth_error() {
+        return RetryDecision::EmitToSession(clone_error(err));
+    }
+    if err.is_encrypted_content_error() {
+        return RetryDecision::EmitToSession(clone_error(err));
+    }
+
+    // 413 Payload Too Large: strip inline images and try once. The
+    // caller checks if there are images left after the strip; if not,
+    // upgrade to Fatal.
+    if err.is_payload_too_large() {
+        return RetryDecision::RetryWithImageStrip;
+    }
+
+    // Image processing errors (direct 400 or proxy-wrapped 500): strip
+    // images and retry, same recovery as 413.
+    if err.is_image_processing_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).
+    if err.is_context_length_error() {
+        return RetryDecision::Fatal(clone_error(err));
+    }
+
+    // Doom-loop failures: always Retry with near-immediate backoff. The
+    // recovery loop intercepts these BEFORE classification and runs its own
+    // budget (`policy.max_retries`, enforced by disarming the abort); this
+    // arm only keeps classification total so a stray doom failure through
+    // any other path can never be Fatal.
+    if matches!(err, SamplingError::DoomLoopDetected { .. }) {
+        return RetryDecision::Retry {
+            backoff: doom_loop_backoff(retry_count + 1),
+        };
+    }
+
+    // Rate-limited (429): cap retries at the rate-limit threshold to
+    // avoid burning long waits.
+    if err.is_rate_limited() {
+        let next_attempt = retry_count + 1;
+        let effective_cap = max_retries.min(rate_limit_threshold);
+        if next_attempt >= effective_cap {
+            return RetryDecision::Fatal(clone_error(err));
+        }
+        let backoff = err
+            .retry_after()
+            .map(Duration::from_secs)
+            .unwrap_or_else(|| retry_backoff_with_jitter(next_attempt));
+        return RetryDecision::RetryWithBackoff {
+            backoff,
+            is_rate_limited: true,
+        };
+    }
+
+    // Generic retryable transport / 5xx errors. First retry rebuilds
+    // the HTTP client with HTTP/1.1 to escape poisoned HTTP/2 pools;
+    // later retries just back off.
+    if err.is_retryable() {
+        let next_attempt = retry_count + 1;
+        if next_attempt >= max_retries {
+            return RetryDecision::Fatal(clone_error(err));
+        }
+        let backoff = err
+            .retry_after()
+            .map(Duration::from_secs)
+            .unwrap_or_else(|| retry_backoff_with_jitter(next_attempt));
+        if next_attempt == 1 {
+            return RetryDecision::RetryWithClientRebuild { backoff };
+        }
+        return RetryDecision::Retry { backoff };
+    }
+
+    // Everything else is fatal.
+    RetryDecision::Fatal(clone_error(err))
+}
+
+/// Build a human-readable, telemetry-friendly description of a sampling
+/// error.
+///
+/// `retry_count`, when present, is rendered as a "Request failed after
+/// N retries." prefix. The function is pure string formatting: no
+/// logging, no I/O, no allocation beyond the produced `String`.
+pub fn format_sampling_error(err: &SamplingError, retry_count: Option<u32>) -> String {
+    let retry_prefix = match retry_count {
+        Some(count) => format!("Request failed after {} retries. ", count),
+        None => String::new(),
+    };
+
+    match err {
+        SamplingError::Auth(msg) => {
+            format!(
+                "{}Authentication failed: {}. Please check your API key configuration.",
+                retry_prefix, msg
+            )
+        }
+        SamplingError::InvalidConfiguration(msg) => {
+            format!(
+                "{}Invalid configuration: {}. Please check your model settings.",
+                retry_prefix, msg
+            )
+        }
+
+        SamplingError::Http(e) => {
+            let mut details = Vec::new();
+            if e.is_timeout() {
+                details.push("timeout".to_string());
+            }
+            if e.is_connect() {
+                details.push("connection failed".to_string());
+            }
+            if let Some(status) = e.status() {
+                details.push(format!("status {}", status));
+            }
+            if let Some(url) = e.url() {
+                details.push(format!("url: {}", url));
+            }
+            let detail_str = if details.is_empty() {
+                e.to_string()
+            } else {
+                format!("{} ({})", e, details.join(", "))
+            };
+            format!(
+                "{}HTTP request failed: {}. This may be a network issue or the API endpoint may be unavailable.",
+                retry_prefix, detail_str
+            )
+        }
+        SamplingError::Serialization(e) => {
+            format!(
+                "{}Failed to parse API response at line {} column {}: {}. This indicates an unexpected response format from the server.",
+                retry_prefix,
+                e.line(),
+                e.column(),
+                e
+            )
+        }
+        SamplingError::Api {
+            status, message, ..
+        } => {
+            let status_hint = match status.as_u16() {
+                400 => " (bad request - check your input)",
+                401 | 403 => " (authentication issue - check your API key)",
+                404 => " (endpoint not found - check model configuration)",
+                413 => " (request too large - try /compact or start new session)",
+                429 => " (rate limited - please wait and retry)",
+                500 => " (server internal error)",
+                #[allow(clippy::manual_range_patterns)]
+                502 | 503 | 504 => " (server unavailable - please retry)",
+                _ => "",
+            };
+            format!(
+                "{}API error (HTTP {}{}): {}",
+                retry_prefix,
+                status.as_u16(),
+                status_hint,
+                message
+            )
+        }
+        SamplingError::EventStreamError(msg) => {
+            format!(
+                "{}Event stream error: {}. The connection to the server was interrupted.",
+                retry_prefix, msg
+            )
+        }
+        SamplingError::StreamError {
+            error_type,
+            message,
+        } => {
+            format!(
+                "{}Server stream error ({}): {}. The server encountered an error while streaming the response.",
+                retry_prefix, error_type, message
+            )
+        }
+        SamplingError::IdleTimeout { elapsed_secs } => {
+            format!(
+                "{}Model stopped responding after {}s. The model may be overloaded or stuck. Try again or use a different model.",
+                retry_prefix, elapsed_secs
+            )
+        }
+        SamplingError::EmptyResponse { context } => {
+            format!(
+                "{}Empty response from model ({}): model={}, had_reasoning={}, finish_reason={}, completion_tokens={}",
+                retry_prefix,
+                context.reason,
+                context.model,
+                context.had_reasoning,
+                context.finish_reason_str(),
+                context.completion_tokens.unwrap_or(0),
+            )
+        }
+        SamplingError::MaxTokensTruncation => {
+            format!("{}Response truncated by max_tokens.", retry_prefix)
+        }
+        SamplingError::DoomLoopDetected { triggers, .. } => {
+            format!(
+                "{}Server detected a reasoning loop ({}); resampling the response.",
+                retry_prefix,
+                triggers.join(", ")
+            )
+        }
+    }
+}
+
+/// Reconstruct an owned [`SamplingError`] from a borrowed one.
+///
+/// `SamplingError` does not implement `Clone` because its `Http` and
+/// `Serialization` variants wrap non-`Clone` types. The retry loop
+/// only borrows the error during classification, then needs to surface
+/// it; this helper produces a faithful copy where possible. `Http`
+/// falls back to a structured `EventStreamError` (still retryable, like
+/// the original transport error). `Serialization` must stay
+/// `Serialization`: laundering it into `EventStreamError` would flip a
+/// fatal response-parse failure into a retryable one and burn the full
+/// retry budget re-generating a response that fails the same way.
+pub(crate) fn clone_error(err: &SamplingError) -> SamplingError {
+    match err {
+        SamplingError::Auth(msg) => SamplingError::Auth(msg.clone()),
+        SamplingError::InvalidConfiguration(msg) => SamplingError::InvalidConfiguration(msg),
+        SamplingError::Http(e) => {
+            // reqwest::Error is not Clone; preserve the rendered message
+            // as an EventStreamError (the closest retryable transport
+            // variant) so callers see an equivalent description.
+            SamplingError::EventStreamError(e.to_string())
+        }
+        SamplingError::Serialization(e) => {
+            // serde_json::Error is not Clone; its Display already carries the
+            // original line/column exactly once.
+            SamplingError::serialization_message(e)
+        }
+        SamplingError::Api {
+            status,
+            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 {
+            error_type,
+            message,
+        } => SamplingError::StreamError {
+            error_type: error_type.clone(),
+            message: message.clone(),
+        },
+        SamplingError::IdleTimeout { elapsed_secs } => SamplingError::IdleTimeout {
+            elapsed_secs: *elapsed_secs,
+        },
+        SamplingError::EmptyResponse { context } => SamplingError::EmptyResponse {
+            context: context.clone(),
+        },
+        SamplingError::MaxTokensTruncation => SamplingError::MaxTokensTruncation,
+        SamplingError::DoomLoopDetected {
+            triggers,
+            aborted_at_chunk,
+        } => SamplingError::DoomLoopDetected {
+            triggers: triggers.clone(),
+            aborted_at_chunk: *aborted_at_chunk,
+        },
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+    use reqwest::StatusCode;
+
+    fn api_err(status: StatusCode, message: &str) -> SamplingError {
+        SamplingError::Api {
+            status,
+            message: message.to_string(),
+            model_metadata: None,
+            retry_after_secs: None,
+            should_retry: None,
+        }
+    }
+
+    fn api_err_with_retry_after(status: StatusCode, retry_after: u64) -> SamplingError {
+        SamplingError::Api {
+            status,
+            message: "x".to_string(),
+            model_metadata: None,
+            retry_after_secs: Some(retry_after),
+            should_retry: None,
+        }
+    }
+
+    #[test]
+    fn resolve_max_retries_env_override_takes_precedence() {
+        assert_eq!(resolve_max_retries_with_env(Some("9"), Some(3)), 9);
+    }
+
+    #[test]
+    fn resolve_max_retries_falls_back_to_model() {
+        assert_eq!(resolve_max_retries_with_env(None, Some(7)), 7);
+    }
+
+    #[test]
+    fn resolve_max_retries_default() {
+        assert_eq!(
+            resolve_max_retries_with_env(None, None),
+            DEFAULT_MAX_RETRIES
+        );
+    }
+
+    #[test]
+    fn resolve_max_retries_invalid_env_falls_through() {
+        assert_eq!(resolve_max_retries_with_env(Some("abc"), Some(4)), 4);
+    }
+
+    #[test]
+    fn backoff_first_retry_is_around_two_seconds() {
+        let backoff = retry_backoff_with_jitter(1);
+        // Base 2000ms +/- 20% jitter (400ms range).
+        assert!(
+            backoff >= Duration::from_millis(1600) && backoff <= Duration::from_millis(2400),
+            "first retry backoff out of range: {:?}",
+            backoff
+        );
+    }
+
+    #[test]
+    fn backoff_doubles_then_caps_at_thirty_seconds() {
+        // retry_count=2: base 4s
+        let r2 = retry_backoff_with_jitter(2);
+        assert!(r2 >= Duration::from_millis(3200) && r2 <= Duration::from_millis(4800));
+
+        // retry_count=10: base would be 2^10 * 2000 = 2.048s but capped to 30s
+        let r10 = retry_backoff_with_jitter(10);
+        assert!(r10 >= Duration::from_millis(24_000) && r10 <= Duration::from_millis(36_000));
+    }
+
+    #[test]
+    fn backoff_zero_retry_count_is_well_defined() {
+        // retry_count = 0 corresponds to "before the first retry"; ensure
+        // it does not panic and stays in the lowest backoff bucket.
+        let backoff = retry_backoff_with_jitter(0);
+        assert!(backoff >= Duration::from_millis(1600) && backoff <= Duration::from_millis(2400));
+    }
+
+    #[test]
+    fn classify_auth_error_emits_to_session() {
+        let err = SamplingError::Auth("bad token".into());
+        match classify_error(&err, 0, 5, RATE_LIMIT_RETRY_THRESHOLD) {
+            RetryDecision::EmitToSession(SamplingError::Auth(_)) => {}
+            other => panic!("expected EmitToSession(Auth), got {other:?}"),
+        }
+    }
+
+    #[test]
+    fn classify_unauthorized_emits_to_session() {
+        let err = api_err(StatusCode::UNAUTHORIZED, "no");
+        match classify_error(&err, 0, 5, RATE_LIMIT_RETRY_THRESHOLD) {
+            RetryDecision::EmitToSession(SamplingError::Api { status, .. }) => {
+                assert_eq!(status, StatusCode::UNAUTHORIZED);
+            }
+            other => panic!("expected EmitToSession(Api 401), got {other:?}"),
+        }
+    }
+
+    #[test]
+    fn classify_encrypted_content_emits_to_session() {
+        let err = api_err(
+            StatusCode::BAD_REQUEST,
+            "Could not decrypt the provided encrypted_content",
+        );
+        match classify_error(&err, 0, 5, RATE_LIMIT_RETRY_THRESHOLD) {
+            RetryDecision::EmitToSession(_) => {}
+            other => panic!("expected EmitToSession, got {other:?}"),
+        }
+    }
+
+    #[test]
+    fn classify_payload_too_large_strips_images() {
+        let err = api_err(StatusCode::PAYLOAD_TOO_LARGE, "too big");
+        assert!(matches!(
+            classify_error(&err, 0, 5, RATE_LIMIT_RETRY_THRESHOLD),
+            RetryDecision::RetryWithImageStrip
+        ));
+    }
+
+    #[test]
+    fn classify_image_processing_error_400_strips_images() {
+        let err = api_err(StatusCode::BAD_REQUEST, "Could not process image");
+        assert!(matches!(
+            classify_error(&err, 0, 5, RATE_LIMIT_RETRY_THRESHOLD),
+            RetryDecision::RetryWithImageStrip
+        ));
+    }
+
+    #[test]
+    fn classify_image_processing_error_500_wrapped_strips_images() {
+        let err = api_err(
+            StatusCode::INTERNAL_SERVER_ERROR,
+            "upstream: 400 Bad Request: Could not process image",
+        );
+        assert!(matches!(
+            classify_error(&err, 0, 5, RATE_LIMIT_RETRY_THRESHOLD),
+            RetryDecision::RetryWithImageStrip
+        ));
+    }
+
+    #[test]
+    fn classify_image_processing_error_takes_priority_over_5xx_retry() {
+        // A 500 wrapping "Could not process image" is retryable by status
+        // code alone — verify the image-processing guard intercepts first.
+        let err = api_err(
+            StatusCode::INTERNAL_SERVER_ERROR,
+            "Could not process image: bad format",
+        );
+        assert!(
+            err.is_retryable(),
+            "500 is retryable without the image-processing guard"
+        );
+        assert!(matches!(
+            classify_error(&err, 0, 5, RATE_LIMIT_RETRY_THRESHOLD),
+            RetryDecision::RetryWithImageStrip
+        ));
+    }
+
+    #[test]
+    fn classify_rate_limited_uses_retry_after() {
+        let err = api_err_with_retry_after(StatusCode::TOO_MANY_REQUESTS, 7);
+        match classify_error(&err, 0, 5, RATE_LIMIT_RETRY_THRESHOLD) {
+            RetryDecision::RetryWithBackoff {
+                backoff,
+                is_rate_limited,
+            } => {
+                assert!(is_rate_limited);
+                assert_eq!(backoff, Duration::from_secs(7));
+            }
+            other => panic!("expected RetryWithBackoff, got {other:?}"),
+        }
+    }
+
+    #[test]
+    fn classify_rate_limited_capped_at_threshold() {
+        let err = api_err(StatusCode::TOO_MANY_REQUESTS, "slow");
+        // retry_count=1, threshold=2 -> next_attempt=2 >= 2 -> Fatal.
+        match classify_error(&err, 1, 5, RATE_LIMIT_RETRY_THRESHOLD) {
+            RetryDecision::Fatal(SamplingError::Api { status, .. }) => {
+                assert_eq!(status, StatusCode::TOO_MANY_REQUESTS);
+            }
+            other => panic!("expected Fatal at threshold, got {other:?}"),
+        }
+    }
+
+    #[test]
+    fn classify_5xx_first_retry_rebuilds_client() {
+        let err = api_err(StatusCode::INTERNAL_SERVER_ERROR, "boom");
+        match classify_error(&err, 0, 5, RATE_LIMIT_RETRY_THRESHOLD) {
+            RetryDecision::RetryWithClientRebuild { backoff } => {
+                assert!(backoff >= Duration::from_millis(1600));
+            }
+            other => panic!("expected RetryWithClientRebuild, got {other:?}"),
+        }
+    }
+
+    #[test]
+    fn classify_5xx_subsequent_retry_uses_plain_retry() {
+        let err = api_err(StatusCode::BAD_GATEWAY, "boom");
+        match classify_error(&err, 1, 5, RATE_LIMIT_RETRY_THRESHOLD) {
+            RetryDecision::Retry { backoff } => {
+                assert!(backoff >= Duration::from_millis(3200));
+            }
+            other => panic!("expected Retry, got {other:?}"),
+        }
+    }
+
+    #[test]
+    fn classify_5xx_exhausted_retries_is_fatal() {
+        let err = api_err(StatusCode::SERVICE_UNAVAILABLE, "boom");
+        match classify_error(&err, 4, 5, RATE_LIMIT_RETRY_THRESHOLD) {
+            RetryDecision::Fatal(SamplingError::Api { .. }) => {}
+            other => panic!("expected Fatal, got {other:?}"),
+        }
+    }
+
+    #[test]
+    fn classify_event_stream_error_is_retryable() {
+        let err = SamplingError::EventStreamError("connection reset".into());
+        match classify_error(&err, 0, 5, RATE_LIMIT_RETRY_THRESHOLD) {
+            RetryDecision::RetryWithClientRebuild { .. } => {}
+            other => panic!("expected RetryWithClientRebuild, got {other:?}"),
+        }
+    }
+
+    #[test]
+    fn classify_stream_error_is_retryable() {
+        let err = SamplingError::StreamError {
+            error_type: "transient".into(),
+            message: "x".into(),
+        };
+        match classify_error(&err, 0, 5, RATE_LIMIT_RETRY_THRESHOLD) {
+            RetryDecision::RetryWithClientRebuild { .. } => {}
+            other => panic!("expected RetryWithClientRebuild for StreamError, got {other:?}"),
+        }
+    }
+
+    #[test]
+    fn classify_idle_timeout_is_fatal() {
+        let err = SamplingError::IdleTimeout { elapsed_secs: 300 };
+        match classify_error(&err, 0, 5, RATE_LIMIT_RETRY_THRESHOLD) {
+            RetryDecision::Fatal(SamplingError::IdleTimeout { elapsed_secs: 300 }) => {}
+            other => panic!("expected Fatal(IdleTimeout), got {other:?}"),
+        }
+    }
+
+    #[test]
+    fn classify_invalid_config_is_fatal() {
+        let err = SamplingError::InvalidConfiguration("missing model");
+        assert!(matches!(
+            classify_error(&err, 0, 5, RATE_LIMIT_RETRY_THRESHOLD),
+            RetryDecision::Fatal(SamplingError::InvalidConfiguration(_))
+        ));
+    }
+
+    #[test]
+    fn classify_api_400_non_encrypted_is_fatal() {
+        let err = api_err(StatusCode::BAD_REQUEST, "Invalid model parameter");
+        assert!(matches!(
+            classify_error(&err, 0, 5, RATE_LIMIT_RETRY_THRESHOLD),
+            RetryDecision::Fatal(_)
+        ));
+    }
+
+    fn serialization_err() -> SamplingError {
+        SamplingError::Serialization(serde_json::from_str::<i32>("not a number").unwrap_err())
+    }
+
+    /// Regression: `clone_error` used to launder `Serialization` into the
+    /// retryable `EventStreamError`, turning a deterministic parse failure
+    /// into a full-budget retry storm.
+    #[test]
+    fn clone_error_preserves_serialization_and_non_retryability() {
+        let cloned = clone_error(&serialization_err());
+        assert!(
+            matches!(cloned, SamplingError::Serialization(_)),
+            "expected Serialization, got {cloned:?}"
+        );
+        assert!(!cloned.is_retryable());
+        assert!(
+            cloned.to_string().contains("line 1 column"),
+            "original position text must survive the clone: {cloned}"
+        );
+    }
+
+    #[test]
+    fn classify_serialization_is_fatal_on_first_attempt() {
+        match classify_error(&serialization_err(), 0, 15, RATE_LIMIT_RETRY_THRESHOLD) {
+            RetryDecision::Fatal(SamplingError::Serialization(_)) => {}
+            other => panic!("expected Fatal(Serialization) on attempt 1, got {other:?}"),
+        }
+    }
+
+    #[test]
+    fn format_includes_retry_prefix_when_count_present() {
+        let err = SamplingError::Auth("bad".into());
+        let s = format_sampling_error(&err, Some(3));
+        assert!(s.starts_with("Request failed after 3 retries."));
+    }
+
+    #[test]
+    fn format_omits_retry_prefix_when_count_absent() {
+        let err = SamplingError::Auth("bad".into());
+        let s = format_sampling_error(&err, None);
+        assert!(!s.starts_with("Request failed after"));
+        assert!(s.starts_with("Authentication failed:"));
+    }
+
+    #[test]
+    fn format_includes_status_hint_for_known_codes() {
+        let err = api_err(StatusCode::PAYLOAD_TOO_LARGE, "big");
+        let s = format_sampling_error(&err, None);
+        assert!(s.contains("HTTP 413"));
+        assert!(s.contains("request too large"));
+    }
+
+    #[test]
+    fn format_idle_timeout_includes_elapsed_secs() {
+        let err = SamplingError::IdleTimeout { elapsed_secs: 240 };
+        let s = format_sampling_error(&err, None);
+        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.
+        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),
+            RetryDecision::Fatal(_)
+        ));
+    }
+
+    #[test]
+    fn should_retry_true_falls_through_to_existing_logic() {
+        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),
+            RetryDecision::RetryWithClientRebuild { .. }
+        ));
+    }
+
+    #[test]
+    fn classify_doom_loop_detected_is_retry_with_immediate_backoff() {
+        let err = SamplingError::DoomLoopDetected {
+            triggers: vec!["tail_repetition:8@thinking".into()],
+            aborted_at_chunk: None,
+        };
+        // Whatever the counters say, classification is Retry — the recovery
+        // loop owns the budget by disarming the abort when it is spent.
+        for retry_count in [0, 5, 99] {
+            match classify_error(&err, retry_count, 2, RATE_LIMIT_RETRY_THRESHOLD) {
+                RetryDecision::Retry { backoff } => {
+                    assert!(backoff <= Duration::from_millis(250), "near-immediate");
+                }
+                other => panic!("expected Retry, got {other:?}"),
+            }
+        }
+    }
+
+    #[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/xai-grok-sampler/src/sampling_log.rs

@@ -0,0 +1,37 @@
+//! Sampling log — emits `tracing` events with `target: "sampling_log"`.
+//! A dedicated layer in `xai-grok-telemetry` routes these to
+//! `~/.grok/logs/sampling.jsonl`. Enable with `--log-sampling`.
+
+use crate::types::RequestId;
+
+pub const TARGET: &str = "sampling_log";
+
+#[derive(Debug, Clone)]
+pub struct AuthInfo {
+    pub auth_type: &'static str,
+    pub auth_prefix: Option<String>,
+}
+
+pub fn request_span(
+    request_id: &RequestId,
+    model: &str,
+    api_backend: &str,
+    base_url: &str,
+    auth: &AuthInfo,
+) -> tracing::Span {
+    tracing::info_span!(
+        target: TARGET,
+        "sampling_request",
+        request_id = %request_id,
+        model = model,
+        api_backend = api_backend,
+        base_url = base_url,
+        auth_type = auth.auth_type,
+        auth_prefix = auth.auth_prefix.as_deref().unwrap_or(""),
+        // Recorded from `SamplerConfig` / response usage as the request
+        // progresses; `field::Empty` lets callers `record()` them later.
+        reasoning_effort = tracing::field::Empty,
+        output_tokens = tracing::field::Empty,
+        reasoning_tokens = tracing::field::Empty,
+    )
+}

crates/codegen/xai-grok-sampler/src/shared_http.rs

@@ -0,0 +1,156 @@
+//! Process-wide shared `reqwest::Client`s for sampling requests.
+//!
+//! Sharing one client across all `SamplingClient` instances is safe because
+//! the builders below take no config-derived input: auth, extra headers, base
+//! URL, and User-Agent are all applied per-request in `SamplingClient::post`.
+//! Stale-connection exposure is bounded by HTTP/2 keepalive pings (15s
+//! interval, 5s timeout, while idle), the 90s idle-pool eviction, and the
+//! first-retry HTTP/1.1 rebuild escape hatch (that client never pools, so
+//! every use opens a fresh connection).
+//!
+//! Wire-level behavior (connection reuse, header isolation, pool-less http1
+//! fallback, kill switch) is pinned by the `shared_http_wire` and
+//! `shared_http_kill_switch` integration binaries, which own their process
+//! environment.
+
+use std::sync::OnceLock;
+use std::time::Duration;
+
+static SHARED_H2: OnceLock<reqwest::Client> = OnceLock::new();
+static SHARED_HTTP1: OnceLock<reqwest::Client> = OnceLock::new();
+
+/// Kill switch: `GROK_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.
+fn sharing_disabled() -> bool {
+    static DISABLED: OnceLock<bool> = OnceLock::new();
+    *DISABLED.get_or_init(|| {
+        let disabled = match std::env::var("GROK_SAMPLER_SHARED_CLIENT") {
+            Ok(v) => v == "0" || v.eq_ignore_ascii_case("false"),
+            Err(_) => false,
+        };
+        if disabled {
+            tracing::info!("sampler HTTP client sharing disabled via GROK_SAMPLER_SHARED_CLIENT");
+        }
+        disabled
+    })
+}
+
+/// Clone the shared client out of `cell`, building it on first use. Build
+/// failures are not cached: on `Err` the cell stays empty and the next call
+/// retries. A racing loser's freshly built client is simply dropped.
+fn shared(
+    cell: &OnceLock<reqwest::Client>,
+    build: fn() -> Result<reqwest::Client, reqwest::Error>,
+    disabled: bool,
+) -> Result<reqwest::Client, reqwest::Error> {
+    if disabled {
+        return build();
+    }
+    if let Some(client) = cell.get() {
+        return Ok(client.clone());
+    }
+    let built = build()?;
+    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.
+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("GROK_POOL_MAX_IDLE")
+        .ok()
+        .and_then(|v| v.parse().ok())
+        .unwrap_or(2);
+    let pool_idle_timeout_secs: u64 = std::env::var("GROK_POOL_IDLE_TIMEOUT_SECS")
+        .ok()
+        .and_then(|v| v.parse().ok())
+        .unwrap_or(90);
+    let connect_timeout_secs: u64 = std::env::var("GROK_CONNECT_TIMEOUT_SECS")
+        .ok()
+        .and_then(|v| v.parse().ok())
+        .unwrap_or(10);
+
+    reqwest::Client::builder()
+        .pool_max_idle_per_host(pool_max_idle)
+        .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.
+fn build_http_client_http1() -> Result<reqwest::Client, reqwest::Error> {
+    let connect_timeout_secs: u64 = std::env::var("GROK_CONNECT_TIMEOUT_SECS")
+        .ok()
+        .and_then(|v| v.parse().ok())
+        .unwrap_or(10);
+
+    reqwest::Client::builder()
+        .pool_max_idle_per_host(0)
+        .pool_idle_timeout(Duration::from_secs(0))
+        .connect_timeout(Duration::from_secs(connect_timeout_secs))
+        .tcp_nodelay(true)
+        .http1_only()
+        .build()
+}
+
+#[cfg(test)]
+mod tests {
+    use std::sync::OnceLock;
+    use std::sync::atomic::{AtomicUsize, Ordering};
+
+    use super::shared;
+
+    static BUILD_CALLS: AtomicUsize = AtomicUsize::new(0);
+
+    /// Fails on the first call (a real `reqwest::Error`, no I/O), then builds.
+    fn flaky_build() -> Result<reqwest::Client, reqwest::Error> {
+        if BUILD_CALLS.fetch_add(1, Ordering::SeqCst) == 0 {
+            return Err(reqwest::Proxy::all("not a proxy url").unwrap_err());
+        }
+        reqwest::Client::builder().build()
+    }
+
+    #[test]
+    fn shared_does_not_cache_build_failures() {
+        static CELL: OnceLock<reqwest::Client> = OnceLock::new();
+        assert!(shared(&CELL, flaky_build, false).is_err());
+        assert!(CELL.get().is_none(), "failure must leave the cell empty");
+        assert!(shared(&CELL, flaky_build, false).is_ok());
+        assert!(CELL.get().is_some(), "success must populate the cell");
+        assert!(shared(&CELL, flaky_build, false).is_ok());
+        assert_eq!(
+            BUILD_CALLS.load(Ordering::SeqCst),
+            2,
+            "third call must reuse the cached client, not rebuild"
+        );
+    }
+
+    #[test]
+    fn shared_disabled_bypasses_cell() {
+        static CELL: OnceLock<reqwest::Client> = OnceLock::new();
+        assert!(shared(&CELL, || reqwest::Client::builder().build(), true).is_ok());
+        assert!(
+            CELL.get().is_none(),
+            "disabled mode must never touch the cell"
+        );
+    }
+}

crates/codegen/xai-grok-sampler/src/stream/chat_completions.rs

@@ -0,0 +1,774 @@
+//! Layer-2 stream transform for the Chat Completions API.
+//!
+//! Consumes a raw `ChatCompletionChunk` stream and produces
+//! [`SamplingEvent`]s. Pure: no I/O, no shell coupling.
+
+use std::collections::BTreeMap;
+use std::time::{Duration, Instant};
+
+use futures_util::StreamExt;
+use futures_util::stream::{BoxStream, Stream};
+
+use xai_grok_sampling_types::{
+    AssistantItem, ChatCompletionChunk, ConversationItem, ConversationResponse,
+    ResponseModelMetadata, SamplingError, StopReason, TokenUsage, ToolCall,
+};
+
+use crate::events::{SamplingChannel, SamplingErrorInfo, SamplingEvent};
+use crate::metrics::InferenceLatencyStats;
+use crate::types::RequestId;
+
+/// Transform a raw Chat Completions chunk stream into a stream of
+/// [`SamplingEvent`]s.
+///
+/// The output stream emits exactly one terminal event per request:
+/// [`SamplingEvent::Completed`] on normal stream end, or
+/// [`SamplingEvent::Failed`] on error / idle timeout. Callers must not
+/// consume past the terminal event (the implementation `return`s after
+/// yielding it).
+///
+/// `idle_timeout` covers two cases:
+/// 1. The transport stops yielding chunks at all (`tokio::time::timeout`).
+/// 2. The transport keeps yielding empty / keepalive chunks but no
+///    meaningful content (separate `last_content_chunk_at` timer).
+///
+/// Both produce `SamplingEvent::Failed { kind: IdleTimeout }`.
+pub fn stream_chat_completions<'a>(
+    raw_stream: BoxStream<'a, Result<ChatCompletionChunk, SamplingError>>,
+    model_metadata: Option<ResponseModelMetadata>,
+    request_id: RequestId,
+    idle_timeout: Duration,
+) -> impl Stream<Item = SamplingEvent> + Send + 'a {
+    async_stream::stream! {
+        let stream_start = Instant::now();
+        let mut chunk_timestamps: Vec<Instant> = Vec::new();
+
+        // Emit StreamStarted before reading any chunks so subscribers
+        // can record TTFB / TTLB baselines.
+        yield SamplingEvent::StreamStarted {
+            request_id: request_id.clone(),
+            timestamp_ms: chrono::Utc::now().timestamp_millis(),
+        };
+
+        if let Some(metadata) = model_metadata {
+            yield SamplingEvent::ModelMetadata {
+                request_id: request_id.clone(),
+                metadata,
+            };
+        }
+
+        // Per-response accumulators
+        let mut first_chunk_seen = false;
+        let mut first_choice_seen = false;
+        let mut first_token_emitted = false;
+        let mut model: String = String::new();
+        let mut model_fingerprint: Option<String> = None;
+        let mut usage: Option<TokenUsage> = None;
+        let mut cost_usd_ticks: Option<i64> = None;
+        let mut finish_reason: Option<StopReason> = None;
+
+        let mut content_acc = String::new();
+        let mut reasoning_acc = String::new();
+        // Tool call deltas keyed by positional index. Each entry is
+        // (id, name, arguments_buffer); the first chunk for an index
+        // carries id+name and starts the arguments buffer, subsequent
+        // chunks append to arguments only.
+        let mut tool_call_acc: BTreeMap<u32, (String, String, String)> = BTreeMap::new();
+
+        // Index counter spanning text + reasoning chunks (matches the
+        // shell's chunk_index used for notification correlation).
+        let mut chunk_index: u64 = 0;
+        // Separate counter for AgentMessageChunk (text-only) emissions;
+        // mirrored onto ConversationResponse.message_chunks_emitted so
+        // downstream can detect lost-streaming-events scenarios.
+        let mut message_chunk_count: u64 = 0;
+
+        // Content-aware idle timer: the outer
+        // `tokio::time::timeout(idle_timeout, stream.next())` already
+        // catches "transport stops yielding chunks". This second timer
+        // catches the more subtle case where the model keeps emitting
+        // keepalive / empty-delta SSE events that satisfy the outer
+        // timer but make no real progress -- some inference engines
+        // do exactly that.
+        let mut last_content_chunk_at = Instant::now();
+
+        let mut stream = raw_stream;
+        loop {
+            let next = match tokio::time::timeout(idle_timeout, stream.next()).await {
+                Ok(Some(next)) => next,
+                Ok(None) => break, // stream ended normally
+                Err(_elapsed) => {
+                    let err = SamplingError::IdleTimeout {
+                        elapsed_secs: idle_timeout.as_secs(),
+                    };
+                    yield SamplingEvent::Failed {
+                        request_id: request_id.clone(),
+                        error: SamplingErrorInfo::from(&err),
+                    };
+                    return;
+                }
+            };
+            let chunk = match next {
+                Ok(chunk) => chunk,
+                Err(err) => {
+                    yield SamplingEvent::Failed {
+                        request_id: request_id.clone(),
+                        error: SamplingErrorInfo::from(&err),
+                    };
+                    return;
+                }
+            };
+
+            if !first_chunk_seen {
+                model = chunk.model.clone();
+                model_fingerprint = chunk
+                    .system_fingerprint
+                    .clone()
+                    .filter(|s| !s.is_empty());
+                first_chunk_seen = true;
+            }
+
+            if let Some(u) = chunk.usage.clone() {
+                // Wire cost is cumulative for the response, so last-write-wins.
+                // Never clobber a known cost with missing/unreported.
+                let chunk_cost = xai_grok_sampling_types::reported_cost_ticks(u.cost_in_usd_ticks);
+                cost_usd_ticks = match (cost_usd_ticks, chunk_cost) {
+                    (_, Some(n)) => Some(n),
+                    (prev, None) => prev,
+                };
+                usage = Some(u.into());
+            }
+
+            // Track whether this chunk carried meaningful content.
+            // Set inside the choices loop and checked at the end.
+            let mut chunk_has_content = false;
+
+            for choice in chunk.choices.into_iter() {
+                first_choice_seen = true;
+                if let Some(fr) = choice.finish_reason {
+                    finish_reason = Some(fr.into());
+                    chunk_has_content = true;
+                }
+
+                let delta = choice.delta;
+
+                if let Some(text) = delta.content
+                    && !text.is_empty()
+                {
+                    if !first_token_emitted {
+                        first_token_emitted = true;
+                        yield SamplingEvent::FirstToken {
+                            request_id: request_id.clone(),
+                        };
+                    }
+                    chunk_has_content = true;
+                    chunk_timestamps.push(Instant::now());
+                    chunk_index += 1;
+                    message_chunk_count += 1;
+                    content_acc.push_str(&text);
+                    yield SamplingEvent::ChannelToken {
+                        request_id: request_id.clone(),
+                        channel: SamplingChannel::Text,
+                        text,
+                        chunk_index,
+                    };
+                }
+
+                if let Some(thought) = delta.reasoning_content
+                    && !thought.is_empty()
+                {
+                    if !first_token_emitted {
+                        first_token_emitted = true;
+                        yield SamplingEvent::FirstToken {
+                            request_id: request_id.clone(),
+                        };
+                    }
+                    chunk_has_content = true;
+                    chunk_index += 1;
+                    reasoning_acc.push_str(&thought);
+                    yield SamplingEvent::ChannelToken {
+                        request_id: request_id.clone(),
+                        channel: SamplingChannel::Reasoning,
+                        text: thought,
+                        chunk_index,
+                    };
+                }
+
+                for tc_delta in delta.tool_calls.into_iter() {
+                    chunk_has_content = true;
+
+                    let entry = tool_call_acc
+                        .entry(tc_delta.index)
+                        .or_insert_with(|| (String::new(), String::new(), String::new()));
+
+                    let mut id_for_event: Option<String> = None;
+                    let mut name_for_event: Option<String> = None;
+                    let mut args_for_event: Option<String> = None;
+
+                    if let Some(id) = tc_delta.id {
+                        entry.0 = id.clone();
+                        id_for_event = Some(id);
+                    }
+                    if let Some(func) = tc_delta.function {
+                        if let Some(name) = func.name {
+                            entry.1 = name.clone();
+                            name_for_event = Some(name);
+                        }
+                        if let Some(args) = func.arguments {
+                            entry.2.push_str(&args);
+                            args_for_event = Some(args);
+                        }
+                    }
+
+                    yield SamplingEvent::ToolCallDelta {
+                        request_id: request_id.clone(),
+                        tool_index: tc_delta.index,
+                        id: id_for_event,
+                        name: name_for_event,
+                        arguments_delta: args_for_event,
+                    };
+                }
+            }
+
+            if chunk_has_content {
+                last_content_chunk_at = Instant::now();
+            } else if last_content_chunk_at.elapsed() > idle_timeout {
+                let err = SamplingError::IdleTimeout {
+                    elapsed_secs: idle_timeout.as_secs(),
+                };
+                yield SamplingEvent::Failed {
+                    request_id: request_id.clone(),
+                    error: SamplingErrorInfo::from(&err),
+                };
+                return;
+            }
+        }
+
+        // ── 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),
+            })
+            .collect();
+
+        // Honor tool calls by overriding the stop reason if the model
+        // forgot to set it (mirrors the shell's behavior).
+        if !tool_calls.is_empty() {
+            finish_reason = Some(StopReason::ToolCalls);
+        }
+
+        // Build the trailing Assistant + any reasoning sibling.
+        let mut items: Vec<ConversationItem> = Vec::new();
+        if first_choice_seen {
+            if !reasoning_acc.is_empty() {
+                items.push(ConversationItem::Reasoning(
+                    xai_grok_sampling_types::synthesized_reasoning_item(reasoning_acc),
+                ));
+            }
+            items.push(ConversationItem::Assistant(AssistantItem {
+                content: std::sync::Arc::<str>::from(content_acc),
+                tool_calls,
+                model_id: Some(model),
+                model_fingerprint,
+                // Chat Completions does not echo the applied reasoning effort.
+                reasoning_effort: None,
+            }));
+        } else {
+            items.push(ConversationItem::assistant(""));
+        }
+
+        let stream_end = Instant::now();
+        let metrics =
+            InferenceLatencyStats::from_timestamps(stream_start, &chunk_timestamps, stream_end);
+
+        let response = ConversationResponse {
+            items,
+            stop_reason: finish_reason,
+            usage,
+            cost_usd_ticks,
+            message_chunks_emitted: message_chunk_count,
+            doom_loop_signals: Vec::new(),
+            stop_message: None,
+        };
+
+        yield SamplingEvent::Completed {
+            request_id: request_id.clone(),
+            response: Box::new(response),
+            metrics,
+        };
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+    use futures_util::stream;
+    use std::pin::pin;
+    use xai_grok_sampling_types::{
+        ChatChunkChoice, ChatChunkDelta, FinishReason, Role, ToolCallDelta as ChunkToolCallDelta,
+        ToolCallFunctionDelta, Usage, rs,
+    };
+
+    fn rid() -> RequestId {
+        RequestId::from("test-req")
+    }
+
+    fn make_chunk(deltas: Vec<ChatChunkDelta>) -> ChatCompletionChunk {
+        ChatCompletionChunk {
+            id: "chunk-1".into(),
+            object: "chat.completion.chunk".into(),
+            created: 0,
+            model: "test-model".into(),
+            choices: deltas
+                .into_iter()
+                .enumerate()
+                .map(|(i, delta)| ChatChunkChoice {
+                    index: i as u32,
+                    delta,
+                    finish_reason: None,
+                })
+                .collect(),
+            usage: None,
+            system_fingerprint: None,
+        }
+    }
+
+    fn text_chunk(text: &str) -> ChatCompletionChunk {
+        make_chunk(vec![ChatChunkDelta {
+            role: Some(Role::Assistant),
+            content: Some(text.to_string()),
+            reasoning_content: None,
+            tool_calls: vec![],
+            tool_call_id: None,
+        }])
+    }
+
+    fn final_chunk(reason: FinishReason) -> ChatCompletionChunk {
+        let mut chunk = make_chunk(vec![ChatChunkDelta::default()]);
+        chunk.choices[0].finish_reason = Some(reason);
+        chunk
+    }
+
+    async fn collect(s: impl Stream<Item = SamplingEvent>) -> Vec<SamplingEvent> {
+        let mut out = Vec::new();
+        let mut s = pin!(s);
+        while let Some(ev) = s.next().await {
+            out.push(ev);
+        }
+        out
+    }
+
+    #[tokio::test]
+    async fn empty_stream_yields_started_then_completed() {
+        let raw = stream::iter(Vec::<Result<ChatCompletionChunk, SamplingError>>::new()).boxed();
+        let events = collect(stream_chat_completions(
+            raw,
+            None,
+            rid(),
+            Duration::from_secs(60),
+        ))
+        .await;
+
+        assert_eq!(events.len(), 2);
+        assert!(matches!(events[0], SamplingEvent::StreamStarted { .. }));
+        match &events[1] {
+            SamplingEvent::Completed { response, .. } => {
+                assert!(response.is_empty());
+            }
+            other => panic!("expected Completed, got {other:?}"),
+        }
+    }
+
+    #[tokio::test]
+    async fn text_only_stream_emits_first_token_then_channel_tokens_then_completed() {
+        let chunks: Vec<Result<ChatCompletionChunk, SamplingError>> = vec![
+            Ok(text_chunk("Hello, ")),
+            Ok(text_chunk("world!")),
+            Ok(final_chunk(FinishReason::Stop)),
+        ];
+        let raw = stream::iter(chunks).boxed();
+        let events = collect(stream_chat_completions(
+            raw,
+            None,
+            rid(),
+            Duration::from_secs(60),
+        ))
+        .await;
+
+        // Expected sequence: StreamStarted, FirstToken, ChannelToken(Text)
+        // x 2, Completed.
+        assert!(matches!(events[0], SamplingEvent::StreamStarted { .. }));
+        assert!(matches!(events[1], SamplingEvent::FirstToken { .. }));
+
+        let text_tokens: Vec<&str> = events
+            .iter()
+            .filter_map(|e| match e {
+                SamplingEvent::ChannelToken {
+                    channel: SamplingChannel::Text,
+                    text,
+                    ..
+                } => Some(text.as_str()),
+                _ => None,
+            })
+            .collect();
+        assert_eq!(text_tokens, vec!["Hello, ", "world!"]);
+
+        match events.last().unwrap() {
+            SamplingEvent::Completed { response, .. } => {
+                let a = response.assistant().expect("assistant item present");
+                assert_eq!(a.content.as_ref(), "Hello, world!");
+                assert_eq!(response.stop_reason, Some(StopReason::Stop));
+                assert_eq!(response.message_chunks_emitted, 2);
+            }
+            other => panic!("expected Completed, got {other:?}"),
+        }
+    }
+
+    #[tokio::test]
+    async fn reasoning_chunk_emits_reasoning_channel_and_first_token_once() {
+        let mut reasoning_chunk = make_chunk(vec![ChatChunkDelta {
+            role: Some(Role::Assistant),
+            content: None,
+            reasoning_content: Some("thinking...".into()),
+            tool_calls: vec![],
+            tool_call_id: None,
+        }]);
+        reasoning_chunk.choices[0].finish_reason = None;
+
+        let chunks: Vec<Result<ChatCompletionChunk, SamplingError>> = vec![
+            Ok(reasoning_chunk),
+            Ok(text_chunk("done")),
+            Ok(final_chunk(FinishReason::Stop)),
+        ];
+        let raw = stream::iter(chunks).boxed();
+        let events = collect(stream_chat_completions(
+            raw,
+            None,
+            rid(),
+            Duration::from_secs(60),
+        ))
+        .await;
+
+        // FirstToken should appear exactly once.
+        let first_token_count = events
+            .iter()
+            .filter(|e| matches!(e, SamplingEvent::FirstToken { .. }))
+            .count();
+        assert_eq!(first_token_count, 1);
+
+        let mut saw_reasoning = false;
+        let mut saw_text = false;
+        for e in &events {
+            if let SamplingEvent::ChannelToken { channel, text, .. } = e {
+                match channel {
+                    SamplingChannel::Reasoning => {
+                        assert_eq!(text, "thinking...");
+                        saw_reasoning = true;
+                    }
+                    SamplingChannel::Text => {
+                        assert_eq!(text, "done");
+                        saw_text = true;
+                    }
+                }
+            }
+        }
+        assert!(saw_reasoning && saw_text);
+
+        match events.last().unwrap() {
+            SamplingEvent::Completed { response, .. } => {
+                let r = response
+                    .reasoning_items()
+                    .next()
+                    .expect("reasoning sibling preserved");
+                let rs::SummaryPart::SummaryText(t) = &r.summary[0];
+                assert_eq!(t.text, "thinking...");
+            }
+            other => panic!("expected Completed, got {other:?}"),
+        }
+    }
+
+    #[tokio::test]
+    async fn tool_call_stream_emits_deltas_and_assembles_final_call() {
+        // First chunk has id + name + part of arguments.
+        let chunk1 = make_chunk(vec![ChatChunkDelta {
+            role: None,
+            content: None,
+            reasoning_content: None,
+            tool_calls: vec![ChunkToolCallDelta {
+                index: 0,
+                id: Some("call_abc".into()),
+                kind: Some("function".into()),
+                function: Some(ToolCallFunctionDelta {
+                    name: Some("do_thing".into()),
+                    arguments: Some("{\"x\":".into()),
+                }),
+            }],
+            tool_call_id: None,
+        }]);
+        // Second chunk has only argument fragment.
+        let chunk2 = make_chunk(vec![ChatChunkDelta {
+            role: None,
+            content: None,
+            reasoning_content: None,
+            tool_calls: vec![ChunkToolCallDelta {
+                index: 0,
+                id: None,
+                kind: None,
+                function: Some(ToolCallFunctionDelta {
+                    name: None,
+                    arguments: Some("1}".into()),
+                }),
+            }],
+            tool_call_id: None,
+        }]);
+
+        let raw = stream::iter::<Vec<Result<ChatCompletionChunk, SamplingError>>>(vec![
+            Ok(chunk1),
+            Ok(chunk2),
+        ])
+        .boxed();
+        let events = collect(stream_chat_completions(
+            raw,
+            None,
+            rid(),
+            Duration::from_secs(60),
+        ))
+        .await;
+
+        let deltas: Vec<_> = events
+            .iter()
+            .filter_map(|e| match e {
+                SamplingEvent::ToolCallDelta {
+                    tool_index,
+                    id,
+                    name,
+                    arguments_delta,
+                    ..
+                } => Some((
+                    *tool_index,
+                    id.clone(),
+                    name.clone(),
+                    arguments_delta.clone(),
+                )),
+                _ => None,
+            })
+            .collect();
+
+        assert_eq!(deltas.len(), 2);
+        assert_eq!(deltas[0].0, 0);
+        assert_eq!(deltas[0].1.as_deref(), Some("call_abc"));
+        assert_eq!(deltas[0].2.as_deref(), Some("do_thing"));
+        assert_eq!(deltas[0].3.as_deref(), Some("{\"x\":"));
+        assert_eq!(deltas[1].1, None);
+        assert_eq!(deltas[1].2, None);
+        assert_eq!(deltas[1].3.as_deref(), Some("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_abc");
+                assert_eq!(calls[0].name, "do_thing");
+                assert_eq!(calls[0].arguments.as_ref(), "{\"x\":1}");
+                // Tool calls force ToolCalls stop reason.
+                assert_eq!(response.stop_reason, Some(StopReason::ToolCalls));
+            }
+            other => panic!("expected Completed, got {other:?}"),
+        }
+    }
+
+    #[tokio::test]
+    async fn mid_stream_error_yields_failed_no_completed() {
+        let chunks: Vec<Result<ChatCompletionChunk, SamplingError>> = vec![
+            Ok(text_chunk("hi")),
+            Err(SamplingError::EventStreamError("conn reset".into())),
+        ];
+        let raw = stream::iter(chunks).boxed();
+        let events = collect(stream_chat_completions(
+            raw,
+            None,
+            rid(),
+            Duration::from_secs(60),
+        ))
+        .await;
+
+        assert!(
+            events
+                .iter()
+                .any(|e| matches!(e, SamplingEvent::Failed { .. }))
+        );
+        assert!(
+            !events
+                .iter()
+                .any(|e| matches!(e, SamplingEvent::Completed { .. }))
+        );
+    }
+
+    #[tokio::test(start_paused = true)]
+    async fn idle_timeout_when_stream_stalls() {
+        // A stream that yields one chunk then hangs forever.
+        let raw = stream::iter(vec![Ok(text_chunk("hello"))])
+            .chain(stream::pending())
+            .boxed();
+        let events = collect(stream_chat_completions(
+            raw,
+            None,
+            rid(),
+            Duration::from_millis(100),
+        ))
+        .await;
+
+        // Stream should emit StreamStarted, FirstToken, ChannelToken
+        // then Failed(IdleTimeout) when the stall hits the deadline.
+        match events.last().unwrap() {
+            SamplingEvent::Failed { error, .. } => {
+                assert_eq!(error.kind, crate::events::SamplingErrorKind::IdleTimeout);
+            }
+            other => panic!("expected Failed(IdleTimeout), got {other:?}"),
+        }
+    }
+
+    #[tokio::test]
+    async fn model_metadata_yielded_after_stream_started() {
+        let raw = stream::iter(Vec::<Result<ChatCompletionChunk, SamplingError>>::new()).boxed();
+        let metadata = ResponseModelMetadata {
+            context_window: Some(8192),
+            max_completion_tokens: Some(4096),
+            models_etag: None,
+        };
+        let events = collect(stream_chat_completions(
+            raw,
+            Some(metadata.clone()),
+            rid(),
+            Duration::from_secs(60),
+        ))
+        .await;
+
+        assert!(matches!(events[0], SamplingEvent::StreamStarted { .. }));
+        match &events[1] {
+            SamplingEvent::ModelMetadata { metadata: m, .. } => {
+                assert_eq!(m.context_window, Some(8192));
+                assert_eq!(m.max_completion_tokens, Some(4096));
+            }
+            other => panic!("expected ModelMetadata second, got {other:?}"),
+        }
+    }
+
+    #[tokio::test]
+    async fn usage_is_extracted_from_chunk() {
+        let mut chunk_with_usage = make_chunk(vec![ChatChunkDelta::default()]);
+        chunk_with_usage.usage = Some(Usage {
+            prompt_tokens: 100,
+            completion_tokens: 50,
+            total_tokens: 150,
+            prompt_tokens_details: None,
+            completion_tokens_details: None,
+            cost_in_usd_ticks: None,
+        });
+
+        let chunks: Vec<Result<ChatCompletionChunk, SamplingError>> = vec![
+            Ok(text_chunk("ok")),
+            Ok(chunk_with_usage),
+            Ok(final_chunk(FinishReason::Stop)),
+        ];
+        let raw = stream::iter(chunks).boxed();
+        let events = collect(stream_chat_completions(
+            raw,
+            None,
+            rid(),
+            Duration::from_secs(60),
+        ))
+        .await;
+
+        match events.last().unwrap() {
+            SamplingEvent::Completed { response, .. } => {
+                let u = response.usage.as_ref().expect("usage extracted");
+                assert_eq!(u.prompt_tokens, 100);
+                assert_eq!(u.completion_tokens, 50);
+                assert_eq!(u.total_tokens, 150);
+            }
+            other => panic!("expected Completed, got {other:?}"),
+        }
+    }
+
+    /// Server-reported cost lands on the response; the REST mapper's `0`
+    /// backfill means "unreported" and must yield `None`.
+    #[tokio::test]
+    async fn cost_is_extracted_and_zero_is_unreported() {
+        for (wire, expected) in [(Some(78), Some(78)), (Some(0), None), (None, None)] {
+            let mut chunk_with_usage = make_chunk(vec![ChatChunkDelta::default()]);
+            chunk_with_usage.usage = Some(Usage {
+                prompt_tokens: 10,
+                completion_tokens: 5,
+                total_tokens: 15,
+                prompt_tokens_details: None,
+                completion_tokens_details: None,
+                cost_in_usd_ticks: wire,
+            });
+            let chunks: Vec<Result<ChatCompletionChunk, SamplingError>> = vec![
+                Ok(text_chunk("ok")),
+                Ok(chunk_with_usage),
+                Ok(final_chunk(FinishReason::Stop)),
+            ];
+            let raw = stream::iter(chunks).boxed();
+            let events = collect(stream_chat_completions(
+                raw,
+                None,
+                rid(),
+                Duration::from_secs(60),
+            ))
+            .await;
+            match events.last().unwrap() {
+                SamplingEvent::Completed { response, .. } => {
+                    assert_eq!(response.cost_usd_ticks, expected, "wire {wire:?}");
+                }
+                other => panic!("expected Completed, got {other:?}"),
+            }
+        }
+    }
+
+    #[tokio::test]
+    async fn later_missing_cost_does_not_clobber_earlier_ticks() {
+        let mut first = make_chunk(vec![ChatChunkDelta::default()]);
+        first.usage = Some(Usage {
+            prompt_tokens: 10,
+            completion_tokens: 5,
+            total_tokens: 15,
+            prompt_tokens_details: None,
+            completion_tokens_details: None,
+            cost_in_usd_ticks: Some(99),
+        });
+        let mut second = make_chunk(vec![ChatChunkDelta::default()]);
+        second.usage = Some(Usage {
+            prompt_tokens: 12,
+            completion_tokens: 6,
+            total_tokens: 18,
+            prompt_tokens_details: None,
+            completion_tokens_details: None,
+            cost_in_usd_ticks: Some(0),
+        });
+        let chunks: Vec<Result<ChatCompletionChunk, SamplingError>> = vec![
+            Ok(text_chunk("ok")),
+            Ok(first),
+            Ok(second),
+            Ok(final_chunk(FinishReason::Stop)),
+        ];
+        let raw = stream::iter(chunks).boxed();
+        let events = collect(stream_chat_completions(
+            raw,
+            None,
+            rid(),
+            Duration::from_secs(60),
+        ))
+        .await;
+        match events.last().unwrap() {
+            SamplingEvent::Completed { response, .. } => {
+                assert_eq!(response.cost_usd_ticks, Some(99));
+            }
+            other => panic!("expected Completed, got {other:?}"),
+        }
+    }
+}

crates/codegen/xai-grok-sampler/src/stream/collect.rs

@@ -0,0 +1,182 @@
+//! Buffered consumer for [`SamplingEvent`] streams.
+//!
+//! Drains a Layer-2 event stream into the final
+//! `(ConversationResponse, InferenceLatencyStats)` pair. Used by
+//! callers that don't need streaming UI updates (e.g., compaction,
+//! `/btw`, dream-model calls).
+
+use futures_util::StreamExt;
+use futures_util::stream::Stream;
+
+use xai_grok_sampling_types::ConversationResponse;
+
+use crate::events::{SamplingErrorInfo, SamplingErrorKind, SamplingEvent};
+use crate::metrics::InferenceLatencyStats;
+
+/// Drain a [`SamplingEvent`] stream, returning the final response.
+///
+/// Returns `Ok((response, metrics))` on the first
+/// [`SamplingEvent::Completed`] and `Err(error)` on the first
+/// [`SamplingEvent::Failed`]. Intermediate events (deltas, retries,
+/// metadata) are silently consumed -- this function is for callers
+/// that only need the final result.
+///
+/// If the stream ends without yielding either terminal event,
+/// returns an `Err` of kind [`SamplingErrorKind::Api`] indicating
+/// truncation. The Layer-2 transforms guarantee a terminal event in
+/// every successful return path, so this only fires for streams that
+/// are dropped mid-flight (e.g., the producer panicked or the
+/// underlying `tokio::spawn` was cancelled).
+pub async fn collect_response(
+    stream: impl Stream<Item = SamplingEvent>,
+) -> Result<(ConversationResponse, InferenceLatencyStats), SamplingErrorInfo> {
+    tokio::pin!(stream);
+
+    while let Some(event) = stream.next().await {
+        match event {
+            SamplingEvent::Completed {
+                response, metrics, ..
+            } => return Ok((*response, metrics)),
+            SamplingEvent::Failed { error, .. } => return Err(error),
+            // Drop intermediate events; this is a buffered collector.
+            _ => {}
+        }
+    }
+
+    Err(SamplingErrorInfo {
+        kind: SamplingErrorKind::Api,
+        status_code: None,
+        message: "stream ended without Completed or Failed".to_string(),
+        is_retryable: false,
+        retry_after_secs: None,
+        model_metadata: None,
+        empty_response_context: None,
+        doom_loop_triggers: None,
+        doom_loop_aborted_at_chunk: None,
+    })
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+    use futures_util::stream;
+    use xai_grok_sampling_types::{ConversationItem, SamplingError, StopReason};
+
+    use crate::events::SamplingChannel;
+    use crate::stream::stream_chat_completions;
+    use crate::types::RequestId;
+    use std::time::Duration;
+    use xai_grok_sampling_types::{
+        ChatChunkChoice, ChatChunkDelta, ChatCompletionChunk, FinishReason, Role,
+    };
+
+    fn rid() -> RequestId {
+        RequestId::from("collect-test")
+    }
+
+    fn text_chunk(text: &str) -> ChatCompletionChunk {
+        ChatCompletionChunk {
+            id: "chunk".into(),
+            object: "chat.completion.chunk".into(),
+            created: 0,
+            model: "test-model".into(),
+            choices: vec![ChatChunkChoice {
+                index: 0,
+                delta: ChatChunkDelta {
+                    role: Some(Role::Assistant),
+                    content: Some(text.to_string()),
+                    reasoning_content: None,
+                    tool_calls: vec![],
+                    tool_call_id: None,
+                },
+                finish_reason: None,
+            }],
+            usage: None,
+            system_fingerprint: None,
+        }
+    }
+
+    fn final_chunk() -> ChatCompletionChunk {
+        ChatCompletionChunk {
+            id: "chunk".into(),
+            object: "chat.completion.chunk".into(),
+            created: 0,
+            model: "test-model".into(),
+            choices: vec![ChatChunkChoice {
+                index: 0,
+                delta: ChatChunkDelta::default(),
+                finish_reason: Some(FinishReason::Stop),
+            }],
+            usage: None,
+            system_fingerprint: None,
+        }
+    }
+
+    #[tokio::test]
+    async fn happy_path_returns_response_and_metrics() {
+        let chunks: Vec<Result<ChatCompletionChunk, SamplingError>> =
+            vec![Ok(text_chunk("hello")), Ok(final_chunk())];
+        let raw = stream::iter(chunks).boxed();
+        let events = stream_chat_completions(raw, None, rid(), Duration::from_secs(60));
+
+        let (response, _metrics) = collect_response(events)
+            .await
+            .expect("happy path returns Ok");
+        let a = response.assistant().expect("assistant item present");
+        assert_eq!(a.content.as_ref(), "hello");
+        assert_eq!(response.stop_reason, Some(StopReason::Stop));
+    }
+
+    #[tokio::test]
+    async fn failure_path_returns_error() {
+        let chunks: Vec<Result<ChatCompletionChunk, SamplingError>> = vec![
+            Ok(text_chunk("partial")),
+            Err(SamplingError::EventStreamError("boom".into())),
+        ];
+        let raw = stream::iter(chunks).boxed();
+        let events = stream_chat_completions(raw, None, rid(), Duration::from_secs(60));
+
+        let err = collect_response(events).await.expect_err("error returned");
+        assert!(err.message.contains("boom"));
+    }
+
+    #[tokio::test]
+    async fn truncated_stream_returns_error() {
+        let truncated = stream::iter(vec![SamplingEvent::StreamStarted {
+            request_id: rid(),
+            timestamp_ms: 0,
+        }]);
+        let err = collect_response(truncated)
+            .await
+            .expect_err("truncated stream returns Err");
+        assert_eq!(err.kind, SamplingErrorKind::Api);
+        assert!(err.message.contains("stream ended without"));
+    }
+
+    #[tokio::test]
+    async fn intermediate_events_are_dropped() {
+        let token = SamplingEvent::ChannelToken {
+            request_id: rid(),
+            channel: SamplingChannel::Text,
+            text: "hi".into(),
+            chunk_index: 1,
+        };
+        let completed = SamplingEvent::Completed {
+            request_id: rid(),
+            response: Box::new(ConversationResponse {
+                items: vec![ConversationItem::assistant("hi")],
+                stop_reason: Some(StopReason::Stop),
+                usage: None,
+                cost_usd_ticks: None,
+                message_chunks_emitted: 1,
+                doom_loop_signals: Vec::new(),
+                stop_message: None,
+            }),
+            metrics: InferenceLatencyStats::default(),
+        };
+        let s = stream::iter(vec![token, completed]);
+        let (response, _) = collect_response(s).await.expect("ok");
+        let a = response.assistant().expect("assistant item present");
+        assert_eq!(a.content.as_ref(), "hi");
+    }
+}

crates/codegen/xai-grok-sampler/src/stream/messages.rs

@@ -0,0 +1,530 @@
+//! Layer-2 stream transform for the Anthropic Messages API.
+//!
+//! Consumes a raw `MessageStreamEvent` stream and produces
+//! [`SamplingEvent`]s. Pure: no I/O, no shell coupling.
+
+use std::collections::BTreeMap;
+use std::time::{Duration, Instant};
+
+use futures_util::StreamExt;
+use futures_util::stream::{BoxStream, Stream};
+
+use xai_grok_sampling_types::messages::{self, MessageStreamEvent};
+use xai_grok_sampling_types::{
+    AssistantItem, ConversationItem, ConversationResponse, ResponseModelMetadata, SamplingError,
+    StopReason, TokenUsage, ToolCall, rs,
+};
+
+use crate::events::{SamplingChannel, SamplingErrorInfo, SamplingEvent};
+use crate::metrics::InferenceLatencyStats;
+use crate::types::RequestId;
+
+/// Returns whether a Messages API event reflects real model progress
+/// rather than a liveness-only heartbeat (Ping).
+pub(crate) fn messages_event_has_meaningful_content(event: &MessageStreamEvent) -> bool {
+    match event {
+        MessageStreamEvent::Ping => false,
+        MessageStreamEvent::MessageStart { .. }
+        | MessageStreamEvent::MessageDelta { .. }
+        | MessageStreamEvent::MessageStop
+        | MessageStreamEvent::ContentBlockStart { .. }
+        | MessageStreamEvent::ContentBlockDelta { .. }
+        | MessageStreamEvent::ContentBlockStop { .. }
+        | MessageStreamEvent::Error { .. } => true,
+    }
+}
+
+/// Per-block streaming accumulator. The Anthropic Messages API reports
+/// content as a sequence of indexed blocks (text / thinking /
+/// tool_use), each with start / delta / stop events. We accumulate
+/// per-index and finalize each block on `ContentBlockStop`.
+struct BlockState {
+    block_type: BlockType,
+    text_acc: String,
+    tool_name: String,
+    tool_id: String,
+    args_acc: String,
+    thinking_acc: String,
+    signature: String,
+}
+
+#[derive(Debug, Clone, Copy, PartialEq, Eq)]
+enum BlockType {
+    Text,
+    ToolUse,
+    Thinking,
+}
+
+/// Transform a raw Anthropic Messages API stream into a stream of
+/// [`SamplingEvent`]s.
+///
+/// Yields exactly one terminal event ([`SamplingEvent::Completed`] or
+/// [`SamplingEvent::Failed`]) per request. Server-side `Error` events
+/// translate to `SamplingError::Api { status: 500, .. }` so the actor's
+/// retry loop treats them as retryable transport-level errors.
+pub fn stream_messages<'a>(
+    raw_stream: BoxStream<'a, Result<MessageStreamEvent, SamplingError>>,
+    model_metadata: Option<ResponseModelMetadata>,
+    request_id: RequestId,
+    idle_timeout: Duration,
+) -> impl Stream<Item = SamplingEvent> + Send + 'a {
+    async_stream::stream! {
+        use messages::{ContentBlock, StreamDelta};
+
+        let stream_start = Instant::now();
+        let mut chunk_timestamps: Vec<Instant> = Vec::new();
+
+        yield SamplingEvent::StreamStarted {
+            request_id: request_id.clone(),
+            timestamp_ms: chrono::Utc::now().timestamp_millis(),
+        };
+
+        if let Some(metadata) = model_metadata {
+            yield SamplingEvent::ModelMetadata {
+                request_id: request_id.clone(),
+                metadata,
+            };
+        }
+
+        // Per-block accumulators keyed by content block index.
+        let mut blocks: BTreeMap<u32, BlockState> = BTreeMap::new();
+
+        // Final-message-level accumulators
+        let mut final_model: Option<String> = None;
+        // Anthropic Messages API `input_tokens` is the uncached portion; cache hits and writes are reported
+        // in separate buckets and must be summed for the true total prompt size.
+        let mut final_input_tokens: u32 = 0;
+        let mut final_cache_read_input_tokens: u32 = 0;
+        let mut final_cache_creation_input_tokens: u32 = 0;
+        let mut final_output_tokens: u32 = 0;
+        let mut final_stop_reason: Option<StopReason> = None;
+        let mut final_stop_message: Option<String> = None;
+
+        // Assistant-response accumulators (built up as ContentBlockStop
+        // events fire). Reasoning is collected into a synthesized
+        // `rs::ReasoningItem` and emitted as a sibling
+        // `ConversationItem::Reasoning` before the trailing Assistant.
+        let mut assistant_text = String::new();
+        let mut assistant_tool_calls: Vec<ToolCall> = Vec::new();
+        let mut assistant_reasoning: Option<rs::ReasoningItem> = None;
+
+        // Index counters
+        let mut chunk_index: u64 = 0;
+        let mut message_chunk_count: u64 = 0;
+        let mut first_token_emitted = false;
+        let mut last_content_chunk_at = Instant::now();
+
+        // Tool-call index counter for per-tool deltas (separate from
+        // the block index, which can be interleaved with text/thinking
+        // blocks).
+        let mut next_tool_index: u32 = 0;
+        let mut block_to_tool_index: BTreeMap<u32, u32> = BTreeMap::new();
+
+        let mut stream = raw_stream;
+        loop {
+            let event_result = match tokio::time::timeout(idle_timeout, stream.next()).await {
+                Ok(Some(event_result)) => event_result,
+                Ok(None) => break,
+                Err(_elapsed) => {
+                    let err = SamplingError::IdleTimeout {
+                        elapsed_secs: idle_timeout.as_secs(),
+                    };
+                    yield SamplingEvent::Failed {
+                        request_id: request_id.clone(),
+                        error: SamplingErrorInfo::from(&err),
+                    };
+                    return;
+                }
+            };
+
+            let event = match event_result {
+                Ok(event) => event,
+                Err(err) => {
+                    yield SamplingEvent::Failed {
+                        request_id: request_id.clone(),
+                        error: SamplingErrorInfo::from(&err),
+                    };
+                    return;
+                }
+            };
+
+            let event_has_content = messages_event_has_meaningful_content(&event);
+
+            match event {
+                MessageStreamEvent::MessageStart { message } => {
+                    final_model = Some(message.model.clone());
+                    final_input_tokens = message.usage.input_tokens;
+                    final_cache_read_input_tokens = message.usage.cache_read_input_tokens;
+                    final_cache_creation_input_tokens = message.usage.cache_creation_input_tokens;
+                }
+
+                MessageStreamEvent::ContentBlockStart {
+                    index,
+                    content_block,
+                } => match content_block {
+                    ContentBlock::Thinking {
+                        thinking,
+                        signature,
+                    } => {
+                        blocks.insert(
+                            index,
+                            BlockState {
+                                block_type: BlockType::Thinking,
+                                text_acc: String::new(),
+                                tool_name: String::new(),
+                                tool_id: String::new(),
+                                args_acc: String::new(),
+                                thinking_acc: thinking.clone(),
+                                signature: signature.clone(),
+                            },
+                        );
+                        if !first_token_emitted {
+                            first_token_emitted = true;
+                            yield SamplingEvent::FirstToken {
+                                request_id: request_id.clone(),
+                            };
+                        }
+                    }
+                    ContentBlock::Text { text, .. } => {
+                        blocks.insert(
+                            index,
+                            BlockState {
+                                block_type: BlockType::Text,
+                                text_acc: text.clone(),
+                                tool_name: String::new(),
+                                tool_id: String::new(),
+                                args_acc: String::new(),
+                                thinking_acc: String::new(),
+                                signature: String::new(),
+                            },
+                        );
+                        if !first_token_emitted {
+                            first_token_emitted = true;
+                            yield SamplingEvent::FirstToken {
+                                request_id: request_id.clone(),
+                            };
+                        }
+                    }
+                    ContentBlock::ToolUse {
+                        id,
+                        name,
+                        input: _,
+                    } => {
+                        let tool_index = next_tool_index;
+                        next_tool_index += 1;
+                        block_to_tool_index.insert(index, tool_index);
+
+                        blocks.insert(
+                            index,
+                            BlockState {
+                                block_type: BlockType::ToolUse,
+                                text_acc: String::new(),
+                                tool_name: name.clone(),
+                                tool_id: id.clone(),
+                                // Anthropic Messages API streams arguments via
+                                // InputJsonDelta events; starting from
+                                // "{}" then appending fragments would
+                                // produce invalid JSON.
+                                args_acc: String::new(),
+                                thinking_acc: String::new(),
+                                signature: String::new(),
+                            },
+                        );
+
+                        // Emit initial id+name so subscribers can pre-allocate
+                        // UI for the tool call before arguments stream in.
+                        yield SamplingEvent::ToolCallDelta {
+                            request_id: request_id.clone(),
+                            tool_index,
+                            id: Some(id),
+                            name: Some(name),
+                            arguments_delta: None,
+                        };
+                    }
+                    _ => {} // Image / ToolResult are not expected in assistant streams.
+                },
+
+                MessageStreamEvent::ContentBlockDelta { index, delta } => {
+                    if let Some(state) = blocks.get_mut(&index) {
+                        match delta {
+                            StreamDelta::ThinkingDelta { thinking } => {
+                                if !thinking.is_empty() {
+                                    state.thinking_acc.push_str(&thinking);
+                                    if !first_token_emitted {
+                                        first_token_emitted = true;
+                                        yield SamplingEvent::FirstToken {
+                                            request_id: request_id.clone(),
+                                        };
+                                    }
+                                    chunk_index += 1;
+                                    yield SamplingEvent::ChannelToken {
+                                        request_id: request_id.clone(),
+                                        channel: SamplingChannel::Reasoning,
+                                        text: thinking,
+                                        chunk_index,
+                                    };
+                                }
+                            }
+                            StreamDelta::SignatureDelta { signature } => {
+                                state.signature = signature;
+                            }
+                            StreamDelta::TextDelta { text } => {
+                                if !text.is_empty() {
+                                    state.text_acc.push_str(&text);
+                                    if !first_token_emitted {
+                                        first_token_emitted = true;
+                                        yield SamplingEvent::FirstToken {
+                                            request_id: request_id.clone(),
+                                        };
+                                    }
+                                    chunk_timestamps.push(Instant::now());
+                                    chunk_index += 1;
+                                    message_chunk_count += 1;
+                                    yield SamplingEvent::ChannelToken {
+                                        request_id: request_id.clone(),
+                                        channel: SamplingChannel::Text,
+                                        text,
+                                        chunk_index,
+                                    };
+                                }
+                            }
+                            StreamDelta::InputJsonDelta { partial_json } => {
+                                state.args_acc.push_str(&partial_json);
+                                if let Some(&tool_index) = block_to_tool_index.get(&index) {
+                                    yield SamplingEvent::ToolCallDelta {
+                                        request_id: request_id.clone(),
+                                        tool_index,
+                                        id: None,
+                                        name: None,
+                                        arguments_delta: Some(partial_json),
+                                    };
+                                }
+                            }
+                        }
+                    }
+                }
+
+                MessageStreamEvent::ContentBlockStop { index } => {
+                    if let Some(state) = blocks.remove(&index) {
+                        match state.block_type {
+                            BlockType::Text => {
+                                if !state.text_acc.is_empty() {
+                                    if !assistant_text.is_empty() {
+                                        assistant_text.push('\n');
+                                    }
+                                    assistant_text.push_str(&state.text_acc);
+                                }
+                            }
+                            BlockType::Thinking => {
+                                if !state.thinking_acc.is_empty() || !state.signature.is_empty() {
+                                    // Anthropic Messages API `Thinking` blocks uniquely
+                                    // carry an encrypted `signature` distinct
+                                    // from the text; either field may be
+                                    // empty. Build directly rather than via
+                                    // `synthesized_reasoning_item` since the
+                                    // helper assumes a non-empty summary.
+                                    let summary = if state.thinking_acc.is_empty() {
+                                        vec![]
+                                    } else {
+                                        vec![rs::SummaryPart::SummaryText(
+                                            rs::SummaryTextContent {
+                                                text: state.thinking_acc,
+                                            },
+                                        )]
+                                    };
+                                    let encrypted_content = if state.signature.is_empty() {
+                                        None
+                                    } else {
+                                        Some(state.signature)
+                                    };
+                                    assistant_reasoning = Some(rs::ReasoningItem {
+                                        id: String::new(),
+                                        summary,
+                                        content: None,
+                                        encrypted_content,
+                                        status: None,
+                                    });
+                                }
+                            }
+                            BlockType::ToolUse => {
+                                assistant_tool_calls.push(ToolCall {
+                                    id: std::sync::Arc::<str>::from(state.tool_id),
+                                    name: state.tool_name,
+                                    arguments: std::sync::Arc::<str>::from(state.args_acc),
+                                });
+                            }
+                        }
+                    }
+                }
+
+                MessageStreamEvent::MessageDelta { delta, usage } => {
+                    // Normalize the provider's stop detail to a plain message;
+                    // the shell logs it when it surfaces a refusal.
+                    if let Some(details) = delta.stop_details {
+                        final_stop_message = details.explanation;
+                    }
+                    final_stop_reason = delta.stop_reason.map(|sr| match sr {
+                        messages::StopReason::EndTurn => StopReason::Stop,
+                        messages::StopReason::MaxTokens => StopReason::Length,
+                        messages::StopReason::StopSequence => StopReason::Stop,
+                        messages::StopReason::ToolUse => StopReason::ToolCalls,
+                        // The model declined to continue; whatever streamed is
+                        // the complete response, so end the turn cleanly.
+                        messages::StopReason::Refusal => StopReason::ContentFilter,
+                        messages::StopReason::PauseTurn => {
+                            // Anthropic Messages API expects a resend-to-continue; we end the
+                            // turn instead, so leave a triage trail.
+                            tracing::warn!(
+                                wire_stop_reason = "pause_turn",
+                                "pause_turn ended the turn like stop (no auto-continue)"
+                            );
+                            StopReason::Stop
+                        }
+                        messages::StopReason::ModelContextWindowExceeded => {
+                            // Output-side overflow on a successful stream: stays in the
+                            // max_tokens truncation class — compact-on-error recovery needs
+                            // an Api error carrying model metadata plus a prompt-side
+                            // overflow, neither of which exists here.
+                            tracing::warn!(
+                                wire_stop_reason = "model_context_window_exceeded",
+                                "context window hit mid-generation; surfacing as max_tokens truncation"
+                            );
+                            StopReason::Length
+                        }
+                        messages::StopReason::Unknown(wire) => {
+                            tracing::warn!(
+                                wire_stop_reason = %wire,
+                                "unrecognized stop_reason in messages stream; treating as stop"
+                            );
+                            StopReason::Stop
+                        }
+                    });
+                    final_output_tokens = usage.output_tokens;
+                    // Optional on the delta; preserve message_start values when omitted.
+                    if let Some(input) = usage.input_tokens {
+                        final_input_tokens = input;
+                    }
+                    if let Some(cache_read) = usage.cache_read_input_tokens {
+                        final_cache_read_input_tokens = cache_read;
+                    }
+                    if let Some(cache_creation) = usage.cache_creation_input_tokens {
+                        final_cache_creation_input_tokens = cache_creation;
+                    }
+                }
+
+                MessageStreamEvent::MessageStop => {
+                    // Final message complete; the loop exits naturally
+                    // when the underlying stream ends.
+                }
+
+                MessageStreamEvent::Ping => {
+                    // Liveness only, no action; the inner timeout was
+                    // already reset above by the successful `next()`.
+                }
+
+                MessageStreamEvent::Error { error } => {
+                    let error_message = format!("{}: {}", error.r#type, error.message);
+                    let err = SamplingError::Api {
+                        status: reqwest::StatusCode::INTERNAL_SERVER_ERROR,
+                        message: error_message,
+                        model_metadata: None,
+                        retry_after_secs: None,
+                        should_retry: None,
+                    };
+                    yield SamplingEvent::Failed {
+                        request_id: request_id.clone(),
+                        error: SamplingErrorInfo::from(&err),
+                    };
+                    return;
+                }
+            }
+
+            if event_has_content {
+                last_content_chunk_at = Instant::now();
+            } else if last_content_chunk_at.elapsed() > idle_timeout {
+                let err = SamplingError::IdleTimeout {
+                    elapsed_secs: idle_timeout.as_secs(),
+                };
+                yield SamplingEvent::Failed {
+                    request_id: request_id.clone(),
+                    error: SamplingErrorInfo::from(&err),
+                };
+                return;
+            }
+        }
+
+        if final_stop_reason == Some(StopReason::Length) {
+            yield SamplingEvent::Failed {
+                request_id: request_id.clone(),
+                error: SamplingErrorInfo::from(&SamplingError::MaxTokensTruncation),
+            };
+            return;
+        }
+
+        // ── 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
+            .saturating_add(final_cache_read_input_tokens)
+            .saturating_add(final_cache_creation_input_tokens);
+        let usage = if total_prompt_tokens > 0 || final_output_tokens > 0 {
+            Some(TokenUsage {
+                prompt_tokens: total_prompt_tokens,
+                completion_tokens: final_output_tokens,
+                total_tokens: total_prompt_tokens.saturating_add(final_output_tokens),
+                reasoning_tokens: 0,
+                cached_prompt_tokens: final_cache_read_input_tokens,
+            })
+        } else {
+            None
+        };
+
+        let stop_reason = if !assistant_tool_calls.is_empty() {
+            // Completed tool_use blocks win even over Refusal: the calls are
+            // real model output the agent loop must resolve.
+            Some(StopReason::ToolCalls)
+        } else {
+            final_stop_reason
+        };
+
+        let assistant_item = ConversationItem::Assistant(AssistantItem {
+            content: std::sync::Arc::<str>::from(assistant_text),
+            tool_calls: assistant_tool_calls,
+            model_id: Some(model_id),
+            model_fingerprint: None,
+            // The Messages API does not echo the applied reasoning effort.
+            reasoning_effort: None,
+        });
+
+        let mut items: Vec<ConversationItem> = Vec::new();
+        if let Some(r) = assistant_reasoning {
+            items.push(ConversationItem::Reasoning(r));
+        }
+        items.push(assistant_item);
+
+        let stream_end = Instant::now();
+        let metrics =
+            InferenceLatencyStats::from_timestamps(stream_start, &chunk_timestamps, stream_end);
+
+        let response = ConversationResponse {
+            items,
+            stop_reason,
+            usage,
+            // Anthropic Messages API carries no cost on the wire.
+            cost_usd_ticks: None,
+            message_chunks_emitted: message_chunk_count,
+            doom_loop_signals: Vec::new(),
+            stop_message: final_stop_message,
+        };
+
+        yield SamplingEvent::Completed {
+            request_id: request_id.clone(),
+            response: Box::new(response),
+            metrics,
+        };
+    }
+}
+
+#[cfg(test)]
+#[path = "messages_tests.rs"]
+mod tests;

crates/codegen/xai-grok-sampler/src/stream/messages_tests.rs

@@ -0,0 +1,655 @@
+//! Unit tests for the [`super`] Messages L2 stream transform. Extracted
+//! from `messages.rs` so the implementation reads top-to-bottom; wired in
+//! via `#[path = "messages_tests.rs"] mod tests;` in messages.rs.
+
+use super::*;
+use futures_util::stream;
+use std::pin::pin;
+use xai_grok_sampling_types::messages::{
+    ContentBlock, MessageDeltaBody, MessageDeltaUsage, MessagesResponse, MessagesUsage,
+    StreamDelta, StreamError,
+};
+
+fn rid() -> RequestId {
+    RequestId::from("msg-test")
+}
+
+fn message_start() -> MessageStreamEvent {
+    MessageStreamEvent::MessageStart {
+        message: MessagesResponse {
+            id: "msg_1".into(),
+            r#type: "message".into(),
+            role: "assistant".into(),
+            content: vec![],
+            model: "messages-compatible-model".into(),
+            stop_reason: None,
+            usage: MessagesUsage {
+                input_tokens: 10,
+                output_tokens: 0,
+                cache_creation_input_tokens: 0,
+                cache_read_input_tokens: 0,
+            },
+        },
+    }
+}
+
+fn text_block_start(index: u32) -> MessageStreamEvent {
+    MessageStreamEvent::ContentBlockStart {
+        index,
+        content_block: ContentBlock::Text {
+            text: String::new(),
+            cache_control: None,
+        },
+    }
+}
+
+fn text_delta(index: u32, text: &str) -> MessageStreamEvent {
+    MessageStreamEvent::ContentBlockDelta {
+        index,
+        delta: StreamDelta::TextDelta { text: text.into() },
+    }
+}
+
+fn block_stop(index: u32) -> MessageStreamEvent {
+    MessageStreamEvent::ContentBlockStop { index }
+}
+
+fn message_delta_with_stop(stop: messages::StopReason) -> MessageStreamEvent {
+    MessageStreamEvent::MessageDelta {
+        delta: MessageDeltaBody {
+            stop_reason: Some(stop),
+            stop_details: None,
+        },
+        usage: MessageDeltaUsage {
+            output_tokens: 5,
+            input_tokens: Some(10),
+            cache_read_input_tokens: None,
+            cache_creation_input_tokens: None,
+        },
+    }
+}
+
+/// A refusal `message_delta` carrying a provider `stop_details.explanation`,
+/// mirroring the Anthropic Messages API ToS auto-refusal wire shape.
+fn message_delta_refusal_with_explanation(explanation: &str) -> MessageStreamEvent {
+    MessageStreamEvent::MessageDelta {
+        delta: MessageDeltaBody {
+            stop_reason: Some(messages::StopReason::Refusal),
+            stop_details: Some(messages::StopDetails {
+                r#type: Some("refusal".to_string()),
+                category: Some("frontier_llm".to_string()),
+                explanation: Some(explanation.to_string()),
+            }),
+        },
+        usage: MessageDeltaUsage {
+            output_tokens: 0,
+            input_tokens: Some(10),
+            cache_read_input_tokens: None,
+            cache_creation_input_tokens: None,
+        },
+    }
+}
+
+async fn collect(s: impl Stream<Item = SamplingEvent>) -> Vec<SamplingEvent> {
+    let mut out = Vec::new();
+    let mut s = pin!(s);
+    while let Some(ev) = s.next().await {
+        out.push(ev);
+    }
+    out
+}
+
+#[tokio::test]
+async fn empty_stream_yields_started_then_completed() {
+    let raw = stream::iter(Vec::<Result<MessageStreamEvent, SamplingError>>::new()).boxed();
+    let events = collect(stream_messages(raw, None, rid(), Duration::from_secs(60))).await;
+    assert_eq!(events.len(), 2);
+    assert!(matches!(events[0], SamplingEvent::StreamStarted { .. }));
+    assert!(matches!(events[1], SamplingEvent::Completed { .. }));
+}
+
+#[tokio::test]
+async fn text_block_assembles_into_completed_response() {
+    let events: Vec<Result<MessageStreamEvent, SamplingError>> = vec![
+        Ok(message_start()),
+        Ok(text_block_start(0)),
+        Ok(text_delta(0, "Hello, ")),
+        Ok(text_delta(0, "world!")),
+        Ok(block_stop(0)),
+        Ok(message_delta_with_stop(messages::StopReason::EndTurn)),
+        Ok(MessageStreamEvent::MessageStop),
+    ];
+    let raw = stream::iter(events).boxed();
+    let evs = collect(stream_messages(raw, None, rid(), Duration::from_secs(60))).await;
+
+    let text_tokens: Vec<&str> = evs
+        .iter()
+        .filter_map(|e| match e {
+            SamplingEvent::ChannelToken {
+                channel: SamplingChannel::Text,
+                text,
+                ..
+            } => Some(text.as_str()),
+            _ => None,
+        })
+        .collect();
+    assert_eq!(text_tokens, vec!["Hello, ", "world!"]);
+
+    match evs.last().unwrap() {
+        SamplingEvent::Completed { response, .. } => {
+            let a = response.assistant().expect("assistant item present");
+            assert_eq!(a.content.as_ref(), "Hello, world!");
+            assert_eq!(a.model_id.as_deref(), Some("messages-compatible-model"));
+            assert_eq!(response.stop_reason, Some(StopReason::Stop));
+            let u = response.usage.as_ref().expect("usage extracted");
+            assert_eq!(u.prompt_tokens, 10);
+            assert_eq!(u.completion_tokens, 5);
+        }
+        other => panic!("expected Completed, got {other:?}"),
+    }
+}
+
+#[tokio::test]
+async fn thinking_block_emits_reasoning_channel_and_preserved_in_response() {
+    let thinking_start = MessageStreamEvent::ContentBlockStart {
+        index: 0,
+        content_block: ContentBlock::Thinking {
+            thinking: String::new(),
+            signature: String::new(),
+        },
+    };
+    let thinking_delta = MessageStreamEvent::ContentBlockDelta {
+        index: 0,
+        delta: StreamDelta::ThinkingDelta {
+            thinking: "let me think...".into(),
+        },
+    };
+    let sig_delta = MessageStreamEvent::ContentBlockDelta {
+        index: 0,
+        delta: StreamDelta::SignatureDelta {
+            signature: "abc123".into(),
+        },
+    };
+    let events: Vec<Result<MessageStreamEvent, SamplingError>> = vec![
+        Ok(message_start()),
+        Ok(thinking_start),
+        Ok(thinking_delta),
+        Ok(sig_delta),
+        Ok(block_stop(0)),
+        Ok(MessageStreamEvent::MessageStop),
+    ];
+    let raw = stream::iter(events).boxed();
+    let evs = collect(stream_messages(raw, None, rid(), Duration::from_secs(60))).await;
+
+    let reasoning_tokens: Vec<&str> = evs
+        .iter()
+        .filter_map(|e| match e {
+            SamplingEvent::ChannelToken {
+                channel: SamplingChannel::Reasoning,
+                text,
+                ..
+            } => Some(text.as_str()),
+            _ => None,
+        })
+        .collect();
+    assert_eq!(reasoning_tokens, vec!["let me think..."]);
+
+    match evs.last().unwrap() {
+        SamplingEvent::Completed { response, .. } => {
+            let r = response
+                .reasoning_items()
+                .next()
+                .expect("reasoning sibling preserved");
+            let rs::SummaryPart::SummaryText(t) = &r.summary[0];
+            assert_eq!(t.text, "let me think...");
+            assert_eq!(r.encrypted_content.as_deref(), Some("abc123"));
+        }
+        other => panic!("expected Completed, got {other:?}"),
+    }
+}
+
+#[tokio::test]
+async fn tool_use_block_assembles_into_tool_call() {
+    let tool_start = MessageStreamEvent::ContentBlockStart {
+        index: 0,
+        content_block: ContentBlock::ToolUse {
+            id: "call_xyz".into(),
+            name: "do_thing".into(),
+            input: serde_json::json!({}),
+        },
+    };
+    let arg_delta_1 = MessageStreamEvent::ContentBlockDelta {
+        index: 0,
+        delta: StreamDelta::InputJsonDelta {
+            partial_json: "{\"x\":".into(),
+        },
+    };
+    let arg_delta_2 = MessageStreamEvent::ContentBlockDelta {
+        index: 0,
+        delta: StreamDelta::InputJsonDelta {
+            partial_json: "1}".into(),
+        },
+    };
+    let events: Vec<Result<MessageStreamEvent, SamplingError>> = vec![
+        Ok(message_start()),
+        Ok(tool_start),
+        Ok(arg_delta_1),
+        Ok(arg_delta_2),
+        Ok(block_stop(0)),
+        Ok(message_delta_with_stop(messages::StopReason::ToolUse)),
+        Ok(MessageStreamEvent::MessageStop),
+    ];
+    let raw = stream::iter(events).boxed();
+    let evs = collect(stream_messages(raw, None, rid(), Duration::from_secs(60))).await;
+
+    // Should yield three ToolCallDelta events: id+name, then two
+    // arguments fragments.
+    let deltas: Vec<_> = evs
+        .iter()
+        .filter_map(|e| match e {
+            SamplingEvent::ToolCallDelta {
+                tool_index,
+                id,
+                name,
+                arguments_delta,
+                ..
+            } => Some((
+                *tool_index,
+                id.clone(),
+                name.clone(),
+                arguments_delta.clone(),
+            )),
+            _ => None,
+        })
+        .collect();
+    assert_eq!(deltas.len(), 3);
+    assert_eq!(deltas[0].0, 0);
+    assert_eq!(deltas[0].1.as_deref(), Some("call_xyz"));
+    assert_eq!(deltas[0].2.as_deref(), Some("do_thing"));
+    assert_eq!(deltas[0].3, None);
+    assert_eq!(deltas[1].3.as_deref(), Some("{\"x\":"));
+    assert_eq!(deltas[2].3.as_deref(), Some("1}"));
+
+    match evs.last().unwrap() {
+        SamplingEvent::Completed { response, .. } => {
+            let calls = response.tool_calls();
+            assert_eq!(calls.len(), 1);
+            assert_eq!(calls[0].id.as_ref(), "call_xyz");
+            assert_eq!(calls[0].name, "do_thing");
+            assert_eq!(calls[0].arguments.as_ref(), "{\"x\":1}");
+            assert_eq!(response.stop_reason, Some(StopReason::ToolCalls));
+        }
+        other => panic!("expected Completed, got {other:?}"),
+    }
+}
+
+/// Regression: a stream whose terminal `message_delta` carries
+/// `stop_reason: "refusal"` must complete cleanly — not error out and
+/// discard the already-streamed response.
+#[tokio::test]
+async fn refusal_stop_reason_completes_stream() {
+    let events: Vec<Result<MessageStreamEvent, SamplingError>> = vec![
+        Ok(message_start()),
+        Ok(text_block_start(0)),
+        Ok(text_delta(0, "I can't help with that.")),
+        Ok(block_stop(0)),
+        Ok(message_delta_with_stop(messages::StopReason::Refusal)),
+        Ok(MessageStreamEvent::MessageStop),
+    ];
+    let raw = stream::iter(events).boxed();
+    let evs = collect(stream_messages(raw, None, rid(), Duration::from_secs(60))).await;
+
+    assert!(
+        !evs.iter()
+            .any(|e| matches!(e, SamplingEvent::Failed { .. })),
+        "refusal stream must not yield Failed: {evs:?}"
+    );
+    match evs.last().unwrap() {
+        SamplingEvent::Completed { response, .. } => {
+            let a = response.assistant().expect("assistant item present");
+            assert_eq!(a.content.as_ref(), "I can't help with that.");
+            assert_eq!(response.stop_reason, Some(StopReason::ContentFilter));
+        }
+        other => panic!("expected Completed, got {other:?}"),
+    }
+}
+
+/// A refusal `stop_details.explanation` on the terminal delta must be
+/// normalized onto the completed `ConversationResponse.stop_message` so the
+/// agent loop can surface the provider's reason (empty-turn silence otherwise).
+#[tokio::test]
+async fn refusal_stop_message_flows_to_response() {
+    let explanation = "This request was blocked by the provider's content policy.";
+    let events: Vec<Result<MessageStreamEvent, SamplingError>> = vec![
+        Ok(message_start()),
+        Ok(message_delta_refusal_with_explanation(explanation)),
+        Ok(MessageStreamEvent::MessageStop),
+    ];
+    let raw = stream::iter(events).boxed();
+    let evs = collect(stream_messages(raw, None, rid(), Duration::from_secs(60))).await;
+
+    match evs.last().unwrap() {
+        SamplingEvent::Completed { response, .. } => {
+            assert_eq!(response.stop_reason, Some(StopReason::ContentFilter));
+            assert_eq!(
+                response.stop_message.as_deref(),
+                Some(explanation),
+                "provider explanation normalized onto stop_message"
+            );
+        }
+        other => panic!("expected Completed, got {other:?}"),
+    }
+}
+
+#[tokio::test]
+async fn pause_turn_and_unknown_stop_reasons_complete_as_stop() {
+    for stop in [
+        messages::StopReason::PauseTurn,
+        messages::StopReason::Unknown("mystery_reason".to_string()),
+    ] {
+        let label = format!("{stop:?}");
+        let events: Vec<Result<MessageStreamEvent, SamplingError>> = vec![
+            Ok(message_start()),
+            Ok(text_block_start(0)),
+            Ok(text_delta(0, "partial answer")),
+            Ok(block_stop(0)),
+            Ok(message_delta_with_stop(stop)),
+            Ok(MessageStreamEvent::MessageStop),
+        ];
+        let raw = stream::iter(events).boxed();
+        let evs = collect(stream_messages(raw, None, rid(), Duration::from_secs(60))).await;
+        match evs.last().unwrap() {
+            SamplingEvent::Completed { response, .. } => {
+                assert_eq!(
+                    response.stop_reason,
+                    Some(StopReason::Stop),
+                    "{label} must end the turn like stop"
+                );
+            }
+            other => panic!("{label}: expected Completed, got {other:?}"),
+        }
+    }
+}
+
+/// Pins the model_context_window_exceeded decision: it stays in the
+/// max_tokens truncation class (fatal, non-retryable), not the
+/// context-length Api class.
+#[tokio::test]
+async fn model_context_window_exceeded_fails_as_max_tokens_truncation() {
+    let events: Vec<Result<MessageStreamEvent, SamplingError>> = vec![
+        Ok(message_start()),
+        Ok(text_block_start(0)),
+        Ok(text_delta(0, "truncated answ")),
+        Ok(block_stop(0)),
+        Ok(message_delta_with_stop(
+            messages::StopReason::ModelContextWindowExceeded,
+        )),
+        Ok(MessageStreamEvent::MessageStop),
+    ];
+    let raw = stream::iter(events).boxed();
+    let evs = collect(stream_messages(raw, None, rid(), Duration::from_secs(60))).await;
+
+    assert!(
+        !evs.iter()
+            .any(|e| matches!(e, SamplingEvent::Completed { .. })),
+        "context-window truncation must not complete: {evs:?}"
+    );
+    match evs.last().unwrap() {
+        SamplingEvent::Failed { error, .. } => {
+            assert_eq!(
+                error.kind,
+                crate::events::SamplingErrorKind::MaxTokensTruncation
+            );
+            assert!(!error.is_retryable, "truncation is deterministic");
+        }
+        other => panic!("expected Failed(MaxTokensTruncation), got {other:?}"),
+    }
+}
+
+/// Pins the pre-existing override: completed tool_use blocks beat a terminal
+/// Refusal, so the agent loop still resolves the calls.
+#[tokio::test]
+async fn refusal_after_tool_use_blocks_keeps_tool_calls_stop_reason() {
+    let tool_start = MessageStreamEvent::ContentBlockStart {
+        index: 0,
+        content_block: ContentBlock::ToolUse {
+            id: "call_refused".into(),
+            name: "do_thing".into(),
+            input: serde_json::json!({}),
+        },
+    };
+    let arg_delta = MessageStreamEvent::ContentBlockDelta {
+        index: 0,
+        delta: StreamDelta::InputJsonDelta {
+            partial_json: "{}".into(),
+        },
+    };
+    let events: Vec<Result<MessageStreamEvent, SamplingError>> = vec![
+        Ok(message_start()),
+        Ok(tool_start),
+        Ok(arg_delta),
+        Ok(block_stop(0)),
+        Ok(message_delta_with_stop(messages::StopReason::Refusal)),
+        Ok(MessageStreamEvent::MessageStop),
+    ];
+    let raw = stream::iter(events).boxed();
+    let evs = collect(stream_messages(raw, None, rid(), Duration::from_secs(60))).await;
+
+    match evs.last().unwrap() {
+        SamplingEvent::Completed { response, .. } => {
+            assert_eq!(response.tool_calls().len(), 1);
+            assert_eq!(
+                response.stop_reason,
+                Some(StopReason::ToolCalls),
+                "tool_use blocks must win over the refusal stop_reason"
+            );
+        }
+        other => panic!("expected Completed, got {other:?}"),
+    }
+}
+
+#[tokio::test]
+async fn server_error_event_yields_failed_500() {
+    let err_event = MessageStreamEvent::Error {
+        error: StreamError {
+            r#type: "overloaded_error".into(),
+            message: "rate limit hit".into(),
+        },
+    };
+    let raw = stream::iter(vec![Ok(message_start()), Ok(err_event)]).boxed();
+    let evs = collect(stream_messages(raw, None, rid(), Duration::from_secs(60))).await;
+
+    match evs.last().unwrap() {
+        SamplingEvent::Failed { error, .. } => {
+            assert_eq!(error.kind, crate::events::SamplingErrorKind::Api);
+            assert_eq!(error.status_code, Some(500));
+            assert!(error.message.contains("overloaded_error"));
+        }
+        other => panic!("expected Failed, got {other:?}"),
+    }
+}
+
+#[tokio::test]
+async fn mid_stream_transport_error_yields_failed() {
+    let raw = stream::iter(vec![
+        Ok(message_start()),
+        Err(SamplingError::EventStreamError("conn reset".into())),
+    ])
+    .boxed();
+    let evs = collect(stream_messages(raw, None, rid(), Duration::from_secs(60))).await;
+    assert!(
+        evs.iter()
+            .any(|e| matches!(e, SamplingEvent::Failed { .. }))
+    );
+    assert!(
+        !evs.iter()
+            .any(|e| matches!(e, SamplingEvent::Completed { .. }))
+    );
+}
+
+#[tokio::test(start_paused = true)]
+async fn idle_timeout_when_stream_stalls() {
+    let raw = stream::iter(vec![Ok(message_start())])
+        .chain(stream::pending())
+        .boxed();
+    let evs = collect(stream_messages(
+        raw,
+        None,
+        rid(),
+        Duration::from_millis(100),
+    ))
+    .await;
+
+    match evs.last().unwrap() {
+        SamplingEvent::Failed { error, .. } => {
+            assert_eq!(error.kind, crate::events::SamplingErrorKind::IdleTimeout);
+        }
+        other => panic!("expected Failed(IdleTimeout), got {other:?}"),
+    }
+}
+
+#[tokio::test]
+async fn model_metadata_yielded_after_stream_started() {
+    let raw = stream::iter(vec![Ok(MessageStreamEvent::MessageStop)]).boxed();
+    let metadata = ResponseModelMetadata {
+        context_window: Some(200_000),
+        ..Default::default()
+    };
+    let evs = collect(stream_messages(
+        raw,
+        Some(metadata),
+        rid(),
+        Duration::from_secs(60),
+    ))
+    .await;
+
+    assert!(matches!(evs[0], SamplingEvent::StreamStarted { .. }));
+    assert!(matches!(evs[1], SamplingEvent::ModelMetadata { .. }));
+}
+
+#[test]
+fn meaningful_content_classifier_treats_ping_as_keepalive() {
+    assert!(!messages_event_has_meaningful_content(
+        &MessageStreamEvent::Ping
+    ));
+    assert!(messages_event_has_meaningful_content(
+        &MessageStreamEvent::MessageStop
+    ));
+}
+
+// ── Token usage: Anthropic Messages API cache-bucket accounting ────────────
+
+fn message_start_with_cache(
+    input: u32,
+    cache_read: u32,
+    cache_creation: u32,
+) -> MessageStreamEvent {
+    MessageStreamEvent::MessageStart {
+        message: MessagesResponse {
+            id: "msg_cache".into(),
+            r#type: "message".into(),
+            role: "assistant".into(),
+            content: vec![],
+            model: "messages-compatible-model".into(),
+            stop_reason: None,
+            usage: MessagesUsage {
+                input_tokens: input,
+                output_tokens: 0,
+                cache_creation_input_tokens: cache_creation,
+                cache_read_input_tokens: cache_read,
+            },
+        },
+    }
+}
+
+fn message_delta_with_cache(
+    output: u32,
+    input: Option<u32>,
+    cache_read: Option<u32>,
+    cache_creation: Option<u32>,
+) -> MessageStreamEvent {
+    MessageStreamEvent::MessageDelta {
+        delta: MessageDeltaBody {
+            stop_reason: Some(messages::StopReason::EndTurn),
+            stop_details: None,
+        },
+        usage: MessageDeltaUsage {
+            output_tokens: output,
+            input_tokens: input,
+            cache_read_input_tokens: cache_read,
+            cache_creation_input_tokens: cache_creation,
+        },
+    }
+}
+
+/// Helper: drive a minimal stream with the supplied usage events and
+/// pluck the `TokenUsage` out of the terminal `Completed` event.
+async fn usage_from_stream(events: Vec<MessageStreamEvent>) -> TokenUsage {
+    let raw = stream::iter(
+        events
+            .into_iter()
+            .map(Ok::<_, SamplingError>)
+            .collect::<Vec<_>>(),
+    )
+    .boxed();
+    let evs = collect(stream_messages(raw, None, rid(), Duration::from_secs(60))).await;
+    match evs.last().expect("at least one event") {
+        SamplingEvent::Completed { response, .. } => response
+            .usage
+            .clone()
+            .expect("usage should be emitted when prompt or output tokens > 0"),
+        other => panic!("expected Completed, got {other:?}"),
+    }
+}
+
+#[tokio::test]
+async fn prompt_tokens_sums_all_three_anthropic_buckets() {
+    // prompt_tokens = uncached + cache_read + cache_creation;
+    // cached_prompt_tokens = cache_read only (writes aren't a hit).
+    let usage = usage_from_stream(vec![
+        message_start_with_cache(100, 5000, 200),
+        text_block_start(0),
+        text_delta(0, "ok"),
+        block_stop(0),
+        message_delta_with_cache(7, None, None, None),
+        MessageStreamEvent::MessageStop,
+    ])
+    .await;
+
+    assert_eq!(usage.prompt_tokens, 100 + 5000 + 200);
+    assert_eq!(usage.cached_prompt_tokens, 5000);
+    assert_eq!(usage.completion_tokens, 7);
+    assert_eq!(usage.total_tokens, 100 + 5000 + 200 + 7);
+}
+
+#[tokio::test]
+async fn message_delta_cache_fields_override_message_start() {
+    // Providers can report zero cache at message_start and emit the real
+    // values on the final delta; honor the delta when present.
+    let usage = usage_from_stream(vec![
+        message_start_with_cache(10, 0, 0),
+        message_delta_with_cache(4, Some(10), Some(900), Some(50)),
+        MessageStreamEvent::MessageStop,
+    ])
+    .await;
+
+    assert_eq!(usage.prompt_tokens, 10 + 900 + 50);
+    assert_eq!(usage.cached_prompt_tokens, 900);
+    assert_eq!(usage.completion_tokens, 4);
+}
+
+#[tokio::test]
+async fn pure_cache_hit_with_zero_uncached_still_emits_usage() {
+    // 100% cache hit: Anthropic Messages API reports input_tokens=0 with cache_read>0.
+    // The emit-guard must still fire so callers see the cached cost.
+    let usage = usage_from_stream(vec![
+        message_start_with_cache(0, 2500, 0),
+        message_delta_with_cache(1, None, None, None),
+        MessageStreamEvent::MessageStop,
+    ])
+    .await;
+
+    assert_eq!(usage.prompt_tokens, 2500);
+    assert_eq!(usage.cached_prompt_tokens, 2500);
+    assert_eq!(usage.total_tokens, 2501);
+}

crates/codegen/xai-grok-sampler/src/stream/mod.rs

@@ -0,0 +1,19 @@
+//! 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.
+
+pub mod chat_completions;
+pub mod collect;
+pub mod messages;
+pub mod responses;
+
+pub use chat_completions::stream_chat_completions;
+pub use collect::collect_response;
+pub use messages::stream_messages;
+pub use responses::stream_responses;

crates/codegen/xai-grok-sampler/src/types.rs

@@ -0,0 +1,75 @@
+//! Core sampler types.
+
+use std::fmt;
+
+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`].
+#[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
+    }
+}
+
+impl fmt::Display for RequestId {
+    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
+        f.write_str(&self.0)
+    }
+}
+
+impl From<String> for RequestId {
+    fn from(value: String) -> Self {
+        Self(value)
+    }
+}
+
+impl From<&str> for RequestId {
+    fn from(value: &str) -> Self {
+        Self(value.to_string())
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    #[test]
+    fn from_string_roundtrips() {
+        let id: RequestId = String::from("abc-123").into();
+        assert_eq!(id.as_str(), "abc-123");
+    }
+
+    #[test]
+    fn from_str_roundtrips() {
+        let id: RequestId = "xyz-789".into();
+        assert_eq!(id.as_str(), "xyz-789");
+    }
+
+    #[test]
+    fn display_matches_inner_string() {
+        let id: RequestId = "display-me".into();
+        assert_eq!(format!("{id}"), "display-me");
+    }
+
+    #[test]
+    fn random_produces_unique_values() {
+        let a = RequestId::random();
+        let b = RequestId::random();
+        assert_ne!(a, b, "two random IDs must differ");
+        // UUIDv4 strings are 36 characters (8-4-4-4-12 hex with hyphens).
+        assert_eq!(a.as_str().len(), 36);
+    }
+}

crates/codegen/xai-grok-sampler/tests/shared_http_kill_switch.rs

@@ -0,0 +1,29 @@
+//! Kill-switch test in its own integration binary: a separate test binary is
+//! a separate process under cargo test, nextest, and Bazel alike, so the env
+//! write below cannot poison other tests and lands before the crate's
+//! once-per-process kill-switch latch first resolves.
+
+mod support;
+
+use std::sync::atomic::Ordering;
+use std::time::Duration;
+
+use support::{send_one, test_config};
+use xai_grok_sampler::SamplingClient;
+use xai_grok_test_support::spawn_counting_server;
+
+#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
+async fn kill_switch_builds_fresh_client_per_sampling_client() {
+    // Safety: the only test in this binary, set before any client exists; no
+    // concurrent env reads are possible.
+    unsafe { std::env::set_var("GROK_SAMPLER_SHARED_CLIENT", "0") };
+    let (base_url, accepts, _heads) = spawn_counting_server().await;
+    let a = SamplingClient::new(test_config(&base_url, "token-a")).unwrap();
+    let b = SamplingClient::new(test_config(&base_url, "token-b")).unwrap();
+    send_one(&a).await;
+    // Same check-in pause as the reuse test: a (hypothetically) shared pool
+    // would now yield 1 accept, so asserting 2 pins the kill switch.
+    tokio::time::sleep(Duration::from_millis(50)).await;
+    send_one(&b).await;
+    assert_eq!(accepts.load(Ordering::SeqCst), 2);
+}

crates/codegen/xai-grok-sampler/tests/shared_http_wire.rs

@@ -0,0 +1,82 @@
+//! Wire-level tests for the process-wide shared sampling client: connection
+//! reuse across `SamplingClient`s, per-config header isolation, and the
+//! pool-less HTTP/1.1 fallback. These live in their own integration binary
+//! (one process under cargo test, nextest, and Bazel alike) so the
+//! environment they pin cannot leak into, or be poisoned by, other tests.
+
+mod support;
+
+use std::sync::Once;
+use std::sync::atomic::Ordering;
+use std::time::Duration;
+
+use support::{send_one, test_config};
+use xai_grok_sampler::SamplingClient;
+use xai_grok_test_support::spawn_counting_server;
+
+/// Pin the env these assertions depend on before any client is built, so
+/// ambient shell exports (`GROK_SAMPLER_SHARED_CLIENT=0`,
+/// `GROK_POOL_MAX_IDLE=0`) cannot flip the expected pooling behavior.
+fn pin_env() {
+    static PIN: Once = Once::new();
+    PIN.call_once(|| {
+        // Safety: runs before any test builds a client or reads these vars;
+        // racing tests block on the Once, and the crate latches the kill
+        // switch and pool knobs only at first client construction.
+        unsafe {
+            std::env::remove_var("GROK_SAMPLER_SHARED_CLIENT");
+            std::env::set_var("GROK_POOL_MAX_IDLE", "2");
+            std::env::set_var("GROK_POOL_IDLE_TIMEOUT_SECS", "90");
+        }
+    });
+}
+
+#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
+async fn two_sampling_clients_share_one_connection() {
+    pin_env();
+    let (base_url, accepts, _heads) = spawn_counting_server().await;
+    let a = SamplingClient::new(test_config(&base_url, "token-a")).unwrap();
+    let b = SamplingClient::new(test_config(&base_url, "token-b")).unwrap();
+    send_one(&a).await;
+    // Brief pause so the idle connection is checked back into the pool.
+    tokio::time::sleep(Duration::from_millis(50)).await;
+    send_one(&b).await;
+    assert_eq!(accepts.load(Ordering::SeqCst), 1);
+}
+
+#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
+async fn shared_client_keeps_per_config_headers_isolated() {
+    pin_env();
+    let (base_url, _accepts, heads) = spawn_counting_server().await;
+    let mut cfg_a = test_config(&base_url, "token-a");
+    cfg_a
+        .extra_headers
+        .insert("x-test-extra".to_string(), "isolated-a".to_string());
+    let mut cfg_b = test_config(&base_url, "token-b");
+    cfg_b
+        .extra_headers
+        .insert("x-test-extra".to_string(), "isolated-b".to_string());
+    let a = SamplingClient::new(cfg_a).unwrap();
+    let b = SamplingClient::new(cfg_b).unwrap();
+    send_one(&a).await;
+    send_one(&b).await;
+
+    let heads = heads.lock().unwrap();
+    assert_eq!(heads.len(), 2);
+    assert!(heads[0].contains("Bearer token-a") && heads[0].contains("isolated-a"));
+    assert!(!heads[0].contains("token-b") && !heads[0].contains("isolated-b"));
+    assert!(heads[1].contains("Bearer token-b") && heads[1].contains("isolated-b"));
+    assert!(!heads[1].contains("token-a") && !heads[1].contains("isolated-a"));
+}
+
+#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
+async fn shared_http1_fallback_never_pools() {
+    pin_env();
+    let (base_url, accepts, _heads) = spawn_counting_server().await;
+    let mut cfg = test_config(&base_url, "token-a");
+    cfg.force_http1 = true;
+    let client = SamplingClient::new(cfg).unwrap();
+    send_one(&client).await;
+    send_one(&client).await;
+    assert_eq!(accepts.load(Ordering::SeqCst), 2);
+}

crates/codegen/xai-grok-sampler/tests/support/mod.rs

@@ -0,0 +1,32 @@
+//! Sampler-specific helpers for the shared-HTTP-client integration binaries:
+//! config + request drivers for real `SamplingClient`s. The generic
+//! connection-counting server lives in `xai_grok_test_support`.
+
+use std::sync::Arc;
+
+use xai_grok_sampler::{SamplerConfig, SamplingClient};
+use xai_grok_sampling_types::{ContentPart, ConversationItem, ConversationRequest, UserItem};
+
+pub fn test_config(base_url: &str, api_key: &str) -> SamplerConfig {
+    SamplerConfig {
+        api_key: Some(api_key.to_string()),
+        base_url: base_url.to_string(),
+        model: "test-model".to_string(),
+        ..SamplerConfig::default()
+    }
+}
+
+/// Drive one POST through the client; the canned `{}` body is not a valid
+/// completion, but only the wire-level request matters here.
+pub async fn send_one(client: &SamplingClient) {
+    let request = ConversationRequest {
+        items: vec![ConversationItem::User(UserItem {
+            content: vec![ContentPart::Text {
+                text: Arc::<str>::from("hi"),
+            }],
+            ..Default::default()
+        })],
+        ..Default::default()
+    };
+    let _ = client.conversation(request).await;
+}

crates/codegen/xai-grok-sampler/tests/test_actor.rs

@@ -0,0 +1,951 @@
+//! Integration tests for the M4 actor + request_task layer.
+//!
+//! Tests are integration-style (in `tests/`) rather than unit tests
+//! because they require a real `tokio::runtime` and a mock HTTP
+//! server (axum) to talk to the `SamplingClient`. Happy-path SSE
+//! payloads come from `xai_grok_test_support::sse`.
+
+use std::net::SocketAddr;
+use std::sync::Arc;
+use std::sync::atomic::{AtomicU32, Ordering};
+use std::time::Duration;
+
+use axum::Router;
+use axum::http::StatusCode;
+use axum::response::sse::{Event, Sse};
+use axum::routing::post;
+use futures_util::stream::{self, StreamExt};
+use indexmap::IndexMap;
+use serde_json::json;
+use tokio::net::TcpListener;
+use tokio::sync::{mpsc, oneshot};
+
+use xai_grok_sampler::{
+    ApiBackend, RequestId, RetryPolicy, SamplerActor, SamplerConfig, SamplingChannel,
+    SamplingErrorKind, SamplingEvent,
+};
+use xai_grok_sampling_types::{
+    ConversationItem, ConversationRequest, DoomLoopRecoveryPolicy, UserItem,
+};
+use xai_grok_test_support::{SseEvent, sse};
+
+// ---------------------------------------------------------------------------
+// Mock server harness
+// ---------------------------------------------------------------------------
+
+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;
+        });
+        // Give the server a moment to start.
+        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(());
+    }
+}
+
+// ---------------------------------------------------------------------------
+// Config + request helpers
+// ---------------------------------------------------------------------------
+
+fn test_config(base_url: String, model: &str) -> SamplerConfig {
+    SamplerConfig {
+        api_key: Some("test-key".into()),
+        base_url,
+        model: model.into(),
+        max_completion_tokens: Some(1024),
+        temperature: None,
+        top_p: None,
+        api_backend: ApiBackend::ChatCompletions,
+        auth_scheme: Default::default(),
+        extra_headers: IndexMap::new(),
+        context_window: 128_000,
+        force_http1: false,
+        // Keep retries minimal so tests don't take forever.
+        max_retries: Some(2),
+        stream_tool_calls: false,
+        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,
+        compactions_remaining: None,
+        compaction_at_tokens: None,
+        doom_loop_recovery: None,
+        header_injector: None,
+    }
+}
+
+fn user_request(text: &str) -> ConversationRequest {
+    ConversationRequest {
+        items: vec![ConversationItem::User(UserItem {
+            content: vec![xai_grok_sampling_types::ContentPart::Text {
+                text: std::sync::Arc::<str>::from(text),
+            }],
+            synthetic_reason: None,
+            ..Default::default()
+        })],
+        ..Default::default()
+    }
+}
+
+// ---------------------------------------------------------------------------
+// SSE generators
+// ---------------------------------------------------------------------------
+
+/// Render test-helper [`SseEvent`]s (optional `event:` name + `data:`) as
+/// axum SSE events for this file's router-based harness.
+fn sse_events_to_axum(events: Vec<SseEvent>) -> Vec<Event> {
+    events
+        .into_iter()
+        .map(|e| {
+            let ev = Event::default().data(e.data);
+            match e.event {
+                Some(name) => ev.event(name),
+                None => ev,
+            }
+        })
+        .collect()
+}
+
+fn text_chunk_event(content: &str, finish: bool) -> Event {
+    let chunk = json!({
+        "id": "chatcmpl-test",
+        "object": "chat.completion.chunk",
+        "created": 0,
+        "model": "test-model",
+        "choices": [{
+            "index": 0,
+            "delta": { "role": "assistant", "content": content },
+            "finish_reason": if finish { json!("stop") } else { json!(null) }
+        }]
+    });
+    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() {
+    let (event_tx, _event_rx) = mpsc::unbounded_channel();
+    let cfg = test_config("http://127.0.0.1:0/v1".into(), "test-model");
+    let handle = SamplerActor::spawn(cfg, RetryPolicy::default(), event_tx);
+    assert_eq!(handle.active_count().await, 0);
+    handle.cancel(RequestId::from("nonexistent"));
+    // Re-querying should still be 0 (cancel of unknown id is no-op).
+    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() {
+    let app = Router::new().route(
+        "/v1/chat/completions",
+        post(|| async {
+            let events = sse::chat_completion_events("hello world", "test-model");
+            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 cfg = test_config(server.base_url(), "test-model");
+    let handle = SamplerActor::spawn(cfg, RetryPolicy::default(), event_tx);
+
+    let rid = RequestId::from("req-1");
+    handle.submit(rid.clone(), user_request("hi"));
+
+    let events = drain_until_terminal(&mut event_rx, Duration::from_secs(5)).await;
+    server.shutdown();
+
+    assert!(matches!(events[0], SamplingEvent::StreamStarted { .. }));
+    assert!(
+        events
+            .iter()
+            .any(|e| matches!(e, SamplingEvent::FirstToken { .. }))
+    );
+
+    let texts: Vec<&str> = events
+        .iter()
+        .filter_map(|e| match e {
+            SamplingEvent::ChannelToken {
+                channel: SamplingChannel::Text,
+                text,
+                ..
+            } => Some(text.as_str()),
+            _ => None,
+        })
+        .collect();
+    assert_eq!(texts.join(""), "hello world");
+
+    match events.last().unwrap() {
+        SamplingEvent::Completed {
+            request_id,
+            response,
+            ..
+        } => {
+            assert_eq!(request_id, &rid);
+            if let Some(a) = response.assistant() {
+                assert_eq!(a.content.as_ref(), "hello world");
+            } else {
+                panic!("expected Assistant message");
+            }
+        }
+        other => panic!("expected Completed, got {other:?}"),
+    }
+}
+
+// ---------------------------------------------------------------------------
+// submit_and_collect
+// ---------------------------------------------------------------------------
+
+#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
+async fn submit_and_collect_returns_response() {
+    let app = Router::new().route(
+        "/v1/chat/completions",
+        post(|| async {
+            let events = sse::chat_completion_events("collected response", "test-model");
+            Sse::new(stream::iter(
+                events.into_iter().map(Ok::<_, std::convert::Infallible>),
+            ))
+        }),
+    );
+    let server = MockServer::spawn(app).await;
+    let (event_tx, _event_rx) = mpsc::unbounded_channel();
+    let cfg = test_config(server.base_url(), "test-model");
+    let handle = SamplerActor::spawn(cfg, RetryPolicy::default(), event_tx);
+
+    let rid = RequestId::from("req-collect");
+    let result = handle
+        .submit_and_collect(rid, user_request("hi"))
+        .await
+        .expect("collected ok");
+    server.shutdown();
+
+    let (response, _metrics) = result;
+    let a = response.assistant().expect("assistant item present");
+    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() {
+    // Server that yields one chunk then hangs.
+    let app = Router::new().route(
+        "/v1/chat/completions",
+        post(|| async {
+            let stream = stream::iter(vec![Ok::<_, std::convert::Infallible>(text_chunk_event(
+                "starting", false,
+            ))])
+            .chain(stream::pending());
+            Sse::new(stream)
+        }),
+    );
+    let server = MockServer::spawn(app).await;
+    let (event_tx, mut event_rx) = mpsc::unbounded_channel();
+    let cfg = test_config(server.base_url(), "test-model");
+    let handle = SamplerActor::spawn(cfg, RetryPolicy::default(), event_tx);
+
+    let rid = RequestId::from("req-cancel");
+    handle.submit(rid.clone(), user_request("hi"));
+
+    // Wait for the first token to arrive so we know the request is in flight.
+    let _ = await_event_matching(
+        &mut event_rx,
+        |e| matches!(e, SamplingEvent::FirstToken { .. }),
+        Duration::from_secs(5),
+    )
+    .await
+    .expect("first token");
+
+    handle.cancel(rid.clone());
+
+    // Expect a Failed event with the cancellation message.
+    let failed = await_event_matching(
+        &mut event_rx,
+        |e| matches!(e, SamplingEvent::Failed { .. }),
+        Duration::from_secs(5),
+    )
+    .await
+    .expect("Failed event after cancel");
+
+    if let SamplingEvent::Failed { error, .. } = failed {
+        assert!(error.message.contains("cancelled"));
+    }
+
+    // Wait briefly for the task to clean up.
+    tokio::time::sleep(Duration::from_millis(200)).await;
+    assert_eq!(handle.active_count().await, 0);
+    server.shutdown();
+}
+
+// ---------------------------------------------------------------------------
+// Concurrent requests
+// ---------------------------------------------------------------------------
+
+#[tokio::test(flavor = "multi_thread", worker_threads = 4)]
+async fn two_concurrent_requests_complete_with_correct_request_ids() {
+    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);
+                let events = sse::chat_completion_events(&format!("response-{n}"), "test-model");
+                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 cfg = test_config(server.base_url(), "test-model");
+    let handle = SamplerActor::spawn(cfg, RetryPolicy::default(), event_tx);
+
+    let rid_a = RequestId::from("req-a");
+    let rid_b = RequestId::from("req-b");
+    handle.submit(rid_a.clone(), user_request("a"));
+    handle.submit(rid_b.clone(), user_request("b"));
+
+    // Drain until we see Completed for both.
+    let mut completed_a = false;
+    let mut completed_b = false;
+    let deadline = tokio::time::Instant::now() + Duration::from_secs(5);
+    while !(completed_a && completed_b) {
+        let now = tokio::time::Instant::now();
+        if now >= deadline {
+            panic!(
+                "timed out waiting for both requests to complete: a={completed_a}, b={completed_b}"
+            );
+        }
+        let remaining = deadline - now;
+        match tokio::time::timeout(remaining, event_rx.recv()).await {
+            Ok(Some(SamplingEvent::Completed { request_id, .. })) if request_id == rid_a => {
+                completed_a = true;
+            }
+            Ok(Some(SamplingEvent::Completed { request_id, .. })) if request_id == rid_b => {
+                completed_b = true;
+            }
+            Ok(Some(_)) => {}
+            Ok(None) => panic!("event channel closed"),
+            Err(_) => panic!("timeout"),
+        }
+    }
+    server.shutdown();
+}
+
+// ---------------------------------------------------------------------------
+// Retry on transient transport error
+// ---------------------------------------------------------------------------
+
+#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
+async fn retries_on_500_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 {
+                    // First attempt: server error.
+                    Err::<Sse<_>, (StatusCode, String)>((
+                        StatusCode::INTERNAL_SERVER_ERROR,
+                        json!({ "error": { "message": "transient" } }).to_string(),
+                    ))
+                } else {
+                    // Subsequent attempts: success.
+                    let events = sse::chat_completion_events("ok", "test-model");
+                    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();
+    // Lots of retries available; backoff is jittered around 2s on first
+    // retry, so this test takes a bit to run.
+    let cfg = test_config(server.base_url(), "test-model");
+    let handle = SamplerActor::spawn(cfg, RetryPolicy::default(), event_tx);
+
+    let rid = RequestId::from("req-retry");
+    handle.submit(rid.clone(), user_request("hi"));
+
+    let events = drain_until_terminal(&mut event_rx, Duration::from_secs(15)).await;
+    server.shutdown();
+
+    let saw_retrying = events
+        .iter()
+        .any(|e| matches!(e, SamplingEvent::Retrying { .. }));
+    assert!(saw_retrying, "expected at least one Retrying event");
+
+    match events.last().unwrap() {
+        SamplingEvent::Completed { response, .. } => {
+            if let Some(a) = response.assistant() {
+                assert_eq!(a.content.as_ref(), "ok");
+            }
+        }
+        other => panic!("expected Completed after retry, got {other:?}"),
+    }
+
+    assert!(
+        counter.load(Ordering::SeqCst) >= 2,
+        "server hit at least twice"
+    );
+}
+
+// ---------------------------------------------------------------------------
+// Rate limit exhausts threshold
+// ---------------------------------------------------------------------------
+
+#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
+async fn rate_limit_exhausts_at_threshold_and_yields_failed() {
+    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 {
+                counter.fetch_add(1, Ordering::SeqCst);
+                Err::<
+                    Sse<
+                        futures_util::stream::Iter<
+                            std::vec::IntoIter<Result<Event, std::convert::Infallible>>,
+                        >,
+                    >,
+                    (StatusCode, String),
+                >((
+                    StatusCode::TOO_MANY_REQUESTS,
+                    json!({ "error": { "message": "slow down" } }).to_string(),
+                ))
+            }
+        }),
+    );
+    let server = MockServer::spawn(app).await;
+    let (event_tx, mut event_rx) = mpsc::unbounded_channel();
+    let cfg = test_config(server.base_url(), "test-model");
+    let handle = SamplerActor::spawn(cfg, RetryPolicy::default(), event_tx);
+
+    let rid = RequestId::from("req-429");
+    handle.submit(rid.clone(), user_request("hi"));
+
+    let events = drain_until_terminal(&mut event_rx, Duration::from_secs(60)).await;
+    server.shutdown();
+
+    match events.last().unwrap() {
+        SamplingEvent::Failed { error, .. } => {
+            assert_eq!(error.kind, SamplingErrorKind::RateLimited);
+            assert_eq!(error.status_code, Some(429));
+        }
+        other => panic!("expected Failed(RateLimited), got {other:?}"),
+    }
+
+    let hits = counter.load(Ordering::SeqCst);
+    // RATE_LIMIT_RETRY_THRESHOLD = 2, so the actor stops after two
+    // attempts (the first attempt + one retry that also 429s = 2
+    // hits). Allow a small slack in case scheduling fires a third
+    // attempt before the threshold check.
+    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() {
+    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 {
+                counter.fetch_add(1, Ordering::SeqCst);
+                Err::<
+                    Sse<
+                        futures_util::stream::Iter<
+                            std::vec::IntoIter<Result<Event, std::convert::Infallible>>,
+                        >,
+                    >,
+                    (StatusCode, String),
+                >((StatusCode::UNAUTHORIZED, "unauthorized".to_string()))
+            }
+        }),
+    );
+    let server = MockServer::spawn(app).await;
+    let (event_tx, mut event_rx) = mpsc::unbounded_channel();
+    let cfg = test_config(server.base_url(), "test-model");
+    let handle = SamplerActor::spawn(cfg, RetryPolicy::default(), event_tx);
+
+    let rid = RequestId::from("req-auth");
+    handle.submit(rid.clone(), user_request("hi"));
+
+    let events = drain_until_terminal(&mut event_rx, Duration::from_secs(5)).await;
+    server.shutdown();
+
+    // Auth errors are session-owned -- `classify_error` returns
+    // `EmitToSession` so the actor emits Failed immediately without
+    // retrying.
+    assert!(
+        !events
+            .iter()
+            .any(|e| matches!(e, SamplingEvent::Retrying { .. }))
+    );
+    match events.last().unwrap() {
+        SamplingEvent::Failed { error, .. } => {
+            assert_eq!(error.kind, SamplingErrorKind::Auth);
+        }
+        other => panic!("expected Failed(Auth), got {other:?}"),
+    }
+    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");
+    cfg.api_backend = ApiBackend::Messages;
+    cfg
+}
+
+/// Regression for the refusal-stop_reason incident: a well-formed stream
+/// terminated by `stop_reason: "refusal"` must produce a successful
+/// completion from EXACTLY ONE request — no retry storm.
+#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
+async fn messages_refusal_stream_completes_with_single_request() {
+    let counter = Arc::new(AtomicU32::new(0));
+    let counter_handler = Arc::clone(&counter);
+    let app = Router::new().route(
+        "/v1/messages",
+        post(move || {
+            let counter = Arc::clone(&counter_handler);
+            async move {
+                counter.fetch_add(1, Ordering::SeqCst);
+                let events = sse::messages_api_events(
+                    "I can't help with that.",
+                    "messages-compatible-model",
+                    "refusal",
+                );
+                Sse::new(stream::iter(
+                    events.into_iter().map(Ok::<_, std::convert::Infallible>),
+                ))
+            }
+        }),
+    );
+    let server = MockServer::spawn(app).await;
+    let (event_tx, _event_rx) = mpsc::unbounded_channel();
+    let handle = SamplerActor::spawn(
+        messages_config(server.base_url()),
+        RetryPolicy::default(),
+        event_tx,
+    );
+
+    let result = handle
+        .submit_and_collect(RequestId::from("req-refusal"), user_request("hi"))
+        .await;
+    server.shutdown();
+
+    let (response, _metrics) = result.expect("refusal-terminated turn must complete");
+    let a = response.assistant().expect("assistant item present");
+    assert_eq!(a.content.as_ref(), "I can't help with that.");
+    assert_eq!(
+        counter.load(Ordering::SeqCst),
+        1,
+        "refusal must not trigger retries"
+    );
+}
+
+/// Empty-bodied refusal: `message_start → message_delta(refusal) →
+/// message_stop` with zero content blocks must complete from exactly one
+/// request — the content-less response must not be classified as a retryable
+/// EmptyResponse.
+#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
+async fn messages_empty_refusal_completes_without_retry() {
+    let counter = Arc::new(AtomicU32::new(0));
+    let counter_handler = Arc::clone(&counter);
+    let app = Router::new().route(
+        "/v1/messages",
+        post(move || {
+            let counter = Arc::clone(&counter_handler);
+            async move {
+                counter.fetch_add(1, Ordering::SeqCst);
+                let mut events =
+                    sse::messages_api_events("", "messages-compatible-model", "refusal");
+                // Drop the content block events; keep start/delta/stop only.
+                events.drain(1..4);
+                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(
+        messages_config(server.base_url()),
+        RetryPolicy::default(),
+        event_tx,
+    );
+
+    handle.submit(RequestId::from("req-empty-refusal"), user_request("hi"));
+    let events = drain_until_terminal(&mut event_rx, Duration::from_secs(10)).await;
+    server.shutdown();
+
+    assert!(
+        !events
+            .iter()
+            .any(|e| matches!(e, SamplingEvent::Retrying { .. })),
+        "content-less refusal must not be retried"
+    );
+    match events.last().unwrap() {
+        SamplingEvent::Completed { response, .. } => {
+            assert_eq!(
+                response.stop_reason,
+                Some(xai_grok_sampling_types::StopReason::ContentFilter)
+            );
+        }
+        other => panic!("expected Completed, got {other:?}"),
+    }
+    assert_eq!(counter.load(Ordering::SeqCst), 1, "exactly one request");
+}
+
+/// A mid-stream event that fails serde (after a valid `message_start`) is a
+/// deterministic response-parse failure: Fatal on the first attempt, surfaced
+/// as a non-retryable Serialization error — never a retry storm.
+#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
+async fn messages_unparseable_event_is_fatal_without_retry() {
+    let counter = Arc::new(AtomicU32::new(0));
+    let counter_handler = Arc::clone(&counter);
+    let app =
+        Router::new().route(
+            "/v1/messages",
+            post(move || {
+                let counter = Arc::clone(&counter_handler);
+                async move {
+                    counter.fetch_add(1, Ordering::SeqCst);
+                    let mut events =
+                        sse::messages_api_events("hello", "messages-compatible-model", "end_turn");
+                    // Replace the tail with a `message_delta` missing the
+                    // required `delta` field — fails MessageStreamEvent serde.
+                    events.truncate(4);
+                    events.push(Event::default().data(
+                        json!({"type":"message_delta","usage":{"output_tokens":1}}).to_string(),
+                    ));
+                    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(
+        messages_config(server.base_url()),
+        RetryPolicy::default(),
+        event_tx,
+    );
+
+    handle.submit(RequestId::from("req-bad-event"), user_request("hi"));
+    let events = drain_until_terminal(&mut event_rx, Duration::from_secs(10)).await;
+    server.shutdown();
+
+    assert!(
+        !events
+            .iter()
+            .any(|e| matches!(e, SamplingEvent::Retrying { .. })),
+        "serde failures must not be retried"
+    );
+    match events.last().unwrap() {
+        SamplingEvent::Failed { error, .. } => {
+            assert_eq!(error.kind, SamplingErrorKind::Serialization);
+            assert!(!error.is_retryable, "surfaced info must be non-retryable");
+        }
+        other => panic!("expected Failed(Serialization), got {other:?}"),
+    }
+    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() {
+    use std::sync::Mutex;
+
+    let captured_models: Arc<Mutex<Vec<String>>> = Arc::new(Mutex::new(Vec::new()));
+    let captured_handler = Arc::clone(&captured_models);
+    let app = Router::new().route(
+        "/v1/chat/completions",
+        post(move |axum::Json(body): axum::Json<serde_json::Value>| {
+            let captured = Arc::clone(&captured_handler);
+            async move {
+                let model = body
+                    .get("model")
+                    .and_then(|m| m.as_str())
+                    .unwrap_or("")
+                    .to_string();
+                captured.lock().unwrap().push(model);
+                let events = sse::chat_completion_events("ok", "test-model");
+                Sse::new(stream::iter(
+                    events.into_iter().map(Ok::<_, std::convert::Infallible>),
+                ))
+            }
+        }),
+    );
+    let server = MockServer::spawn(app).await;
+    let (event_tx, _event_rx) = mpsc::unbounded_channel();
+    let cfg = test_config(server.base_url(), "model-A");
+    let handle = SamplerActor::spawn(cfg, RetryPolicy::default(), event_tx);
+
+    let _ = handle
+        .submit_and_collect(RequestId::from("req-1"), user_request("hi"))
+        .await
+        .expect("first req ok");
+
+    let mut new_cfg = test_config(server.base_url(), "model-B");
+    new_cfg.api_key = Some("test-key".into());
+    handle.update_config(new_cfg);
+
+    let _ = handle
+        .submit_and_collect(RequestId::from("req-2"), user_request("hi"))
+        .await
+        .expect("second req ok");
+
+    server.shutdown();
+
+    let models = captured_models.lock().unwrap();
+    assert_eq!(
+        models.as_slice(),
+        &["model-A".to_string(), "model-B".to_string()]
+    );
+}
+
+// ---------------------------------------------------------------------------
+// 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");
+    cfg.api_backend = ApiBackend::Responses;
+    cfg.doom_loop_recovery = doom_loop;
+    cfg
+}
+
+/// Server-reported doom-loop triggers flow through the actor rung onto the
+/// completed response, without retries. The trigger is non-confident
+/// (`@response` channel), so the recovery — which resamples only confident
+/// signals — leaves it alone.
+#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
+async fn responses_doom_loop_signals_reach_completed_response() {
+    let counter = Arc::new(AtomicU32::new(0));
+    let counter_handler = Arc::clone(&counter);
+    let app = Router::new().route(
+        "/v1/responses",
+        post(move || {
+            let counter = Arc::clone(&counter_handler);
+            async move {
+                counter.fetch_add(1, Ordering::SeqCst);
+                let events = sse_events_to_axum(sse::responses_api_doom_loop_terminal_only_events(
+                    &["tail_repetition:4@response"],
+                    "some thought",
+                    "an answer",
+                    "test-model",
+                ));
+                Sse::new(stream::iter(
+                    events.into_iter().map(Ok::<_, std::convert::Infallible>),
+                ))
+            }
+        }),
+    );
+    let server = MockServer::spawn(app).await;
+    let (event_tx, _event_rx) = mpsc::unbounded_channel();
+    let handle = SamplerActor::spawn(
+        responses_config(server.base_url(), Some(DoomLoopRecoveryPolicy::default())),
+        RetryPolicy::default(),
+        event_tx,
+    );
+
+    let result = handle
+        .submit_and_collect(RequestId::from("req-doom-signal"), user_request("hi"))
+        .await;
+    server.shutdown();
+
+    let (response, _metrics) = result.expect("a signalled turn still completes");
+    assert_eq!(counter.load(Ordering::SeqCst), 1, "warn-only: no resample");
+    assert_eq!(response.doom_loop_signals.len(), 1);
+    assert_eq!(
+        response.doom_loop_signals[0].raw,
+        "tail_repetition:4@response"
+    );
+    assert_eq!(response.assistant_text(), "an answer");
+}
+
+/// Acceptance spec for the recovery rung: a confident signal
+/// (`tail_repetition:8@thinking` at the default threshold) is resampled once
+/// and the clean second response is accepted, on its own budget.
+#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
+async fn responses_confident_doom_loop_signal_resamples_once() {
+    let counter = Arc::new(AtomicU32::new(0));
+    let counter_handler = Arc::clone(&counter);
+    let app = Router::new().route(
+        "/v1/responses",
+        post(move || {
+            let counter = Arc::clone(&counter_handler);
+            async move {
+                let attempt = counter.fetch_add(1, Ordering::SeqCst);
+                let events = if attempt == 0 {
+                    sse::responses_api_doom_loop_terminal_only_events(
+                        &["tail_repetition:8@thinking"],
+                        "loop loop loop",
+                        "poisoned answer",
+                        "test-model",
+                    )
+                } else {
+                    sse::responses_api_reasoning_and_text_events(
+                        "fresh thought",
+                        "clean answer",
+                        "test-model",
+                    )
+                };
+                let events = sse_events_to_axum(events);
+                Sse::new(stream::iter(
+                    events.into_iter().map(Ok::<_, std::convert::Infallible>),
+                ))
+            }
+        }),
+    );
+    let server = MockServer::spawn(app).await;
+    let (event_tx, _event_rx) = mpsc::unbounded_channel();
+    let handle = SamplerActor::spawn(
+        responses_config(server.base_url(), Some(DoomLoopRecoveryPolicy::default())),
+        RetryPolicy::default(),
+        event_tx,
+    );
+
+    let result = handle
+        .submit_and_collect(RequestId::from("req-doom-resample"), user_request("hi"))
+        .await;
+    server.shutdown();
+
+    let (response, _metrics) = result.expect("recovery accepts the clean resample");
+    assert_eq!(counter.load(Ordering::SeqCst), 2, "exactly one resample");
+    assert_eq!(response.assistant_text(), "clean answer");
+    assert!(
+        response.doom_loop_signals.is_empty(),
+        "the accepted response is the clean resample"
+    );
+}
+
+// ---------------------------------------------------------------------------
+// Helpers for draining the event channel
+// ---------------------------------------------------------------------------
+
+/// Drain the event channel until a terminal event (`Completed` or
+/// `Failed`) is received, or until `deadline` elapses.
+async fn drain_until_terminal(
+    rx: &mut mpsc::UnboundedReceiver<SamplingEvent>,
+    timeout: Duration,
+) -> Vec<SamplingEvent> {
+    let mut out = Vec::new();
+    let start = tokio::time::Instant::now();
+    loop {
+        let elapsed = start.elapsed();
+        if elapsed >= timeout {
+            panic!(
+                "drain_until_terminal timed out after {:?}; got {} events",
+                timeout,
+                out.len()
+            );
+        }
+        let remaining = timeout - elapsed;
+        match tokio::time::timeout(remaining, rx.recv()).await {
+            Ok(Some(ev)) => {
+                let terminal = matches!(
+                    ev,
+                    SamplingEvent::Completed { .. } | SamplingEvent::Failed { .. }
+                );
+                out.push(ev);
+                if terminal {
+                    return out;
+                }
+            }
+            Ok(None) => panic!("event channel closed before terminal event"),
+            Err(_) => panic!(
+                "drain_until_terminal timed out after {:?}; got {} events",
+                timeout,
+                out.len()
+            ),
+        }
+    }
+}
+
+/// Wait for the next event matching `pred`, or return `None` on
+/// timeout.
+async fn await_event_matching(
+    rx: &mut mpsc::UnboundedReceiver<SamplingEvent>,
+    mut pred: impl FnMut(&SamplingEvent) -> bool,
+    timeout: Duration,
+) -> Option<SamplingEvent> {
+    let start = tokio::time::Instant::now();
+    loop {
+        let elapsed = start.elapsed();
+        if elapsed >= timeout {
+            return None;
+        }
+        let remaining = timeout - elapsed;
+        match tokio::time::timeout(remaining, rx.recv()).await {
+            Ok(Some(ev)) => {
+                if pred(&ev) {
+                    return Some(ev);
+                }
+            }
+            Ok(None) => return None,
+            Err(_) => return None,
+        }
+    }
+}