Synced from monorepo
Synced from monorepo Changes: - Temporarily disable session share link creation in the TUI - Do not approve plan on empty Enter from the revise prompt - Expose chat product Skills via ACP available_commands_update - Return immediately from a blocking wait on an already-completed ACP task - Split headless pager module for clearer structure - Stop git worktree prune from removing user registrations on resume - Use compaction sampler tokenizer for item token counts - Opt-in extra root CAs via GROK_EXTRA_CA_BUNDLE - Cancel all session subagents when the user stops - Let the session persistence actor exit when its session ends - Make fullscreen terminal resize much cheaper on long sessions - Report honestly from kill_task when an ACP task does not exist - Hide /usage for external-auth deployments - Forward the history-load trailer’s computer_reason to the client - Remove ineffective no-op tool reminder - Declare slash-command screen-mode support in one place - Keep settings enum picker on the committed value until Enter - Reap a PTY’s full process tree - Stream tool calls from headless mode over ACP - Bridge gateway task lifecycle to ACP for chat session background tasks - Don’t warn about truncated history on a suppressed replay - Fit full-replace summarizer input and recover on context-length errors - Stop dropping agents over an unrecognized frontmatter color - Add /undo as a slash alias for /rewind - Harden sleep/wake token-refresh paths against forced re-login - Add session/list ACP method - Give each sampling backend its own conversion module - Treat an unenrolled child process as a lint error - Suppress the cancelled marker on send-now wake turns - Stop tearing down Roslyn on every edit, and read C# diagnostics Source-Revision: 2a28b4a86cfc4a4c133c35b7fc2a6a9964387c39
This commit is contained in:
grokkybara[bot]
parent
500129c714
commit
dd04f397b1
367 changed files with 29489 additions and 10051 deletions
|
|
@ -7,6 +7,7 @@ description = "Actor-based sampling/inference layer for xAI grok (HTTP streaming
|
|||
|
||||
[dependencies]
|
||||
# Internal
|
||||
xai-grok-extra-ca = { workspace = true }
|
||||
xai-grok-sampling-types = { path = "../xai-grok-sampling-types" }
|
||||
xai-grok-version = { workspace = true }
|
||||
|
||||
|
|
|
|||
|
|
@ -68,18 +68,22 @@ impl SamplingConsumer {
|
|||
}
|
||||
}
|
||||
|
||||
/// 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.
|
||||
/// Bearer-fragment length shared with attribution callbacks: the **last**
|
||||
/// N chars (JWT heads are a shared constant; only the tail distinguishes
|
||||
/// tokens). Must stay in lock-step with `token_suffix` in
|
||||
/// `xai-grok-shell/src/auth/model.rs`, the comparison site.
|
||||
pub const SENT_BEARER_PREFIX_LEN: usize = 12;
|
||||
/// Last [`SENT_BEARER_PREFIX_LEN`] characters of a bearer, char-boundary
|
||||
/// safe (bearer strings are visible-ASCII per the header grammars, but a
|
||||
/// resolver-supplied `String` has no such guarantee -- counting chars from
|
||||
/// the end avoids a byte-index panic on non-ASCII input).
|
||||
pub(crate) fn bearer_tail_fragment(s: &str) -> &str {
|
||||
match s.char_indices().rev().nth(SENT_BEARER_PREFIX_LEN - 1) {
|
||||
Some((i, _)) => &s[i..],
|
||||
None => s,
|
||||
}
|
||||
}
|
||||
|
||||
/// Hook invoked by [`crate::SamplingClient`] at every 401 response site.
|
||||
///
|
||||
/// Implementations are responsible for joining `sent_bearer_prefix`
|
||||
|
|
@ -100,12 +104,13 @@ pub const SENT_BEARER_PREFIX_LEN: usize = 12;
|
|||
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
|
||||
/// `sent_bearer_prefix` is the **last
|
||||
/// [`SENT_BEARER_PREFIX_LEN`] characters** (the tail -- see the
|
||||
/// constant's doc for why the tail, not the head) 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
|
||||
/// Anthropic Messages API backends) and truncates it to that
|
||||
/// fragment **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
|
||||
|
|
@ -118,3 +123,22 @@ pub trait Auth401AttributionCallback: Send + Sync + std::fmt::Debug {
|
|||
|
||||
/// Shared, cheap-to-clone alias for the attribution callback.
|
||||
pub type SharedAttributionCallback = Arc<dyn Auth401AttributionCallback>;
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
#[test]
|
||||
fn bearer_tail_fragment_semantics() {
|
||||
// Tail, not head: JWT heads are a shared constant.
|
||||
assert_eq!(
|
||||
bearer_tail_fragment("eyJ0eXAiOiJh.shared-head.tail-distinct"),
|
||||
"ail-distinct"
|
||||
);
|
||||
assert_eq!(bearer_tail_fragment("abc"), "abc");
|
||||
assert_eq!(bearer_tail_fragment(""), "");
|
||||
assert_eq!(bearer_tail_fragment("123456789012"), "123456789012");
|
||||
// 13 multi-byte chars: a byte-index cut would land mid-char.
|
||||
assert_eq!(bearer_tail_fragment("ééééééééééééé"), "éééééééééééé");
|
||||
}
|
||||
}
|
||||
|
|
|
|||
|
|
@ -29,6 +29,7 @@ use xai_grok_sampling_types::{
|
|||
rs,
|
||||
};
|
||||
|
||||
use crate::attribution::bearer_tail_fragment;
|
||||
use crate::config::{AuthScheme, OriginClientInfo, SamplerConfig};
|
||||
|
||||
// Re-export ApiBackend from the shared types crate for downstream callers.
|
||||
|
|
@ -657,10 +658,15 @@ impl SamplingClient {
|
|||
self.defaults.api_backend.clone()
|
||||
}
|
||||
|
||||
/// POST with default headers. When a bearer_resolver is wired it is the
|
||||
/// sole auth source: a missing live bearer strips default Authorization /
|
||||
/// x-api-key so a hard-expired seed key cannot ride on the wire.
|
||||
fn post(&self, url: impl reqwest::IntoUrl) -> reqwest::RequestBuilder {
|
||||
/// POST with default headers, returning the builder plus the tail
|
||||
/// fragment of the credential actually placed in its headers (`None` =
|
||||
/// no credential) — captured at build time because a record-time
|
||||
/// re-read races with the recovery a 401 triggers.
|
||||
///
|
||||
/// A wired bearer_resolver is the sole auth source: a missing live
|
||||
/// bearer strips default Authorization / x-api-key so a hard-expired
|
||||
/// seed key cannot ride on the wire.
|
||||
fn post(&self, url: impl reqwest::IntoUrl) -> (reqwest::RequestBuilder, Option<String>) {
|
||||
let mut headers = self.default_headers.clone();
|
||||
if let Some(resolver) = &self.bearer_resolver {
|
||||
headers.remove(AUTHORIZATION);
|
||||
|
|
@ -703,55 +709,42 @@ impl SamplingClient {
|
|||
x_api_key_prefix = x_api_key_prefix.as_deref().unwrap_or("none"),
|
||||
);
|
||||
}
|
||||
let sent_bearer = Self::sent_fragment_from_headers(&headers, &self.defaults.auth_scheme);
|
||||
if let Some(injector) = &self.header_injector {
|
||||
injector.inject(&mut headers);
|
||||
}
|
||||
self.http.post(url).headers(headers)
|
||||
(self.http.post(url).headers(headers), sent_bearer)
|
||||
}
|
||||
|
||||
/// Bearer prefix for 401 attribution. When a resolver is wired it is
|
||||
/// authoritative (including `None` ⇒ nothing was sent). Without a resolver,
|
||||
/// fall back to construction-time default headers.
|
||||
/// Tail fragment of the credential in `headers` — `x-api-key`
|
||||
/// (Messages-API scheme) or `Authorization` — per
|
||||
/// [`crate::attribution::SENT_BEARER_PREFIX_LEN`].
|
||||
fn sent_fragment_from_headers(headers: &HeaderMap, scheme: &AuthScheme) -> Option<String> {
|
||||
let raw = match scheme {
|
||||
AuthScheme::XApiKey => headers
|
||||
.get(HeaderName::from_static("x-api-key"))
|
||||
.and_then(|v| v.to_str().ok()),
|
||||
AuthScheme::Bearer => headers
|
||||
.get(AUTHORIZATION)
|
||||
.and_then(|v| v.to_str().ok())
|
||||
.and_then(|s| s.strip_prefix("Bearer ")),
|
||||
};
|
||||
raw.map(|s| bearer_tail_fragment(s).to_string())
|
||||
}
|
||||
|
||||
/// Best-effort *build-time* view of what the next request would carry
|
||||
/// (resolver-authoritative). For request-start diagnostics
|
||||
/// ([`Self::auth_info`]) only — 401 attribution must use the fragment
|
||||
/// captured by [`Self::post`] instead, which cannot race a recovery.
|
||||
fn current_sent_bearer_prefix(&self) -> Option<String> {
|
||||
if self.bearer_resolver.is_some() {
|
||||
return self
|
||||
.bearer_resolver
|
||||
.as_ref()
|
||||
.and_then(|r| r.current_bearer())
|
||||
.map(|mut s| {
|
||||
s.truncate(crate::attribution::SENT_BEARER_PREFIX_LEN.min(s.len()));
|
||||
s
|
||||
});
|
||||
.map(|s| bearer_tail_fragment(&s).to_string());
|
||||
}
|
||||
self.extract_sent_bearer()
|
||||
}
|
||||
|
||||
/// Extract the bearer from `default_headers`, truncated to prefix length.
|
||||
/// Reads `x-api-key` (Anthropic Messages API) or `Authorization` (OpenAI-completions).
|
||||
fn extract_sent_bearer(&self) -> Option<String> {
|
||||
let raw = match self.defaults.auth_scheme {
|
||||
AuthScheme::XApiKey => self
|
||||
.default_headers
|
||||
.get(HeaderName::from_static("x-api-key"))
|
||||
.and_then(|v| v.to_str().ok())
|
||||
.map(|s| s.to_string()),
|
||||
AuthScheme::Bearer => self
|
||||
.default_headers
|
||||
.get(AUTHORIZATION)
|
||||
.and_then(|v| v.to_str().ok())
|
||||
.and_then(|s| s.strip_prefix("Bearer "))
|
||||
.map(|s| s.to_string()),
|
||||
};
|
||||
raw.map(|mut s| {
|
||||
// Truncate in-place so we never materialize a heap-resident
|
||||
// copy of the full bearer outside the local stack of this
|
||||
// function. `String::truncate` operates on byte indices and
|
||||
// panics on a non-char-boundary cut; bearer tokens are
|
||||
// ASCII (per the `Authorization` and `x-api-key` header
|
||||
// grammars) so the byte index is always safe.
|
||||
s.truncate(crate::attribution::SENT_BEARER_PREFIX_LEN.min(s.len()));
|
||||
s
|
||||
})
|
||||
Self::sent_fragment_from_headers(&self.default_headers, &self.defaults.auth_scheme)
|
||||
}
|
||||
|
||||
/// Invoke the optional 401 attribution callback for one logical
|
||||
|
|
@ -761,15 +754,16 @@ impl SamplingClient {
|
|||
/// that saw the status, so higher layers that react to a 401 must
|
||||
/// not emit a duplicate event.
|
||||
///
|
||||
/// The bearer passed to the callback is already truncated to
|
||||
/// [`crate::attribution::SENT_BEARER_PREFIX_LEN`] characters by
|
||||
/// [`Self::extract_sent_bearer`]; the trait contract guarantees
|
||||
/// that callers downstream of this crate never see the full
|
||||
/// bearer.
|
||||
fn record_401_attribution(&self, consumer: crate::attribution::SamplingConsumer) {
|
||||
/// `sent_prefix` is the fragment [`Self::post`] captured for the
|
||||
/// rejected request (already tail-truncated; the full bearer never
|
||||
/// crosses this boundary).
|
||||
fn record_401_attribution(
|
||||
&self,
|
||||
consumer: crate::attribution::SamplingConsumer,
|
||||
sent_prefix: Option<&str>,
|
||||
) {
|
||||
if let Some(cb) = self.attribution_callback.as_ref() {
|
||||
let sent_prefix = self.current_sent_bearer_prefix();
|
||||
cb.record_401(consumer, sent_prefix.as_deref());
|
||||
cb.record_401(consumer, sent_prefix);
|
||||
}
|
||||
}
|
||||
|
||||
|
|
@ -842,7 +836,13 @@ impl SamplingClient {
|
|||
Ok(request)
|
||||
}
|
||||
|
||||
async fn handle_response(&self, response: reqwest::Response) -> Result<ChatCompletionResponse> {
|
||||
/// `sent_bearer` is the fragment [`Self::post`] captured for the
|
||||
/// request that produced `response` (401 attribution).
|
||||
async fn handle_response(
|
||||
&self,
|
||||
response: reqwest::Response,
|
||||
sent_bearer: Option<&str>,
|
||||
) -> Result<ChatCompletionResponse> {
|
||||
let status = response.status();
|
||||
let model_metadata = extract_model_metadata(response.headers());
|
||||
let retry_after_secs = extract_retry_after(response.headers());
|
||||
|
|
@ -851,7 +851,10 @@ impl SamplingClient {
|
|||
|
||||
if !status.is_success() {
|
||||
if status == reqwest::StatusCode::UNAUTHORIZED {
|
||||
self.record_401_attribution(crate::attribution::SamplingConsumer::ChatCompletions);
|
||||
self.record_401_attribution(
|
||||
crate::attribution::SamplingConsumer::ChatCompletions,
|
||||
sent_bearer,
|
||||
);
|
||||
let server_message = user_facing_api_error_message(status, bytes.as_ref());
|
||||
return Err(SamplingError::Auth(format!(
|
||||
"Unauthorized (401): {server_message}"
|
||||
|
|
@ -908,9 +911,8 @@ impl SamplingClient {
|
|||
deployment_id: payload.x_grok_deployment_id.as_deref(),
|
||||
user_id: payload.x_grok_user_id.as_deref(),
|
||||
};
|
||||
let http_request = grok_headers
|
||||
.apply(self.post(self.endpoint("chat/completions")))
|
||||
.json(&payload);
|
||||
let (builder, sent_bearer) = self.post(self.endpoint("chat/completions"));
|
||||
let http_request = grok_headers.apply(builder).json(&payload);
|
||||
|
||||
let response = http_request.send().await.map_err(|e| {
|
||||
// Log at debug level; errors are surfaced to the caller.
|
||||
|
|
@ -918,7 +920,7 @@ impl SamplingClient {
|
|||
e
|
||||
})?;
|
||||
|
||||
self.handle_response(response).await
|
||||
self.handle_response(response, sent_bearer.as_deref()).await
|
||||
}
|
||||
|
||||
/// Start a streaming chat completion request. Returns a stream of typed chunks.
|
||||
|
|
@ -966,8 +968,9 @@ impl SamplingClient {
|
|||
deployment_id: payload.x_grok_deployment_id.as_deref(),
|
||||
user_id: payload.x_grok_user_id.as_deref(),
|
||||
};
|
||||
let (builder, sent_bearer) = self.post(self.endpoint("chat/completions"));
|
||||
let http_request = grok_headers
|
||||
.apply(self.post(self.endpoint("chat/completions")))
|
||||
.apply(builder)
|
||||
.header(ACCEPT, HeaderValue::from_static("text/event-stream"))
|
||||
.json(&streaming_request);
|
||||
|
||||
|
|
@ -1001,6 +1004,7 @@ impl SamplingClient {
|
|||
span.record("error", "unauthorized (401)");
|
||||
self.record_401_attribution(
|
||||
crate::attribution::SamplingConsumer::ChatCompletionsStream,
|
||||
sent_bearer.as_deref(),
|
||||
);
|
||||
let endpoint = self.endpoint("chat/completions");
|
||||
let body = response.bytes().await.unwrap_or_default();
|
||||
|
|
@ -1179,9 +1183,8 @@ impl SamplingClient {
|
|||
// it in post-serialize. This is the last surviving piece of the
|
||||
// old raw_output machinery.
|
||||
xai_grok_sampling_types::patch_reasoning_text_types(&mut request_body);
|
||||
let http_request = grok_headers
|
||||
.apply(self.post(self.endpoint("responses")))
|
||||
.json(&request_body);
|
||||
let (builder, sent_bearer) = self.post(self.endpoint("responses"));
|
||||
let http_request = grok_headers.apply(builder).json(&request_body);
|
||||
|
||||
let response = http_request.send().await.map_err(|e| {
|
||||
tracing::debug!("HTTP request failed: {}", e);
|
||||
|
|
@ -1196,7 +1199,10 @@ impl SamplingClient {
|
|||
|
||||
if !status.is_success() {
|
||||
if status == reqwest::StatusCode::UNAUTHORIZED {
|
||||
self.record_401_attribution(crate::attribution::SamplingConsumer::Responses);
|
||||
self.record_401_attribution(
|
||||
crate::attribution::SamplingConsumer::Responses,
|
||||
sent_bearer.as_deref(),
|
||||
);
|
||||
let endpoint = self.endpoint("responses");
|
||||
let server_message = user_facing_api_error_message(status, bytes.as_ref());
|
||||
return Err(SamplingError::Auth(format!(
|
||||
|
|
@ -1321,8 +1327,9 @@ impl SamplingClient {
|
|||
.defaults
|
||||
.doom_loop_recovery
|
||||
.map(crate::doom_loop::DoomLoopSignalCollector::new);
|
||||
let (builder, sent_bearer) = self.post(self.endpoint("responses"));
|
||||
let mut http_request = grok_headers
|
||||
.apply(self.post(self.endpoint("responses")))
|
||||
.apply(builder)
|
||||
.header(ACCEPT, HeaderValue::from_static("text/event-stream"));
|
||||
if doom_loop.is_some() {
|
||||
// Presence opts in; the server ignores the value.
|
||||
|
|
@ -1355,7 +1362,10 @@ impl SamplingClient {
|
|||
if !status.is_success() {
|
||||
if status == reqwest::StatusCode::UNAUTHORIZED {
|
||||
span.record("error", "unauthorized (401)");
|
||||
self.record_401_attribution(crate::attribution::SamplingConsumer::ResponsesStream);
|
||||
self.record_401_attribution(
|
||||
crate::attribution::SamplingConsumer::ResponsesStream,
|
||||
sent_bearer.as_deref(),
|
||||
);
|
||||
let endpoint = self.endpoint("responses");
|
||||
let body = response.bytes().await.unwrap_or_default();
|
||||
let server_message = user_facing_api_error_message(status, body.as_ref());
|
||||
|
|
@ -1518,9 +1528,8 @@ impl SamplingClient {
|
|||
deployment_id: request.x_grok_deployment_id.as_deref(),
|
||||
user_id: request.x_grok_user_id.as_deref(),
|
||||
};
|
||||
let http_request = grok_headers
|
||||
.apply(self.post(self.endpoint("messages")))
|
||||
.json(&request.inner);
|
||||
let (builder, sent_bearer) = self.post(self.endpoint("messages"));
|
||||
let http_request = grok_headers.apply(builder).json(&request.inner);
|
||||
|
||||
let response = http_request.send().await.map_err(|e| {
|
||||
tracing::debug!("HTTP request failed: {}", e);
|
||||
|
|
@ -1535,7 +1544,10 @@ impl SamplingClient {
|
|||
|
||||
if !status.is_success() {
|
||||
if status == reqwest::StatusCode::UNAUTHORIZED {
|
||||
self.record_401_attribution(crate::attribution::SamplingConsumer::Messages);
|
||||
self.record_401_attribution(
|
||||
crate::attribution::SamplingConsumer::Messages,
|
||||
sent_bearer.as_deref(),
|
||||
);
|
||||
let endpoint = self.endpoint("messages");
|
||||
let server_message = user_facing_api_error_message(status, bytes.as_ref());
|
||||
return Err(SamplingError::Auth(format!(
|
||||
|
|
@ -1625,8 +1637,9 @@ impl SamplingClient {
|
|||
deployment_id: request.x_grok_deployment_id.as_deref(),
|
||||
user_id: request.x_grok_user_id.as_deref(),
|
||||
};
|
||||
let (builder, sent_bearer) = self.post(self.endpoint("messages"));
|
||||
let http_request = grok_headers
|
||||
.apply(self.post(self.endpoint("messages")))
|
||||
.apply(builder)
|
||||
.header(ACCEPT, HeaderValue::from_static("text/event-stream"))
|
||||
.json(&request.inner);
|
||||
|
||||
|
|
@ -1655,7 +1668,10 @@ impl SamplingClient {
|
|||
if !status.is_success() {
|
||||
if status == reqwest::StatusCode::UNAUTHORIZED {
|
||||
span.record("error", "unauthorized (401)");
|
||||
self.record_401_attribution(crate::attribution::SamplingConsumer::MessagesStream);
|
||||
self.record_401_attribution(
|
||||
crate::attribution::SamplingConsumer::MessagesStream,
|
||||
sent_bearer.as_deref(),
|
||||
);
|
||||
let endpoint = self.endpoint("messages");
|
||||
let body = response.bytes().await.unwrap_or_default();
|
||||
let server_message = user_facing_api_error_message(status, body.as_ref());
|
||||
|
|
@ -2307,10 +2323,8 @@ mod tests {
|
|||
let mut config = minimal_config();
|
||||
config.header_injector = Some(std::sync::Arc::new(TestInjector));
|
||||
let client = SamplingClient::new(config).expect("build");
|
||||
let req = client
|
||||
.post("http://localhost/test")
|
||||
.build()
|
||||
.expect("build request");
|
||||
let (builder, _sent) = client.post("http://localhost/test");
|
||||
let req = builder.build().expect("build request");
|
||||
assert!(
|
||||
req.headers().contains_key("traceparent"),
|
||||
"HeaderInjector should inject traceparent into post() requests"
|
||||
|
|
@ -2379,31 +2393,28 @@ mod tests {
|
|||
}
|
||||
}
|
||||
|
||||
/// `extract_sent_bearer` strips the `"Bearer "` prefix off
|
||||
/// `Authorization` for OpenAI-completions backends and truncates the
|
||||
/// remaining bearer to the cross-crate prefix length.
|
||||
/// `post()` strips the `"Bearer "` scheme prefix off `Authorization`
|
||||
/// and captures the tail fragment (see `SENT_BEARER_PREFIX_LEN`).
|
||||
#[test]
|
||||
fn extract_sent_bearer_strips_bearer_prefix_for_openai_compat() {
|
||||
fn post_captures_bearer_tail_for_openai_compat() {
|
||||
let cfg = SamplerConfig {
|
||||
api_key: Some("test-bearer-1234567890".to_string()),
|
||||
api_backend: ApiBackend::ChatCompletions,
|
||||
..minimal_config()
|
||||
};
|
||||
let client = SamplingClient::new(cfg).expect("client should build");
|
||||
let bearer = client.extract_sent_bearer();
|
||||
// Bearer is truncated at the crate boundary -- callers
|
||||
// downstream of this method only ever see the prefix.
|
||||
assert_eq!(bearer.as_deref(), Some("test-bearer-"));
|
||||
let (_builder, bearer) = client.post("https://example.test/v1/chat/completions");
|
||||
assert_eq!(bearer.as_deref(), Some("r-1234567890"));
|
||||
assert_eq!(
|
||||
bearer.as_deref().map(str::len),
|
||||
Some(crate::attribution::SENT_BEARER_PREFIX_LEN),
|
||||
);
|
||||
}
|
||||
|
||||
/// `extract_sent_bearer` reads `x-api-key` for Anthropic Messages API
|
||||
/// and truncates the value to the cross-crate prefix length.
|
||||
/// `post()` captures `x-api-key` for Messages-API backends and keeps
|
||||
/// the value's tail fragment.
|
||||
#[test]
|
||||
fn extract_sent_bearer_reads_x_api_key_for_messages() {
|
||||
fn post_captures_x_api_key_tail_for_messages() {
|
||||
let cfg = SamplerConfig {
|
||||
api_key: Some("anthropic-key-abc123".to_string()),
|
||||
api_backend: ApiBackend::Messages,
|
||||
|
|
@ -2411,24 +2422,68 @@ mod tests {
|
|||
..minimal_config()
|
||||
};
|
||||
let client = SamplingClient::new(cfg).expect("client should build");
|
||||
let bearer = client.extract_sent_bearer();
|
||||
assert_eq!(bearer.as_deref(), Some("anthropic-ke"));
|
||||
let (_builder, bearer) = client.post("https://example.test/v1/messages");
|
||||
assert_eq!(bearer.as_deref(), Some("c-key-abc123"));
|
||||
assert_eq!(
|
||||
bearer.as_deref().map(str::len),
|
||||
Some(crate::attribution::SENT_BEARER_PREFIX_LEN),
|
||||
);
|
||||
}
|
||||
|
||||
/// `extract_sent_bearer` returns `None` when no auth header is set.
|
||||
/// `post()` captures `None` when the request carries no auth header.
|
||||
#[test]
|
||||
fn extract_sent_bearer_returns_none_when_no_header() {
|
||||
fn post_captures_none_when_no_header() {
|
||||
let cfg = SamplerConfig {
|
||||
api_key: None,
|
||||
api_backend: ApiBackend::ChatCompletions,
|
||||
..minimal_config()
|
||||
};
|
||||
let client = SamplingClient::new(cfg).expect("client should build");
|
||||
assert!(client.extract_sent_bearer().is_none());
|
||||
let (_builder, bearer) = client.post("https://example.test/v1/chat/completions");
|
||||
assert!(bearer.is_none());
|
||||
}
|
||||
|
||||
/// The race this design closes: a 401 triggers a recovery that rotates
|
||||
/// the resolver, so a record-time re-read attributes a bearer the
|
||||
/// rejected request never carried. The attributed fragment must be the
|
||||
/// one captured when the request was built.
|
||||
#[test]
|
||||
fn post_capture_is_immune_to_resolver_rotation_after_build() {
|
||||
#[derive(Debug)]
|
||||
struct RotatingResolver(std::sync::Mutex<String>);
|
||||
impl crate::config::BearerResolver for RotatingResolver {
|
||||
fn current_bearer(&self) -> Option<String> {
|
||||
Some(self.0.lock().unwrap().clone())
|
||||
}
|
||||
}
|
||||
|
||||
let resolver = std::sync::Arc::new(RotatingResolver(std::sync::Mutex::new(
|
||||
"rejected-token-oldtail1".to_string(),
|
||||
)));
|
||||
let cfg = SamplerConfig {
|
||||
api_key: None,
|
||||
api_backend: ApiBackend::Responses,
|
||||
bearer_resolver: Some(resolver.clone()),
|
||||
..minimal_config()
|
||||
};
|
||||
let client = SamplingClient::new(cfg).expect("client should build");
|
||||
|
||||
let (_builder, sent_at_build) = client.post("https://example.test/v1/responses");
|
||||
// The 401 kicks recovery; the resolver rotates before the callback runs.
|
||||
*resolver.0.lock().unwrap() = "fresh-token-newtail99".to_string();
|
||||
|
||||
assert_eq!(
|
||||
sent_at_build.as_deref(),
|
||||
Some("ken-oldtail1"),
|
||||
"attribution must describe the bearer the rejected request carried"
|
||||
);
|
||||
// A record-time re-read (the pre-fix behavior) would report the
|
||||
// rotated token instead:
|
||||
assert_eq!(
|
||||
client.current_sent_bearer_prefix().as_deref(),
|
||||
Some("en-newtail99"),
|
||||
"sanity: the build-time capture and a live re-read now differ"
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
|
|
@ -2441,10 +2496,8 @@ mod tests {
|
|||
..minimal_config()
|
||||
};
|
||||
let client = SamplingClient::new(cfg).expect("client should build");
|
||||
let request = client
|
||||
.post("https://example.test/v1/messages")
|
||||
.build()
|
||||
.expect("request should build");
|
||||
let (builder, _sent) = client.post("https://example.test/v1/messages");
|
||||
let request = builder.build().expect("request should build");
|
||||
let auth = request
|
||||
.headers()
|
||||
.get(AUTHORIZATION)
|
||||
|
|
@ -2469,10 +2522,8 @@ mod tests {
|
|||
..minimal_config()
|
||||
};
|
||||
let client = SamplingClient::new(cfg).expect("client should build");
|
||||
let request = client
|
||||
.post("https://example.test/v1/responses")
|
||||
.build()
|
||||
.expect("request should build");
|
||||
let (builder, _sent) = client.post("https://example.test/v1/responses");
|
||||
let request = builder.build().expect("request should build");
|
||||
let auth_count = request.headers().get_all(AUTHORIZATION).iter().count();
|
||||
assert_eq!(
|
||||
auth_count, 1,
|
||||
|
|
@ -2497,10 +2548,8 @@ mod tests {
|
|||
..minimal_config()
|
||||
};
|
||||
let client = SamplingClient::new(cfg).expect("client should build");
|
||||
let request = client
|
||||
.post("https://example.test/v1/messages")
|
||||
.build()
|
||||
.expect("request should build");
|
||||
let (builder, _sent) = client.post("https://example.test/v1/messages");
|
||||
let request = builder.build().expect("request should build");
|
||||
let api_key = request
|
||||
.headers()
|
||||
.get("x-api-key")
|
||||
|
|
@ -2509,27 +2558,10 @@ mod tests {
|
|||
assert!(request.headers().get(AUTHORIZATION).is_none());
|
||||
}
|
||||
|
||||
/// Bearers shorter than the prefix length pass through unchanged.
|
||||
/// Defensive against the truncation logic inadvertently widening
|
||||
/// short bearers (no panics, no zero-padding).
|
||||
/// The callback receives the `post()`-captured fragment only — the
|
||||
/// full bearer never crosses the crate boundary.
|
||||
#[test]
|
||||
fn extract_sent_bearer_short_bearer_passes_through_unchanged() {
|
||||
let cfg = SamplerConfig {
|
||||
api_key: Some("abc".to_string()),
|
||||
api_backend: ApiBackend::ChatCompletions,
|
||||
..minimal_config()
|
||||
};
|
||||
let client = SamplingClient::new(cfg).expect("client should build");
|
||||
assert_eq!(client.extract_sent_bearer().as_deref(), Some("abc"));
|
||||
}
|
||||
|
||||
/// `record_401_attribution` invokes the wired callback with the
|
||||
/// expected `consumer` and the truncated bearer prefix that the
|
||||
/// wire would carry. The key assertion is that the callback
|
||||
/// receives the prefix only -- the full bearer never crosses the
|
||||
/// crate boundary.
|
||||
#[test]
|
||||
fn record_401_attribution_invokes_callback_with_extracted_bearer() {
|
||||
fn record_401_attribution_invokes_callback_with_captured_bearer() {
|
||||
let cb = std::sync::Arc::new(CountingCallback::default());
|
||||
let cb_dyn: crate::attribution::SharedAttributionCallback = cb.clone();
|
||||
let cfg = SamplerConfig {
|
||||
|
|
@ -2540,16 +2572,18 @@ mod tests {
|
|||
..minimal_config()
|
||||
};
|
||||
let client = SamplingClient::new(cfg).expect("client should build");
|
||||
client.record_401_attribution(crate::attribution::SamplingConsumer::ChatCompletionsStream);
|
||||
let (_builder, sent_bearer) = client.post("https://example.test/v1/chat/completions");
|
||||
client.record_401_attribution(
|
||||
crate::attribution::SamplingConsumer::ChatCompletionsStream,
|
||||
sent_bearer.as_deref(),
|
||||
);
|
||||
let calls = cb.invocations.lock().unwrap();
|
||||
assert_eq!(calls.len(), 1);
|
||||
assert_eq!(
|
||||
calls[0].0,
|
||||
crate::attribution::SamplingConsumer::ChatCompletionsStream
|
||||
);
|
||||
// Prefix-only -- the `extra-tail` portion of the bearer is
|
||||
// dropped by `extract_sent_bearer` before the callback fires.
|
||||
assert_eq!(calls[0].1.as_deref(), Some("the-bearer-1"));
|
||||
assert_eq!(calls[0].1.as_deref(), Some("0-extra-tail"));
|
||||
assert_eq!(
|
||||
calls[0].1.as_deref().map(str::len),
|
||||
Some(crate::attribution::SENT_BEARER_PREFIX_LEN),
|
||||
|
|
@ -2602,11 +2636,9 @@ mod tests {
|
|||
..minimal_config()
|
||||
};
|
||||
let client = SamplingClient::new(cfg).expect("client should build");
|
||||
let request = client
|
||||
.post("https://example.test/v1/responses")
|
||||
.body("")
|
||||
.build()
|
||||
.expect("request should build");
|
||||
let (builder, sent) = client.post("https://example.test/v1/responses");
|
||||
let request = builder.body("").build().expect("request should build");
|
||||
assert_eq!(sent, None, "capture must agree: nothing was sent");
|
||||
assert!(
|
||||
request.headers().get(AUTHORIZATION).is_none(),
|
||||
"stale default Authorization must not be sent when resolver is empty"
|
||||
|
|
@ -2638,7 +2670,7 @@ mod tests {
|
|||
let client = SamplingClient::new(cfg).expect("client should build");
|
||||
|
||||
// Build a request to inspect the final headers.
|
||||
let builder = client.post("https://example.test/v1/responses");
|
||||
let (builder, _sent) = client.post("https://example.test/v1/responses");
|
||||
let request = builder.body("").build().expect("request should build");
|
||||
|
||||
let auth_values: Vec<_> = request.headers().get_all(AUTHORIZATION).iter().collect();
|
||||
|
|
@ -2671,7 +2703,10 @@ mod tests {
|
|||
};
|
||||
let client = SamplingClient::new(cfg).expect("client should build");
|
||||
// Must not panic.
|
||||
client.record_401_attribution(crate::attribution::SamplingConsumer::ChatCompletions);
|
||||
client.record_401_attribution(
|
||||
crate::attribution::SamplingConsumer::ChatCompletions,
|
||||
Some("bearer-tail-12"),
|
||||
);
|
||||
}
|
||||
|
||||
/// `response.completed` carrying
|
||||
|
|
|
|||
|
|
@ -56,6 +56,35 @@ pub enum SamplingEvent {
|
|||
arguments_delta: Option<String>,
|
||||
},
|
||||
|
||||
/// The provider opened a response (Messages `message_start`). Carries the
|
||||
/// real message id, model, and input-side token counts exactly as they
|
||||
/// arrive on the wire, before any content. Surfaced in order so partial-mode
|
||||
/// consumers can emit the real `message_start` id/usage instead of a
|
||||
/// synthesized placeholder. Emitted by the Messages L2 transform only; the
|
||||
/// Responses/Chat transforms lack these fields at stream open and emit
|
||||
/// nothing here.
|
||||
///
|
||||
/// `input_tokens` is the uncached prompt portion; the Anthropic Messages API
|
||||
/// reports cache hits and writes in the separate `cache_read_input_tokens`
|
||||
/// and `cache_creation_input_tokens` buckets, both known at `message_start`.
|
||||
ResponseStarted {
|
||||
request_id: RequestId,
|
||||
message_id: String,
|
||||
model: String,
|
||||
input_tokens: u64,
|
||||
cache_read_input_tokens: u64,
|
||||
cache_creation_input_tokens: u64,
|
||||
},
|
||||
|
||||
/// The reasoning (thinking) block finished and its encrypted signature is
|
||||
/// known (Messages thinking `content_block_stop`). Surfaced in order so
|
||||
/// partial-mode consumers can emit `signature_delta` before the thinking
|
||||
/// block's `content_block_stop`. Emitted by the Messages L2 transform only.
|
||||
ReasoningCompleted {
|
||||
request_id: RequestId,
|
||||
signature: String,
|
||||
},
|
||||
|
||||
/// Streaming completed successfully.
|
||||
Completed {
|
||||
request_id: RequestId,
|
||||
|
|
|
|||
|
|
@ -11,7 +11,7 @@
|
|||
//! 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.
|
||||
//! environment. Extra roots: `GROK_EXTRA_CA_BUNDLE` via `xai_grok_extra_ca`.
|
||||
|
||||
use std::sync::OnceLock;
|
||||
use std::time::Duration;
|
||||
|
|
@ -83,16 +83,18 @@ fn build_http_client() -> Result<reqwest::Client, reqwest::Error> {
|
|||
.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()
|
||||
xai_grok_extra_ca::with_extra_root_certificates(
|
||||
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.
|
||||
|
|
@ -103,13 +105,15 @@ fn build_http_client_http1() -> Result<reqwest::Client, reqwest::Error> {
|
|||
.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()
|
||||
xai_grok_extra_ca::with_extra_root_certificates(
|
||||
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)]
|
||||
|
|
|
|||
|
|
@ -292,6 +292,9 @@ pub fn stream_chat_completions<'a>(
|
|||
message_chunks_emitted: message_chunk_count,
|
||||
doom_loop_signals: Vec::new(),
|
||||
stop_message: None,
|
||||
message_id: None,
|
||||
raw_stop_reason: None,
|
||||
stop_sequence: None,
|
||||
};
|
||||
|
||||
yield SamplingEvent::Completed {
|
||||
|
|
|
|||
|
|
@ -171,6 +171,9 @@ mod tests {
|
|||
message_chunks_emitted: 1,
|
||||
doom_loop_signals: Vec::new(),
|
||||
stop_message: None,
|
||||
message_id: None,
|
||||
raw_stop_reason: None,
|
||||
stop_sequence: None,
|
||||
}),
|
||||
metrics: InferenceLatencyStats::default(),
|
||||
};
|
||||
|
|
|
|||
|
|
@ -19,6 +19,22 @@ use crate::events::{SamplingChannel, SamplingErrorInfo, SamplingEvent};
|
|||
use crate::metrics::InferenceLatencyStats;
|
||||
use crate::types::RequestId;
|
||||
|
||||
/// The verbatim wire string for a Messages API stop reason, before it collapses
|
||||
/// into the internal [`StopReason`]. Uses the enum's serde `snake_case`
|
||||
/// renaming so it cannot drift from the wire contract.
|
||||
fn messages_stop_reason_wire(sr: &messages::StopReason) -> String {
|
||||
match serde_json::to_value(sr) {
|
||||
Ok(serde_json::Value::String(s)) => s,
|
||||
other => {
|
||||
debug_assert!(
|
||||
false,
|
||||
"StopReason must serialize to a string, got {other:?}"
|
||||
);
|
||||
"end_turn".to_string()
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// 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 {
|
||||
|
|
@ -99,6 +115,12 @@ pub fn stream_messages<'a>(
|
|||
let mut final_output_tokens: u32 = 0;
|
||||
let mut final_stop_reason: Option<StopReason> = None;
|
||||
let mut final_stop_message: Option<String> = None;
|
||||
let mut final_message_id: Option<String> = None;
|
||||
let mut final_raw_stop_reason: Option<String> = None;
|
||||
// The provider's matched stop sequence (Messages `message_delta.stop_sequence`),
|
||||
// set only on a `stop_sequence`-terminated turn; carried through so the
|
||||
// headless `streaming-messages-json` consumer can echo it.
|
||||
let mut final_stop_sequence: Option<String> = None;
|
||||
|
||||
// Assistant-response accumulators (built up as ContentBlockStop
|
||||
// events fire). Reasoning is collected into a synthesized
|
||||
|
|
@ -152,10 +174,26 @@ pub fn stream_messages<'a>(
|
|||
|
||||
match event {
|
||||
MessageStreamEvent::MessageStart { message } => {
|
||||
final_message_id = Some(message.id.clone());
|
||||
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;
|
||||
// Surface the real id/model/input-usage in order, before any
|
||||
// content, so partial-mode framing emits them on the real
|
||||
// `message_start` instead of a synthesized placeholder.
|
||||
yield SamplingEvent::ResponseStarted {
|
||||
request_id: request_id.clone(),
|
||||
message_id: message.id,
|
||||
model: message.model,
|
||||
input_tokens: u64::from(message.usage.input_tokens),
|
||||
cache_read_input_tokens: u64::from(
|
||||
message.usage.cache_read_input_tokens,
|
||||
),
|
||||
cache_creation_input_tokens: u64::from(
|
||||
message.usage.cache_creation_input_tokens,
|
||||
),
|
||||
};
|
||||
}
|
||||
|
||||
MessageStreamEvent::ContentBlockStart {
|
||||
|
|
@ -237,7 +275,19 @@ pub fn stream_messages<'a>(
|
|||
arguments_delta: None,
|
||||
};
|
||||
}
|
||||
_ => {} // Image / ToolResult are not expected in assistant streams.
|
||||
// Encrypted reasoning the model chose to redact. Deliberately
|
||||
// parse-only: the `RedactedThinking` wire variant exists so a
|
||||
// stream that includes one deserializes instead of failing the
|
||||
// whole event parse and discarding an already-streamed
|
||||
// response, but its opaque `data` blob is not surfaced as a
|
||||
// `SamplingEvent` — forwarding it to the headless reducer's
|
||||
// `redacted_thinking` block would need a new event threaded
|
||||
// through the deferred sampler→shell→reducer hop and handled by
|
||||
// every `SamplingEvent` consumer (TUI included), so it is not
|
||||
// wired. No consumer claims redacted_thinking support.
|
||||
ContentBlock::RedactedThinking { .. } => {}
|
||||
// Image / ToolResult are not expected in assistant streams.
|
||||
_ => {}
|
||||
},
|
||||
|
||||
MessageStreamEvent::ContentBlockDelta { index, delta } => {
|
||||
|
|
@ -312,6 +362,15 @@ pub fn stream_messages<'a>(
|
|||
}
|
||||
}
|
||||
BlockType::Thinking => {
|
||||
// Surface the encrypted signature in order (at the
|
||||
// thinking block's stop) so partial-mode framing can
|
||||
// emit `signature_delta` before its `content_block_stop`.
|
||||
if !state.signature.is_empty() {
|
||||
yield SamplingEvent::ReasoningCompleted {
|
||||
request_id: request_id.clone(),
|
||||
signature: state.signature.clone(),
|
||||
};
|
||||
}
|
||||
if !state.thinking_acc.is_empty() || !state.signature.is_empty() {
|
||||
// Anthropic Messages API `Thinking` blocks uniquely
|
||||
// carry an encrypted `signature` distinct
|
||||
|
|
@ -359,6 +418,14 @@ pub fn stream_messages<'a>(
|
|||
if let Some(details) = delta.stop_details {
|
||||
final_stop_message = details.explanation;
|
||||
}
|
||||
// Keep the exact wire string so consumers can echo it.
|
||||
final_raw_stop_reason =
|
||||
delta.stop_reason.as_ref().map(messages_stop_reason_wire);
|
||||
// The matched stop sequence rides the same terminal delta
|
||||
// (present only on a `stop_sequence` stop); carry it verbatim.
|
||||
if delta.stop_sequence.is_some() {
|
||||
final_stop_sequence = delta.stop_sequence.clone();
|
||||
}
|
||||
final_stop_reason = delta.stop_reason.map(|sr| match sr {
|
||||
messages::StopReason::EndTurn => StopReason::Stop,
|
||||
messages::StopReason::MaxTokens => StopReason::Length,
|
||||
|
|
@ -470,6 +537,7 @@ pub fn stream_messages<'a>(
|
|||
total_tokens: total_prompt_tokens.saturating_add(final_output_tokens),
|
||||
reasoning_tokens: 0,
|
||||
cached_prompt_tokens: final_cache_read_input_tokens,
|
||||
cache_creation_prompt_tokens: final_cache_creation_input_tokens,
|
||||
})
|
||||
} else {
|
||||
None
|
||||
|
|
@ -511,6 +579,9 @@ pub fn stream_messages<'a>(
|
|||
message_chunks_emitted: message_chunk_count,
|
||||
doom_loop_signals: Vec::new(),
|
||||
stop_message: final_stop_message,
|
||||
message_id: final_message_id,
|
||||
raw_stop_reason: final_raw_stop_reason,
|
||||
stop_sequence: final_stop_sequence,
|
||||
};
|
||||
|
||||
yield SamplingEvent::Completed {
|
||||
|
|
|
|||
|
|
@ -58,6 +58,7 @@ fn message_delta_with_stop(stop: messages::StopReason) -> MessageStreamEvent {
|
|||
MessageStreamEvent::MessageDelta {
|
||||
delta: MessageDeltaBody {
|
||||
stop_reason: Some(stop),
|
||||
stop_sequence: None,
|
||||
stop_details: None,
|
||||
},
|
||||
usage: MessageDeltaUsage {
|
||||
|
|
@ -75,6 +76,7 @@ fn message_delta_refusal_with_explanation(explanation: &str) -> MessageStreamEve
|
|||
MessageStreamEvent::MessageDelta {
|
||||
delta: MessageDeltaBody {
|
||||
stop_reason: Some(messages::StopReason::Refusal),
|
||||
stop_sequence: None,
|
||||
stop_details: Some(messages::StopDetails {
|
||||
r#type: Some("refusal".to_string()),
|
||||
category: Some("frontier_llm".to_string()),
|
||||
|
|
@ -141,6 +143,10 @@ async fn text_block_assembles_into_completed_response() {
|
|||
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));
|
||||
// Provider message id and the verbatim wire stop reason survive
|
||||
// onto the response (collapsed `stop_reason` loses the string).
|
||||
assert_eq!(response.message_id.as_deref(), Some("msg_1"));
|
||||
assert_eq!(response.raw_stop_reason.as_deref(), Some("end_turn"));
|
||||
let u = response.usage.as_ref().expect("usage extracted");
|
||||
assert_eq!(u.prompt_tokens, 10);
|
||||
assert_eq!(u.completion_tokens, 5);
|
||||
|
|
@ -208,6 +214,61 @@ async fn thinking_block_emits_reasoning_channel_and_preserved_in_response() {
|
|||
}
|
||||
}
|
||||
|
||||
/// `thinking(sig1) → text → thinking(sig2)` must surface each thinking block's
|
||||
/// OWN signature, in order, on its own `ReasoningCompleted` (emitted at that
|
||||
/// block's stop) — so the per-index signature reaches the headless reducer and
|
||||
/// each block keeps its own signature rather than collapsing to one.
|
||||
#[tokio::test]
|
||||
async fn multiple_thinking_blocks_emit_per_block_signatures_in_order() {
|
||||
let thinking_block = |index: u32, text: &str, sig: &str| {
|
||||
vec![
|
||||
Ok(MessageStreamEvent::ContentBlockStart {
|
||||
index,
|
||||
content_block: ContentBlock::Thinking {
|
||||
thinking: String::new(),
|
||||
signature: String::new(),
|
||||
},
|
||||
}),
|
||||
Ok(MessageStreamEvent::ContentBlockDelta {
|
||||
index,
|
||||
delta: StreamDelta::ThinkingDelta {
|
||||
thinking: text.into(),
|
||||
},
|
||||
}),
|
||||
Ok(MessageStreamEvent::ContentBlockDelta {
|
||||
index,
|
||||
delta: StreamDelta::SignatureDelta {
|
||||
signature: sig.into(),
|
||||
},
|
||||
}),
|
||||
Ok(block_stop(index)),
|
||||
]
|
||||
};
|
||||
let mut events: Vec<Result<MessageStreamEvent, SamplingError>> = vec![Ok(message_start())];
|
||||
events.extend(thinking_block(0, "first", "sig-1"));
|
||||
events.push(Ok(text_block_start(1)));
|
||||
events.push(Ok(text_delta(1, "interlude")));
|
||||
events.push(Ok(block_stop(1)));
|
||||
events.extend(thinking_block(2, "second", "sig-2"));
|
||||
events.push(Ok(MessageStreamEvent::MessageStop));
|
||||
|
||||
let raw = stream::iter(events).boxed();
|
||||
let evs = collect(stream_messages(raw, None, rid(), Duration::from_secs(60))).await;
|
||||
|
||||
let sigs: Vec<&str> = evs
|
||||
.iter()
|
||||
.filter_map(|e| match e {
|
||||
SamplingEvent::ReasoningCompleted { signature, .. } => Some(signature.as_str()),
|
||||
_ => None,
|
||||
})
|
||||
.collect();
|
||||
assert_eq!(
|
||||
sigs,
|
||||
vec!["sig-1", "sig-2"],
|
||||
"each thinking block emits its own signature in order"
|
||||
);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn tool_use_block_assembles_into_tool_call() {
|
||||
let tool_start = MessageStreamEvent::ContentBlockStart {
|
||||
|
|
@ -574,6 +635,7 @@ fn message_delta_with_cache(
|
|||
MessageStreamEvent::MessageDelta {
|
||||
delta: MessageDeltaBody {
|
||||
stop_reason: Some(messages::StopReason::EndTurn),
|
||||
stop_sequence: None,
|
||||
stop_details: None,
|
||||
},
|
||||
usage: MessageDeltaUsage {
|
||||
|
|
@ -621,6 +683,7 @@ async fn prompt_tokens_sums_all_three_anthropic_buckets() {
|
|||
|
||||
assert_eq!(usage.prompt_tokens, 100 + 5000 + 200);
|
||||
assert_eq!(usage.cached_prompt_tokens, 5000);
|
||||
assert_eq!(usage.cache_creation_prompt_tokens, 200);
|
||||
assert_eq!(usage.completion_tokens, 7);
|
||||
assert_eq!(usage.total_tokens, 100 + 5000 + 200 + 7);
|
||||
}
|
||||
|
|
@ -638,6 +701,7 @@ async fn message_delta_cache_fields_override_message_start() {
|
|||
|
||||
assert_eq!(usage.prompt_tokens, 10 + 900 + 50);
|
||||
assert_eq!(usage.cached_prompt_tokens, 900);
|
||||
assert_eq!(usage.cache_creation_prompt_tokens, 50);
|
||||
assert_eq!(usage.completion_tokens, 4);
|
||||
}
|
||||
|
||||
|
|
|
|||
|
|
@ -381,6 +381,25 @@ pub(crate) fn stream_responses_tracked<'a>(
|
|||
ResponseStreamEvent::ResponseWebSearchCallCompleted(_)
|
||||
| ResponseStreamEvent::ResponseWebSearchCallSearching(_) => {}
|
||||
|
||||
// Code interpreter (server-side, like web/x search). Surface it
|
||||
// the same way x_search is: a generic backend tool call that the
|
||||
// shell renders as a client `tool_use` + `user` `tool_result`
|
||||
// split (grok has no HostedTool::CodeInterpreter, so these events
|
||||
// are latent under the current hosted-tool set). The started
|
||||
// event fires on InProgress; the full payload (code + outputs)
|
||||
// rides ResponseOutputItemDone(CodeInterpreterCall) below.
|
||||
ResponseStreamEvent::ResponseCodeInterpreterCallInProgress(ev) => {
|
||||
yield SamplingEvent::BackendToolCallStarted {
|
||||
request_id: request_id.clone(),
|
||||
call_id: ev.item_id.clone(),
|
||||
name: "code_interpreter".to_string(),
|
||||
};
|
||||
}
|
||||
// Interpreting/Completed carry no payload — the result arrives
|
||||
// via ResponseOutputItemDone(CodeInterpreterCall) below.
|
||||
ResponseStreamEvent::ResponseCodeInterpreterCallInterpreting(_)
|
||||
| ResponseStreamEvent::ResponseCodeInterpreterCallCompleted(_) => {}
|
||||
|
||||
// OutputItemDone carries the full result for backend tools.
|
||||
// For WebSearchCall this includes the query and source URLs.
|
||||
// For CustomToolCall this includes x_search results.
|
||||
|
|
@ -409,6 +428,19 @@ pub(crate) fn stream_responses_tracked<'a>(
|
|||
result,
|
||||
};
|
||||
}
|
||||
// Code interpreter: the full call (code + outputs) rides
|
||||
// the done item. Surfaced under the shared "code_interpreter"
|
||||
// name (matching the Started event); the shell renders it via
|
||||
// the client `tool_use` + `user` `tool_result` split.
|
||||
rs::OutputItem::CodeInterpreterCall(ci) => {
|
||||
let result = serde_json::to_value(ci).ok();
|
||||
yield SamplingEvent::BackendToolCallCompleted {
|
||||
request_id: request_id.clone(),
|
||||
call_id: ci.id.clone(),
|
||||
name: "code_interpreter".to_string(),
|
||||
result,
|
||||
};
|
||||
}
|
||||
_ => {}
|
||||
}
|
||||
}
|
||||
|
|
@ -485,6 +517,7 @@ pub(crate) fn stream_responses_tracked<'a>(
|
|||
total_tokens: u.total_tokens,
|
||||
reasoning_tokens: u.output_tokens_details.reasoning_tokens,
|
||||
cached_prompt_tokens: u.input_tokens_details.cached_tokens,
|
||||
cache_creation_prompt_tokens: 0,
|
||||
});
|
||||
|
||||
let cost_usd_ticks = response
|
||||
|
|
@ -543,6 +576,9 @@ pub(crate) fn stream_responses_tracked<'a>(
|
|||
message_chunks_emitted: message_chunk_count,
|
||||
doom_loop_signals,
|
||||
stop_message: None, // not reported on the Responses API
|
||||
message_id: None, // no provider message id on the Responses API
|
||||
raw_stop_reason: None,
|
||||
stop_sequence: None,
|
||||
};
|
||||
|
||||
yield SamplingEvent::Completed {
|
||||
|
|
@ -874,6 +910,67 @@ mod tests {
|
|||
assert!(output_observed.load(Ordering::Relaxed));
|
||||
}
|
||||
|
||||
/// A server-side code-interpreter run surfaces as a generic backend tool
|
||||
/// call (started on InProgress, completed on OutputItemDone) named
|
||||
/// "code_interpreter" — the same shape as x_search — so it is no longer
|
||||
/// silently dropped from the event stream.
|
||||
#[tokio::test]
|
||||
async fn code_interpreter_forwards_backend_tool_call() {
|
||||
let in_progress = rs::ResponseStreamEvent::ResponseCodeInterpreterCallInProgress(
|
||||
rs_types::ResponseCodeInterpreterCallInProgressEvent {
|
||||
sequence_number: 0,
|
||||
output_index: 0,
|
||||
item_id: "ci-1".into(),
|
||||
},
|
||||
);
|
||||
let done = rs::ResponseStreamEvent::ResponseOutputItemDone(
|
||||
rs_types::ResponseOutputItemDoneEvent {
|
||||
sequence_number: 1,
|
||||
output_index: 0,
|
||||
item: rs_types::OutputItem::CodeInterpreterCall(
|
||||
rs_types::CodeInterpreterToolCall {
|
||||
code: Some("print(1)".into()),
|
||||
container_id: "cont-1".into(),
|
||||
id: "ci-1".into(),
|
||||
outputs: None,
|
||||
status: rs_types::CodeInterpreterToolCallStatus::Completed,
|
||||
},
|
||||
),
|
||||
},
|
||||
);
|
||||
let raw = stream::iter(vec![Ok(in_progress), Ok(done), Ok(completed_event())]).boxed();
|
||||
let events = collect(stream_responses(
|
||||
raw,
|
||||
None,
|
||||
rid(),
|
||||
Duration::from_secs(60),
|
||||
None,
|
||||
))
|
||||
.await;
|
||||
|
||||
assert!(
|
||||
events.iter().any(|e| matches!(
|
||||
e,
|
||||
SamplingEvent::BackendToolCallStarted { call_id, name, .. }
|
||||
if call_id == "ci-1" && name == "code_interpreter"
|
||||
)),
|
||||
"expected a code_interpreter BackendToolCallStarted, got {events:?}"
|
||||
);
|
||||
let completed = events.iter().find_map(|e| match e {
|
||||
SamplingEvent::BackendToolCallCompleted {
|
||||
call_id,
|
||||
name,
|
||||
result,
|
||||
..
|
||||
} if name == "code_interpreter" => Some((call_id.clone(), result.clone())),
|
||||
_ => None,
|
||||
});
|
||||
let (call_id, result) = completed.expect("a code_interpreter BackendToolCallCompleted");
|
||||
assert_eq!(call_id, "ci-1");
|
||||
let result = result.expect("serialized code-interpreter payload");
|
||||
assert_eq!(result["code"], "print(1)");
|
||||
}
|
||||
|
||||
fn function_call_added_event(
|
||||
output_index: u32,
|
||||
call_id: &str,
|
||||
|
|
|
|||
Loading…
Reference in a new issue