Publish harness and TUI open-source
initial sync from the monorepo
This commit is contained in:
commit
c68e39f604
2734 changed files with 1437016 additions and 0 deletions
46
crates/codegen/xai-grok-telemetry/src/appender.rs
Normal file
46
crates/codegen/xai-grok-telemetry/src/appender.rs
Normal file
|
|
@ -0,0 +1,46 @@
|
|||
//! Shared non-blocking file appender + worker-guard registry for telemetry file-log layers.
|
||||
|
||||
use std::path::Path;
|
||||
use std::sync::{Mutex, OnceLock};
|
||||
|
||||
use tracing_appender::non_blocking::{NonBlocking, WorkerGuard};
|
||||
|
||||
// Park every worker guard for process lifetime; dropping a guard flushes and
|
||||
// shuts down that file's writer thread, so accumulate (never overwrite) to let
|
||||
// multiple file-log layers coexist.
|
||||
static FILE_LOG_GUARDS: OnceLock<Mutex<Vec<WorkerGuard>>> = OnceLock::new();
|
||||
|
||||
/// Shared non-blocking file writer for telemetry file-log layers. Opens `path`
|
||||
/// in append mode and parks the worker guard for process lifetime so buffered
|
||||
/// logs aren't lost. Sibling loggers (hooks/memory/sampling/instrumentation) can
|
||||
/// migrate onto this in a follow-up.
|
||||
pub(crate) fn non_blocking_file_writer(path: &Path) -> std::io::Result<NonBlocking> {
|
||||
if let Some(parent) = path.parent() {
|
||||
let _ = std::fs::create_dir_all(parent);
|
||||
}
|
||||
|
||||
let file = std::fs::OpenOptions::new()
|
||||
.create(true)
|
||||
.append(true)
|
||||
.open(path)?;
|
||||
|
||||
let (non_blocking, guard) = tracing_appender::non_blocking(file);
|
||||
let guards = FILE_LOG_GUARDS.get_or_init(|| Mutex::new(Vec::new()));
|
||||
// Recover from a poisoned mutex so the guard is always parked; dropping it
|
||||
// would shut down the writer thread and silently lose buffered logs.
|
||||
let mut guards = guards
|
||||
.lock()
|
||||
.unwrap_or_else(|poisoned| poisoned.into_inner());
|
||||
guards.push(guard);
|
||||
Ok(non_blocking)
|
||||
}
|
||||
|
||||
/// Drop all parked worker guards, flushing their non-blocking writers. Call at
|
||||
/// process exit so short-lived runs (e.g. headless `grok -p`) don't lose buffered logs.
|
||||
pub(crate) fn flush_file_log_guards() {
|
||||
if let Some(m) = FILE_LOG_GUARDS.get() {
|
||||
// Recover from a poisoned mutex so exit-flush still drains the guards.
|
||||
let mut guards = m.lock().unwrap_or_else(|poisoned| poisoned.into_inner());
|
||||
guards.clear(); // dropping each WorkerGuard flushes + joins its writer thread
|
||||
}
|
||||
}
|
||||
465
crates/codegen/xai-grok-telemetry/src/client.rs
Normal file
465
crates/codegen/xai-grok-telemetry/src/client.rs
Normal file
|
|
@ -0,0 +1,465 @@
|
|||
//! Core telemetry tracking — product events + Mixpanel.
|
||||
//!
|
||||
//! All calls route through [`track`]. Precedence: env > config > remote config > default.
|
||||
//!
|
||||
//! Extracted from `xai-grok-shell::agent::telemetry::track`. The HTTP client is
|
||||
//! injected via [`init`]/[`init_if_needed`] so this crate avoids depending on
|
||||
//! shell's `User-Agent` builder (which couples to the `permission` module).
|
||||
|
||||
use std::sync::{Arc, Mutex, OnceLock};
|
||||
|
||||
use chrono::{Local, SecondsFormat};
|
||||
use serde_json::json;
|
||||
use xai_mixpanel::Mixpanel;
|
||||
|
||||
use crate::config::{TelemetryConfig, TelemetryMode, deployment_id_from_key};
|
||||
use crate::http::OriginClientInfo;
|
||||
use crate::session_ctx::EmitterOrigin;
|
||||
|
||||
/// Event property map shared by all telemetry modules.
|
||||
pub type Metadata = serde_json::Map<String, serde_json::Value>;
|
||||
|
||||
/// Derive the analytics `event_value` from the full wire `event_name` by stripping
|
||||
/// whichever [`EmitterOrigin`] prefix it carries (`grok-shell-` /
|
||||
/// `grok-workspace-`). Unprefixed names pass through unchanged. Kept in
|
||||
/// lockstep with [`EmitterOrigin::event_prefix`] via [`EmitterOrigin::ALL`],
|
||||
/// so shell events keep their historical stripped value and workspace events
|
||||
/// collapse to the same bare suffix.
|
||||
fn event_value(event_name: &str) -> &str {
|
||||
for origin in EmitterOrigin::ALL {
|
||||
if let Some(suffix) = event_name.strip_prefix(origin.event_prefix()) {
|
||||
return suffix;
|
||||
}
|
||||
}
|
||||
event_name
|
||||
}
|
||||
|
||||
#[derive(Clone)]
|
||||
pub struct TelemetryClient {
|
||||
mode: TelemetryMode,
|
||||
events_url: Option<String>,
|
||||
events_api_key: Option<String>,
|
||||
mixpanel: Option<Arc<Mixpanel>>,
|
||||
user_id: Option<String>,
|
||||
team_id: Option<String>,
|
||||
deployment_id: Option<String>,
|
||||
shell_version: String,
|
||||
client_type: Option<String>,
|
||||
client_version: Option<String>,
|
||||
subscription_tier: Option<String>,
|
||||
http_client: reqwest::Client,
|
||||
}
|
||||
|
||||
impl std::fmt::Debug for TelemetryClient {
|
||||
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
|
||||
f.debug_struct("TelemetryClient")
|
||||
.field("events_url", &self.events_url)
|
||||
.field(
|
||||
"events_api_key",
|
||||
&self.events_api_key.as_ref().map(|_| "***"),
|
||||
)
|
||||
.field("mixpanel", &self.mixpanel.as_ref().map(|_| "configured"))
|
||||
.finish()
|
||||
}
|
||||
}
|
||||
|
||||
impl TelemetryClient {
|
||||
pub fn from_config(
|
||||
config: TelemetryConfig,
|
||||
mode: TelemetryMode,
|
||||
user_id: Option<String>,
|
||||
team_id: Option<String>,
|
||||
deployment_key: Option<String>,
|
||||
origin_client: Option<OriginClientInfo>,
|
||||
shell_version: String,
|
||||
subscription_tier: Option<String>,
|
||||
http_client: reqwest::Client,
|
||||
) -> Self {
|
||||
let mixpanel = if config.mixpanel_enabled {
|
||||
config
|
||||
.mixpanel_token
|
||||
.as_ref()
|
||||
.map(|token| Arc::new(Mixpanel::new(token.as_str())))
|
||||
} else {
|
||||
None
|
||||
};
|
||||
let deployment_id = deployment_key
|
||||
.filter(|s| !s.is_empty())
|
||||
.map(|k| deployment_id_from_key(&k));
|
||||
let (client_type, client_version) = match origin_client {
|
||||
Some(o) => (Some(o.product), o.version),
|
||||
None => (None, None),
|
||||
};
|
||||
|
||||
Self {
|
||||
mode,
|
||||
events_url: config.events_url,
|
||||
events_api_key: config.events_api_key,
|
||||
mixpanel,
|
||||
user_id,
|
||||
team_id,
|
||||
deployment_id,
|
||||
shell_version,
|
||||
client_type,
|
||||
client_version,
|
||||
subscription_tier: subscription_tier.map(|t| normalize_tier(&t)),
|
||||
http_client,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Normalize a subscription tier string to a consistent lowercase_underscore
|
||||
/// format for Mixpanel. Handles both CCP display names ("SuperGrok Heavy")
|
||||
/// and JWT-derived keys ("supergrok_heavy").
|
||||
fn normalize_tier(tier: &str) -> String {
|
||||
match tier {
|
||||
"SuperGrok Heavy" | "supergrok_heavy" => "supergrok_heavy",
|
||||
"SuperGrok" | "supergrok" => "supergrok",
|
||||
"SuperGrok Lite" | "supergrok_lite" => "supergrok_lite",
|
||||
"X Premium+" | "x_premium_plus" => "x_premium_plus",
|
||||
"X Premium" | "x_premium" => "x_premium",
|
||||
"X Basic" | "x_basic" => "x_basic",
|
||||
"Free" | "free" => "free",
|
||||
// Team / console API keys — dedicated Mixpanel segment, not free.
|
||||
"API Key" | "api_key" => "api_key",
|
||||
other => return other.to_ascii_lowercase().replace(' ', "_"),
|
||||
}
|
||||
.to_string()
|
||||
}
|
||||
|
||||
static TELEMETRY_CLIENT: OnceLock<Mutex<Option<TelemetryClient>>> = OnceLock::new();
|
||||
|
||||
/// Returns `true` when telemetry mode is `Enabled`.
|
||||
/// Used by `log_event` — product analytics events only fire in `Enabled` mode.
|
||||
pub fn is_enabled() -> bool {
|
||||
TELEMETRY_CLIENT
|
||||
.get()
|
||||
.and_then(|m| m.lock().ok())
|
||||
.is_some_and(|g| g.as_ref().is_some_and(|c| c.mode.is_enabled()))
|
||||
}
|
||||
|
||||
/// Returns `true` when telemetry mode is `Enabled` or `SessionMetrics`.
|
||||
/// Used by `session_metrics` — lifecycle events fire in both modes.
|
||||
pub fn is_session_metrics_enabled() -> bool {
|
||||
TELEMETRY_CLIENT
|
||||
.get()
|
||||
.and_then(|m| m.lock().ok())
|
||||
.is_some_and(|g| g.as_ref().is_some_and(|c| c.mode.session_metrics_enabled()))
|
||||
}
|
||||
|
||||
pub struct UserContext {
|
||||
pub country: String,
|
||||
pub language: String,
|
||||
pub timestamp: String,
|
||||
}
|
||||
|
||||
impl UserContext {
|
||||
pub fn collect() -> Self {
|
||||
let default_language = whoami::Language::En(whoami::Country::Any);
|
||||
let lang = whoami::langs()
|
||||
.ok()
|
||||
.and_then(|mut langs| langs.next())
|
||||
.unwrap_or(default_language);
|
||||
Self {
|
||||
country: lang.country().to_string(),
|
||||
language: lang.to_string(),
|
||||
timestamp: Local::now().to_rfc3339_opts(SecondsFormat::Millis, true),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Core telemetry emitter. Routes to product events + Mixpanel.
|
||||
pub async fn track(event_name: &str, request_id: &str, ctx: &UserContext, mut metadata: Metadata) {
|
||||
let lock = TELEMETRY_CLIENT.get_or_init(|| Mutex::new(None));
|
||||
let client = {
|
||||
let guard = lock.lock().unwrap_or_else(|err| err.into_inner());
|
||||
match guard.clone() {
|
||||
Some(c) => c,
|
||||
None => return,
|
||||
}
|
||||
};
|
||||
|
||||
let agent_id = crate::id::agent_id();
|
||||
let user_id = client.user_id.as_deref().unwrap_or(&agent_id);
|
||||
metadata.insert("agent_id".into(), json!(agent_id));
|
||||
if let Some(ref team_id) = client.team_id {
|
||||
metadata.insert("team_id".into(), json!(team_id));
|
||||
}
|
||||
if let Some(ref deployment_id) = client.deployment_id {
|
||||
metadata.insert("deployment_id".into(), json!(deployment_id));
|
||||
}
|
||||
metadata.insert("shell_version".into(), json!(client.shell_version));
|
||||
if let Some(ref client_type) = client.client_type {
|
||||
metadata.insert("client_type".into(), json!(client_type));
|
||||
}
|
||||
if let Some(ref client_version) = client.client_version {
|
||||
metadata.insert("client_version".into(), json!(client_version));
|
||||
}
|
||||
if let Some(ref subscription_tier) = client.subscription_tier {
|
||||
metadata.insert("subscription_tier".into(), json!(subscription_tier));
|
||||
}
|
||||
|
||||
// Product events path
|
||||
if let (Some(url), Some(api_key)) = (&client.events_url, &client.events_api_key) {
|
||||
let body = json!({
|
||||
"viewer_context": {
|
||||
"request_id": request_id,
|
||||
"user_attributes": {
|
||||
"user_id": user_id,
|
||||
"user_type": "LoggedIn",
|
||||
"country": ctx.country,
|
||||
"language": ctx.language,
|
||||
"locale": "English",
|
||||
},
|
||||
"device_attributes": {
|
||||
"app_name": "Grok Code",
|
||||
},
|
||||
},
|
||||
"api_key": api_key,
|
||||
"events": [{
|
||||
"event_name": event_name,
|
||||
"event_value": event_value(event_name),
|
||||
"event_metadata": metadata.clone(),
|
||||
"timestamp": ctx.timestamp,
|
||||
}]
|
||||
});
|
||||
let _ = client
|
||||
.http_client
|
||||
.post(url)
|
||||
.header("x-api-key", api_key.as_str())
|
||||
.timeout(std::time::Duration::from_secs(10))
|
||||
.json(&body)
|
||||
.send()
|
||||
.await;
|
||||
}
|
||||
|
||||
// Mixpanel path
|
||||
if let Some(ref mixpanel) = client.mixpanel {
|
||||
let time_secs = chrono::Utc::now().timestamp();
|
||||
let insert_id = format!("{event_name}:{request_id}:{time_secs}");
|
||||
|
||||
// Convert serde_json::Map to HashMap for mixpanel
|
||||
let mut props: std::collections::HashMap<String, serde_json::Value> =
|
||||
metadata.into_iter().collect();
|
||||
props.insert("distinct_id".into(), json!(user_id));
|
||||
props.insert("time".into(), json!(time_secs));
|
||||
props.insert("$insert_id".into(), json!(insert_id));
|
||||
props.insert("app_name".into(), json!("Grok Code"));
|
||||
props.insert("user_type".into(), json!("LoggedIn"));
|
||||
props.insert("country".into(), json!(ctx.country));
|
||||
props.insert("language".into(), json!(ctx.language));
|
||||
props.insert("locale".into(), json!("English"));
|
||||
|
||||
let _ = mixpanel.track(event_name, Some(props)).await;
|
||||
}
|
||||
}
|
||||
|
||||
/// Sync the user's Mixpanel profile once per init. Fire-and-forget.
|
||||
pub fn sync_profile() {
|
||||
let lock = TELEMETRY_CLIENT.get_or_init(|| Mutex::new(None));
|
||||
let client = {
|
||||
let guard = lock.lock().unwrap_or_else(|err| err.into_inner());
|
||||
match guard.clone() {
|
||||
Some(c) => c,
|
||||
None => return,
|
||||
}
|
||||
};
|
||||
|
||||
let Some(mixpanel) = client.mixpanel.clone() else {
|
||||
return;
|
||||
};
|
||||
|
||||
let agent_id = crate::id::agent_id();
|
||||
let user_id = client.user_id.as_deref().unwrap_or(&agent_id).to_owned();
|
||||
|
||||
tokio::spawn(async move {
|
||||
let mut props = std::collections::HashMap::new();
|
||||
props.insert("agent_id".into(), json!(agent_id));
|
||||
props.insert("shell_version".into(), json!(client.shell_version));
|
||||
props.insert("app_name".into(), json!("Grok Code"));
|
||||
if let Some(ref client_type) = client.client_type {
|
||||
props.insert("client_type".into(), json!(client_type));
|
||||
}
|
||||
if let Some(ref client_version) = client.client_version {
|
||||
props.insert("client_version".into(), json!(client_version));
|
||||
}
|
||||
if let Some(ref deployment_id) = client.deployment_id {
|
||||
props.insert("deployment_id".into(), json!(deployment_id));
|
||||
}
|
||||
if let Some(ref team_id) = client.team_id {
|
||||
props.insert("team_id".into(), json!(team_id));
|
||||
}
|
||||
if let Some(ref subscription_tier) = client.subscription_tier {
|
||||
props.insert("subscription_tier".into(), json!(subscription_tier));
|
||||
}
|
||||
let _ = mixpanel.engage(&user_id, props).await;
|
||||
});
|
||||
}
|
||||
|
||||
/// Initialize telemetry client. Safe to call multiple times.
|
||||
///
|
||||
/// - `Disabled` → no client
|
||||
/// - `SessionMetrics` → client active (only `session_metrics::*` events fire)
|
||||
/// - `Enabled` → client active (all events fire)
|
||||
///
|
||||
/// `shell_version` is stamped into every event payload as `shell_version`
|
||||
/// (legacy field name preserved for analytics continuity); shell passes its
|
||||
/// own `CARGO_PKG_VERSION`. `http_client` is owned by the caller (typically
|
||||
/// shell's `shared_client()`) so the shared TLS-warmed pool is reused for
|
||||
/// telemetry posts.
|
||||
pub fn init(
|
||||
config: TelemetryConfig,
|
||||
mode: TelemetryMode,
|
||||
user_id: Option<String>,
|
||||
team_id: Option<String>,
|
||||
deployment_key: Option<String>,
|
||||
origin_client: Option<OriginClientInfo>,
|
||||
shell_version: String,
|
||||
subscription_tier: Option<String>,
|
||||
http_client: reqwest::Client,
|
||||
) {
|
||||
let lock = TELEMETRY_CLIENT.get_or_init(|| Mutex::new(None));
|
||||
let mut guard = lock.lock().unwrap_or_else(|err| err.into_inner());
|
||||
*guard = if mode.is_disabled() {
|
||||
None
|
||||
} else {
|
||||
Some(TelemetryClient::from_config(
|
||||
config,
|
||||
mode,
|
||||
user_id,
|
||||
team_id,
|
||||
deployment_key,
|
||||
origin_client,
|
||||
shell_version,
|
||||
subscription_tier,
|
||||
http_client,
|
||||
))
|
||||
};
|
||||
drop(guard);
|
||||
sync_profile();
|
||||
}
|
||||
|
||||
/// Re-initialize the telemetry client if it was not created at startup
|
||||
/// (e.g. because auth was not yet available). No-op when the client
|
||||
/// is already set, so safe to call unconditionally after auth succeeds.
|
||||
pub fn init_if_needed(
|
||||
config: TelemetryConfig,
|
||||
mode: TelemetryMode,
|
||||
user_id: Option<String>,
|
||||
team_id: Option<String>,
|
||||
deployment_key: Option<String>,
|
||||
origin_client: Option<OriginClientInfo>,
|
||||
shell_version: String,
|
||||
subscription_tier: Option<String>,
|
||||
http_client: reqwest::Client,
|
||||
) {
|
||||
if mode.is_disabled() {
|
||||
return;
|
||||
}
|
||||
let lock = TELEMETRY_CLIENT.get_or_init(|| Mutex::new(None));
|
||||
let mut guard = lock.lock().unwrap_or_else(|err| err.into_inner());
|
||||
if guard.is_none() {
|
||||
*guard = Some(TelemetryClient::from_config(
|
||||
config,
|
||||
mode,
|
||||
user_id,
|
||||
team_id,
|
||||
deployment_key,
|
||||
origin_client,
|
||||
shell_version,
|
||||
subscription_tier,
|
||||
http_client,
|
||||
));
|
||||
drop(guard);
|
||||
sync_profile();
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
/// Shell events must still strip to their bare suffix, byte-for-byte
|
||||
/// identical to the previous `strip_prefix("grok-shell-")` behavior.
|
||||
#[test]
|
||||
fn event_value_strips_shell_prefix() {
|
||||
assert_eq!(event_value("grok-shell-turn"), "turn");
|
||||
assert_eq!(
|
||||
event_value("grok-shell-trace_upload_attempted"),
|
||||
"trace_upload_attempted"
|
||||
);
|
||||
}
|
||||
|
||||
/// Workspace events strip their own prefix to the same bare suffix.
|
||||
#[test]
|
||||
fn event_value_strips_workspace_prefix() {
|
||||
assert_eq!(event_value("grok-workspace-turn"), "turn");
|
||||
}
|
||||
|
||||
/// Names without a known emitter prefix pass through unchanged (preserves
|
||||
/// the old `unwrap_or(event_name)` fallback).
|
||||
#[test]
|
||||
fn event_value_passes_through_unprefixed() {
|
||||
assert_eq!(event_value("turn"), "turn");
|
||||
assert_eq!(event_value(""), "");
|
||||
}
|
||||
|
||||
/// Only the leading emitter prefix is stripped; a suffix that itself looks
|
||||
/// like another prefix is left intact.
|
||||
#[test]
|
||||
fn event_value_strips_only_leading_prefix() {
|
||||
assert_eq!(event_value("grok-shell-workspace-x"), "workspace-x");
|
||||
}
|
||||
|
||||
/// The stripper recovers the bare suffix for every origin the emitter can
|
||||
/// produce — ties `event_value` to `EmitterOrigin::event_prefix`.
|
||||
#[test]
|
||||
fn event_value_round_trips_every_emitter_prefix() {
|
||||
for origin in EmitterOrigin::ALL {
|
||||
let name = format!("{}my_event", origin.event_prefix());
|
||||
assert_eq!(event_value(&name), "my_event");
|
||||
}
|
||||
}
|
||||
|
||||
/// Mixpanel `subscription_tier` must be a stable snake_case key. Free
|
||||
/// users arrive as CCP display `"Free"` or JWT-fallback `"free"`; both
|
||||
/// must land as `"free"` (not omitted / not `"Free"`).
|
||||
#[test]
|
||||
fn normalize_tier_maps_display_and_claim_names() {
|
||||
assert_eq!(normalize_tier("Free"), "free");
|
||||
assert_eq!(normalize_tier("free"), "free");
|
||||
assert_eq!(normalize_tier("SuperGrok"), "supergrok");
|
||||
assert_eq!(normalize_tier("SuperGrok Heavy"), "supergrok_heavy");
|
||||
assert_eq!(normalize_tier("supergrok_heavy"), "supergrok_heavy");
|
||||
assert_eq!(normalize_tier("X Basic"), "x_basic");
|
||||
assert_eq!(normalize_tier("X Premium+"), "x_premium_plus");
|
||||
assert_eq!(normalize_tier("X Premium"), "x_premium");
|
||||
assert_eq!(normalize_tier("SuperGrok Lite"), "supergrok_lite");
|
||||
// API key is a dedicated Mixpanel segment — never free.
|
||||
assert_eq!(normalize_tier("API Key"), "api_key");
|
||||
assert_eq!(normalize_tier("api_key"), "api_key");
|
||||
}
|
||||
|
||||
/// `event_value`'s first-match-wins over `EmitterOrigin::ALL` is only
|
||||
/// correct because the emitter prefixes are mutually exclusive: no origin's
|
||||
/// `event_prefix()` is a prefix of another's. If that invariant ever broke
|
||||
/// (e.g. a future `"grok-shell-ext-"` origin), an earlier `ALL` entry could
|
||||
/// strip a shorter prefix first and yield the wrong `event_value`. Pin the
|
||||
/// invariant so adding such a variant fails the suite rather than silently
|
||||
/// corrupting analytics.
|
||||
#[test]
|
||||
fn emitter_prefixes_are_mutually_exclusive() {
|
||||
for a in EmitterOrigin::ALL {
|
||||
for b in EmitterOrigin::ALL {
|
||||
if a != b {
|
||||
assert!(
|
||||
!a.event_prefix().starts_with(b.event_prefix()),
|
||||
"{a:?} prefix {:?} must not start with {b:?} prefix {:?}",
|
||||
a.event_prefix(),
|
||||
b.event_prefix(),
|
||||
);
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
237
crates/codegen/xai-grok-telemetry/src/config.rs
Normal file
237
crates/codegen/xai-grok-telemetry/src/config.rs
Normal file
|
|
@ -0,0 +1,237 @@
|
|||
//! Telemetry-engine configuration.
|
||||
//!
|
||||
//! Extracted from `xai-grok-shell::agent::config` so the data-collector
|
||||
//! engine can construct a [`TelemetryClient`](crate::client::TelemetryClient)
|
||||
//! without a build-time dependency on the shell.
|
||||
//!
|
||||
//! Shell still re-exports these types from their original paths so existing
|
||||
//! call sites (and `Config` derive impls) compile unchanged.
|
||||
use serde::{Deserialize, Serialize};
|
||||
/// Telemetry mode: `true`/`false` (legacy bool) or `"session_metrics"` (string).
|
||||
///
|
||||
/// - `Disabled` -- nothing sent (enterprise default)
|
||||
/// - `SessionMetrics` -- metadata-only lifecycle events, no content
|
||||
/// - `Enabled` -- full product telemetry (events + Mixpanel)
|
||||
#[derive(Debug, Clone, Copy, PartialEq, Eq, Default)]
|
||||
pub enum TelemetryMode {
|
||||
#[default]
|
||||
Disabled,
|
||||
SessionMetrics,
|
||||
Enabled,
|
||||
}
|
||||
impl TelemetryMode {
|
||||
pub fn is_disabled(&self) -> bool {
|
||||
matches!(self, Self::Disabled)
|
||||
}
|
||||
pub fn is_enabled(&self) -> bool {
|
||||
matches!(self, Self::Enabled)
|
||||
}
|
||||
/// True for both `SessionMetrics` and `Enabled`.
|
||||
pub fn session_metrics_enabled(&self) -> bool {
|
||||
matches!(self, Self::SessionMetrics | Self::Enabled)
|
||||
}
|
||||
pub fn parse(s: &str) -> Option<Self> {
|
||||
match s.trim().to_ascii_lowercase().as_str() {
|
||||
"1" | "true" | "yes" | "on" | "enabled" | "full" => Some(Self::Enabled),
|
||||
"0" | "false" | "no" | "off" | "disabled" => Some(Self::Disabled),
|
||||
"session-metrics" | "session_metrics" => Some(Self::SessionMetrics),
|
||||
_ => None,
|
||||
}
|
||||
}
|
||||
}
|
||||
impl std::fmt::Display for TelemetryMode {
|
||||
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
|
||||
match self {
|
||||
Self::Disabled => write!(f, "false"),
|
||||
Self::SessionMetrics => write!(f, "session_metrics"),
|
||||
Self::Enabled => write!(f, "true"),
|
||||
}
|
||||
}
|
||||
}
|
||||
impl From<bool> for TelemetryMode {
|
||||
fn from(b: bool) -> Self {
|
||||
if b { Self::Enabled } else { Self::Disabled }
|
||||
}
|
||||
}
|
||||
impl serde::Serialize for TelemetryMode {
|
||||
fn serialize<S: serde::Serializer>(&self, serializer: S) -> Result<S::Ok, S::Error> {
|
||||
match self {
|
||||
Self::Disabled => serializer.serialize_bool(false),
|
||||
Self::Enabled => serializer.serialize_bool(true),
|
||||
Self::SessionMetrics => serializer.serialize_str("session_metrics"),
|
||||
}
|
||||
}
|
||||
}
|
||||
/// Wire format for `[features] telemetry`: accepts `true`, `false`, or `"session_metrics"`.
|
||||
#[derive(serde::Deserialize)]
|
||||
#[serde(untagged)]
|
||||
enum TelemetryModeValue {
|
||||
Bool(bool),
|
||||
Str(String),
|
||||
}
|
||||
impl<'de> serde::Deserialize<'de> for TelemetryMode {
|
||||
fn deserialize<D: serde::Deserializer<'de>>(deserializer: D) -> Result<Self, D::Error> {
|
||||
match TelemetryModeValue::deserialize(deserializer)? {
|
||||
TelemetryModeValue::Bool(b) => Ok(Self::from(b)),
|
||||
TelemetryModeValue::Str(s) => Ok(Self::parse(&s).unwrap_or_else(|| {
|
||||
tracing::warn!(
|
||||
value = % s,
|
||||
"TELEMETRY_MODE_UNKNOWN: unrecognized telemetry mode; treating as disabled",
|
||||
);
|
||||
Self::Disabled
|
||||
})),
|
||||
}
|
||||
}
|
||||
}
|
||||
/// Parse an env var as a `TelemetryMode`. Returns `None` if unset or empty.
|
||||
pub fn env_telemetry_mode(name: &str) -> Option<TelemetryMode> {
|
||||
let value = std::env::var(name).ok()?;
|
||||
TelemetryMode::parse(&value)
|
||||
}
|
||||
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||
#[serde(default)]
|
||||
pub struct TelemetryConfig {
|
||||
/// Declared for `serde_ignored`. Actual toggle is `[features] telemetry`.
|
||||
#[serde(default)]
|
||||
pub enabled: Option<bool>,
|
||||
pub events_url: Option<String>,
|
||||
pub events_api_key: Option<String>,
|
||||
pub mixpanel_token: Option<String>,
|
||||
pub mixpanel_enabled: bool,
|
||||
/// `None` = inherit from `[features] telemetry`. `Some(false)` = disable GCS uploads only.
|
||||
pub trace_upload: Option<bool>,
|
||||
/// External OTEL master switch (`= GROK_EXTERNAL_OTEL`, env wins).
|
||||
pub otel_enabled: Option<bool>,
|
||||
/// External OTEL metrics exporter: `otlp` | `console` | `none`.
|
||||
pub otel_metrics_exporter: Option<String>,
|
||||
/// External OTEL logs/events exporter: `otlp` | `console` | `none`.
|
||||
pub otel_logs_exporter: Option<String>,
|
||||
/// External OTLP base endpoint (`/v1/logs`, `/v1/metrics` appended for HTTP).
|
||||
pub otel_endpoint: Option<String>,
|
||||
/// External OTLP transport: `http/protobuf` | `grpc`.
|
||||
#[serde(alias = "otel_transport")]
|
||||
pub otel_protocol: Option<String>,
|
||||
/// External OTEL content gate (admins can pin to `false` via requirements).
|
||||
pub otel_log_user_prompts: Option<bool>,
|
||||
/// External OTEL content gate (admins can pin to `false` via requirements).
|
||||
pub otel_log_tool_details: Option<bool>,
|
||||
}
|
||||
fn internal_defaults() -> (Option<String>, Option<String>, Option<String>, bool) {
|
||||
(None, None, None, false)
|
||||
}
|
||||
fn build_env_default(value: Option<&'static str>) -> Option<String> {
|
||||
value
|
||||
.map(str::trim)
|
||||
.filter(|v| !v.is_empty())
|
||||
.map(str::to_owned)
|
||||
}
|
||||
impl Default for TelemetryConfig {
|
||||
fn default() -> Self {
|
||||
let (baked_url, baked_key, baked_token, baked_enabled) = internal_defaults();
|
||||
let build_url = build_env_default(option_env!("GROK_TELEMETRY_BUILD_EVENTS_URL"));
|
||||
let build_key = build_env_default(option_env!("GROK_TELEMETRY_BUILD_EVENTS_API_KEY"));
|
||||
let build_token = build_env_default(option_env!("GROK_TELEMETRY_BUILD_MIXPANEL_TOKEN"));
|
||||
let mixpanel_enabled = baked_enabled || build_token.is_some();
|
||||
let (events_url, events_api_key, mixpanel_token) = (
|
||||
build_url.or(baked_url),
|
||||
build_key.or(baked_key),
|
||||
build_token.or(baked_token),
|
||||
);
|
||||
Self {
|
||||
enabled: None,
|
||||
events_url,
|
||||
events_api_key,
|
||||
mixpanel_token,
|
||||
mixpanel_enabled,
|
||||
trace_upload: None,
|
||||
otel_enabled: None,
|
||||
otel_metrics_exporter: None,
|
||||
otel_logs_exporter: None,
|
||||
otel_endpoint: None,
|
||||
otel_protocol: None,
|
||||
otel_log_user_prompts: None,
|
||||
otel_log_tool_details: None,
|
||||
}
|
||||
}
|
||||
}
|
||||
impl TelemetryConfig {
|
||||
pub fn apply_env_overrides(&mut self) {
|
||||
self.normalize();
|
||||
if let Some(value) = Self::env_override("GROK_TELEMETRY_EVENTS_URL") {
|
||||
self.events_url = value;
|
||||
}
|
||||
if let Some(value) = Self::env_override("GROK_TELEMETRY_EVENTS_API_KEY") {
|
||||
self.events_api_key = value;
|
||||
}
|
||||
if let Some(value) = Self::env_override("GROK_TELEMETRY_MIXPANEL_TOKEN") {
|
||||
self.mixpanel_token = value;
|
||||
}
|
||||
if let Some(value) = env_bool("GROK_TELEMETRY_MIXPANEL_ENABLED") {
|
||||
self.mixpanel_enabled = value;
|
||||
}
|
||||
if let Some(value) = env_bool("GROK_TELEMETRY_TRACE_UPLOAD") {
|
||||
self.trace_upload = Some(value);
|
||||
}
|
||||
}
|
||||
fn normalize(&mut self) {
|
||||
self.events_url = Self::normalize_optional_string(self.events_url.take());
|
||||
self.events_api_key = Self::normalize_optional_string(self.events_api_key.take());
|
||||
self.mixpanel_token = Self::normalize_optional_string(self.mixpanel_token.take());
|
||||
}
|
||||
fn env_override(name: &str) -> Option<Option<String>> {
|
||||
match std::env::var(name) {
|
||||
Ok(value) => Some(Self::normalize_optional_string(Some(value))),
|
||||
Err(_) => None,
|
||||
}
|
||||
}
|
||||
fn normalize_optional_string(value: Option<String>) -> Option<String> {
|
||||
value.and_then(|raw| {
|
||||
let trimmed = raw.trim();
|
||||
if trimmed.is_empty() {
|
||||
None
|
||||
} else {
|
||||
Some(trimmed.to_string())
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
/// Parse an env var as a boolean. Returns `None` if unset or unrecognized.
|
||||
///
|
||||
/// Local copy of `xai_grok_shell::agent::config::env_bool` so this crate
|
||||
/// stays free of a shell back-edge. Shell keeps its own copy for callers
|
||||
/// outside the telemetry config path.
|
||||
fn env_bool(name: &str) -> Option<bool> {
|
||||
let value = std::env::var(name).ok()?;
|
||||
match value.trim().to_ascii_lowercase().as_str() {
|
||||
"" => None,
|
||||
"1" | "true" | "yes" | "on" | "enabled" => Some(true),
|
||||
"0" | "false" | "no" | "off" | "disabled" => Some(false),
|
||||
_ => None,
|
||||
}
|
||||
}
|
||||
/// Derive a stable deployment ID (UUIDv5) from the deployment key.
|
||||
pub fn deployment_id_from_key(key: &str) -> String {
|
||||
uuid::Uuid::new_v5(&uuid::Uuid::NAMESPACE_OID, key.as_bytes()).to_string()
|
||||
}
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
#[test]
|
||||
fn build_env_default_normalizes() {
|
||||
assert_eq!(build_env_default(None), None);
|
||||
assert_eq!(build_env_default(Some("")), None);
|
||||
assert_eq!(build_env_default(Some(" \t ")), None);
|
||||
assert_eq!(build_env_default(Some(" key ")), Some("key".to_owned()));
|
||||
}
|
||||
#[test]
|
||||
fn default_is_build_env_layer_when_feature_off() {
|
||||
let cfg = TelemetryConfig::default();
|
||||
let url = build_env_default(option_env!("GROK_TELEMETRY_BUILD_EVENTS_URL"));
|
||||
let key = build_env_default(option_env!("GROK_TELEMETRY_BUILD_EVENTS_API_KEY"));
|
||||
let token = build_env_default(option_env!("GROK_TELEMETRY_BUILD_MIXPANEL_TOKEN"));
|
||||
assert_eq!(cfg.mixpanel_enabled, token.is_some());
|
||||
assert_eq!(cfg.events_url, url);
|
||||
assert_eq!(cfg.events_api_key, key);
|
||||
assert_eq!(cfg.mixpanel_token, token);
|
||||
}
|
||||
}
|
||||
14
crates/codegen/xai-grok-telemetry/src/context.rs
Normal file
14
crates/codegen/xai-grok-telemetry/src/context.rs
Normal file
|
|
@ -0,0 +1,14 @@
|
|||
//! Git context collection for telemetry events.
|
||||
|
||||
pub struct GitContext {
|
||||
pub is_git_repo: bool,
|
||||
}
|
||||
|
||||
pub fn collect_git_context(cwd: &str) -> GitContext {
|
||||
use git2::Repository;
|
||||
use std::path::Path;
|
||||
|
||||
GitContext {
|
||||
is_git_repo: Repository::discover(Path::new(cwd)).is_ok(),
|
||||
}
|
||||
}
|
||||
961
crates/codegen/xai-grok-telemetry/src/debug_log.rs
Normal file
961
crates/codegen/xai-grok-telemetry/src/debug_log.rs
Normal file
|
|
@ -0,0 +1,961 @@
|
|||
//! Reusable non-blocking file-logging tracing layers for the `--debug` firehose.
|
||||
//!
|
||||
//! Two install modes, chosen by env precedence (see `resolve_debug_target_inner`):
|
||||
//! - PerSession (`GROK_DEBUG_LOG=1`): a routing layer fans each session's
|
||||
//! firehose to `~/.grok/debug/<session_id>.txt` (one file per session), with a
|
||||
//! `<role>-<pid>.txt` catch-all for events fired outside any session span, and
|
||||
//! a `latest.txt` symlink pointing at the most-recently-opened session file.
|
||||
//! - SingleFile (explicit path via `GROK_LOG_FILE` or `GROK_DEBUG_LOG=<path>`):
|
||||
//! one flat `fmt` file, routing bypassed. Disk IO stays off the tracing hot
|
||||
//! path via `tracing_appender`'s non-blocking writer in both modes.
|
||||
|
||||
use std::collections::HashMap;
|
||||
use std::ffi::OsStr;
|
||||
use std::io::Write as _;
|
||||
use std::path::{Path, PathBuf};
|
||||
use std::sync::{Mutex, MutexGuard};
|
||||
|
||||
use tracing::Subscriber;
|
||||
use tracing::field::{Field, Visit};
|
||||
use tracing_appender::non_blocking::NonBlocking;
|
||||
use tracing_subscriber::filter::{EnvFilter, LevelFilter};
|
||||
use tracing_subscriber::layer::{Context, Layer};
|
||||
use tracing_subscriber::registry::LookupSpan;
|
||||
|
||||
use crate::session_ctx::SESSION_ID_FIELD;
|
||||
use xai_grok_config::grok_home;
|
||||
|
||||
/// Which env var requested a single-file debug log (drives filter and diagnostics).
|
||||
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
|
||||
pub(crate) enum DebugSource {
|
||||
GrokLogFile,
|
||||
GrokDebugLog,
|
||||
}
|
||||
|
||||
impl DebugSource {
|
||||
fn label(self) -> &'static str {
|
||||
match self {
|
||||
Self::GrokLogFile => "GROK_LOG_FILE",
|
||||
Self::GrokDebugLog => "GROK_DEBUG_LOG",
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Target for the pager's always-on compact ACP update summary line
|
||||
/// (kind, ids, status, payload sizes).
|
||||
///
|
||||
/// Lives here (not in `xai-grok-pager`) so the firehose directives below and
|
||||
/// the pager's own filter are built from the same constants — a rename can't
|
||||
/// silently desync them into a no-op directive.
|
||||
pub const ACP_UPDATE_TARGET: &str = "acp_update";
|
||||
|
||||
/// Target for the pager's full ACP update payload dump (plain JSON).
|
||||
///
|
||||
/// Off in the pager's release filter; the firehose is the always-available
|
||||
/// subscriber for full payloads, and it writes to disk, where the volume is
|
||||
/// safe. See `xai-grok-pager/src/tracing.rs` for the consumer side.
|
||||
pub const ACP_UPDATE_PAYLOAD_TARGET: &str = "acp_update_payload";
|
||||
|
||||
/// Module path of rmcp 2.1's per-reconnect SSE warn (`sse stream error: ...`),
|
||||
/// which subscribers demote to `error` to drop the flood. Re-check on rmcp bump.
|
||||
pub const RMCP_SSE_NOISE_TARGET: &str = "rmcp::transport::common::client_side_sse";
|
||||
|
||||
// Broad firehose filter for the routing/GROK_DEBUG_LOG sources: capture our
|
||||
// crates at debug regardless of a narrowing RUST_LOG, with deps at info so they
|
||||
// don't flood. Curated first-party allowlist: new grok crates default to `info`
|
||||
// until added here.
|
||||
const FIREHOSE_BASE_DIRECTIVES: &str = "info,xai_grok_pager=debug,xai_grok_shell=debug,xai_grok_tools=debug,xai_grok_telemetry=debug,xai_grok_agent=debug,xai_grok_mcp=debug,xai_acp_lib=debug,sampling_log=off";
|
||||
|
||||
// Full firehose directives: the curated crate list plus the pager's ACP
|
||||
// update target (built from the constant above, not a literal).
|
||||
fn firehose_directives() -> String {
|
||||
format!("{FIREHOSE_BASE_DIRECTIVES},{ACP_UPDATE_TARGET}=debug")
|
||||
}
|
||||
|
||||
// The broad firehose filter, used by both the routing layer and the
|
||||
// GROK_DEBUG_LOG single-file source (mirrors `default_file_filter`).
|
||||
fn firehose_filter() -> EnvFilter {
|
||||
EnvFilter::new(firehose_directives())
|
||||
}
|
||||
|
||||
// RUST_LOG-respecting filter for the GROK_LOG_FILE source: DEBUG default, honor
|
||||
// RUST_LOG, silence sampling_log (preserves GROK_LOG_FILE back-compat).
|
||||
fn default_file_filter() -> EnvFilter {
|
||||
EnvFilter::builder()
|
||||
.with_default_directive(LevelFilter::DEBUG.into())
|
||||
.from_env_lossy()
|
||||
.add_directive(
|
||||
"sampling_log=off"
|
||||
.parse()
|
||||
.expect("static directive is valid"),
|
||||
)
|
||||
}
|
||||
|
||||
// Open `path` as a non-blocking flat `fmt` layer with `filter`; ansi off, target on.
|
||||
fn build_file_layer<S>(path: &Path, filter: EnvFilter) -> std::io::Result<impl Layer<S>>
|
||||
where
|
||||
S: Subscriber + for<'span> LookupSpan<'span>,
|
||||
{
|
||||
let non_blocking = crate::appender::non_blocking_file_writer(path)?;
|
||||
let fmt_layer = tracing_subscriber::fmt::layer()
|
||||
.with_target(true)
|
||||
.with_ansi(false)
|
||||
.with_writer(non_blocking)
|
||||
.with_filter(filter);
|
||||
Ok(fmt_layer)
|
||||
}
|
||||
|
||||
// ── Per-session routing layer ───────────────────────────────────────────────
|
||||
|
||||
/// Filesystem-safe session key. Sanitized once at capture (`on_new_span`) and
|
||||
/// stashed in the span's tracing extensions, so events fired anywhere under the
|
||||
/// span route to the right file without re-sanitizing on the hot path.
|
||||
#[derive(Clone)]
|
||||
struct SessionId(String);
|
||||
|
||||
/// Visits span attributes to pull out the `session_id` field. Production records
|
||||
/// it via `%` (Display → `record_debug`, no quotes); like `EventVisitor`, the
|
||||
/// single `record_debug` impl captures every field type (the other recorders
|
||||
/// default to it).
|
||||
#[derive(Default)]
|
||||
struct SessionIdVisitor(Option<String>);
|
||||
|
||||
impl Visit for SessionIdVisitor {
|
||||
fn record_debug(&mut self, field: &Field, value: &dyn std::fmt::Debug) {
|
||||
if field.name() == SESSION_ID_FIELD {
|
||||
self.0 = Some(format!("{value:?}"));
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Renders an event's message + remaining fields into plain strings. All field
|
||||
/// types funnel through `record_debug` (the trait's other recorders default to
|
||||
/// it), so this one impl captures everything.
|
||||
#[derive(Default)]
|
||||
struct EventVisitor {
|
||||
message: String,
|
||||
fields: String,
|
||||
}
|
||||
|
||||
impl Visit for EventVisitor {
|
||||
fn record_debug(&mut self, field: &Field, value: &dyn std::fmt::Debug) {
|
||||
use std::fmt::Write as _;
|
||||
if field.name() == "message" {
|
||||
let _ = write!(self.message, "{value:?}");
|
||||
} else {
|
||||
let _ = write!(self.fields, " {}={:?}", field.name(), value);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// Format one compact, ANSI-free firehose line. Intentionally NOT byte-identical
|
||||
// to `fmt::Layer`: its `FormatEvent` can't be reused from another layer and a
|
||||
// `MakeWriter` can't see span context, so we render here. Span context is
|
||||
// omitted on purpose — the file name already carries the session id.
|
||||
fn format_event(event: &tracing::Event<'_>) -> String {
|
||||
let meta = event.metadata();
|
||||
let mut visitor = EventVisitor::default();
|
||||
event.record(&mut visitor);
|
||||
let ts = chrono::Utc::now().to_rfc3339_opts(chrono::SecondsFormat::Micros, true);
|
||||
let level = meta.level();
|
||||
let target = meta.target();
|
||||
// Skip the message gap when there's no `message` field so a field-only event
|
||||
// renders "target: k=v" (each field already carries a leading space), not
|
||||
// "target: k=v" with a dangling double space.
|
||||
if visitor.message.is_empty() {
|
||||
format!("{ts} {level} {target}:{}\n", visitor.fields)
|
||||
} else {
|
||||
format!(
|
||||
"{ts} {level} {target}: {}{}\n",
|
||||
visitor.message, visitor.fields
|
||||
)
|
||||
}
|
||||
}
|
||||
|
||||
// Keep per-session file names filesystem-safe. A session id is normally a UUID,
|
||||
// but never let an unexpected value (path separators, `..`) escape the dir.
|
||||
fn sanitize_key(id: &str) -> String {
|
||||
let safe: String = id
|
||||
.chars()
|
||||
.map(|c| {
|
||||
if c.is_ascii_alphanumeric() || matches!(c, '-' | '_' | '.') {
|
||||
c
|
||||
} else {
|
||||
'_'
|
||||
}
|
||||
})
|
||||
.collect();
|
||||
// Map empty / dot-only keys ("", ".", "..", "...") to a constant: those are
|
||||
// filesystem-special, and relying on the `.txt` suffix to neutralize them is
|
||||
// incidental. Make the safety explicit instead.
|
||||
if safe.is_empty() || safe.bytes().all(|b| b == b'.') {
|
||||
return "_".to_owned();
|
||||
}
|
||||
safe
|
||||
}
|
||||
|
||||
// `latest.txt` link + swap-temp name parts, shared by `update_latest_symlink`
|
||||
// (create/rename) and `prune_old_logs` (spare rule + orphan cleanup) so the
|
||||
// sites can never drift. Tests pin the literals on purpose: orphans created by
|
||||
// already-shipped binaries must stay reapable across a rename of these consts.
|
||||
const LATEST_LINK_NAME: &str = "latest.txt";
|
||||
const LATEST_TMP_PREFIX: &str = ".latest.";
|
||||
const LATEST_TMP_SUFFIX: &str = ".tmp";
|
||||
|
||||
/// Repoint `<dir>/latest.txt` at `target` (a sibling session file) for
|
||||
/// `tail -f`. Best-effort and Unix-only; the relative target keeps the link
|
||||
/// valid regardless of the dir's absolute path.
|
||||
#[cfg(unix)]
|
||||
fn update_latest_symlink(dir: &Path, target: &Path) {
|
||||
let Some(name) = target.file_name() else {
|
||||
return;
|
||||
};
|
||||
// Atomic swap: symlink a unique temp then rename it over latest.txt (rename
|
||||
// is atomic on POSIX), so a racing `tail -f` never sees latest.txt missing.
|
||||
// The temp name is keyed by the target file so concurrent opens of different
|
||||
// sessions don't collide on it.
|
||||
let tmp = dir.join(format!(
|
||||
"{LATEST_TMP_PREFIX}{}{LATEST_TMP_SUFFIX}",
|
||||
name.to_string_lossy()
|
||||
));
|
||||
let _ = std::fs::remove_file(&tmp);
|
||||
if std::os::unix::fs::symlink(name, &tmp).is_ok()
|
||||
&& std::fs::rename(&tmp, dir.join(LATEST_LINK_NAME)).is_err()
|
||||
{
|
||||
// Rename failed: remove the temp symlink now rather than leaving an
|
||||
// orphan for prune to reap only after LOG_RETENTION.
|
||||
let _ = std::fs::remove_file(&tmp);
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(not(unix))]
|
||||
fn update_latest_symlink(_dir: &Path, _target: &Path) {}
|
||||
|
||||
/// Per-session sinks plus a single fallback sink, all behind the routing layer's
|
||||
/// mutex. There is no cap or eviction: each distinct session id opens one file +
|
||||
/// non-blocking worker + parked guard that persist for the process lifetime
|
||||
/// (reclaimed only when the process/leader restarts). That is acceptable for an
|
||||
/// opt-in, debug-only firehose; a long-lived `--debug` leader holds one fd per
|
||||
/// session it logs. The central guard parking (`appender`) is what lets
|
||||
/// `flush()` drain these at exit, so we do not reclaim per session.
|
||||
#[derive(Default)]
|
||||
struct SinkMap {
|
||||
sessions: HashMap<String, NonBlocking>,
|
||||
fallback: Option<NonBlocking>,
|
||||
}
|
||||
|
||||
/// Routes the firehose per session: events under a `session` span go to
|
||||
/// `<dir>/<session_id>.txt`; everything else to `<dir>/<role>-<pid>.txt`.
|
||||
struct RoutingLayer {
|
||||
dir: PathBuf,
|
||||
role: String,
|
||||
pid: u32,
|
||||
// The lock is scoped to map access ONLY — file opens (fs + a worker-thread
|
||||
// spawn + the appender's own mutex) run OUTSIDE it, so a tracing event
|
||||
// emitted on the open path can't re-enter and deadlock this non-reentrant
|
||||
// Mutex. Lock-on-write is otherwise fine: the firehose is opt-in/debug-only.
|
||||
sinks: Mutex<SinkMap>,
|
||||
}
|
||||
|
||||
impl RoutingLayer {
|
||||
fn new(dir: PathBuf, role: String, pid: u32) -> Self {
|
||||
Self {
|
||||
dir,
|
||||
role,
|
||||
pid,
|
||||
sinks: Mutex::new(SinkMap::default()),
|
||||
}
|
||||
}
|
||||
|
||||
fn lock(&self) -> MutexGuard<'_, SinkMap> {
|
||||
self.sinks.lock().unwrap_or_else(|p| p.into_inner())
|
||||
}
|
||||
|
||||
// Append `line` to the session's file. `key` is already sanitized. Opens (and
|
||||
// points `latest.txt` at) the file on first use; open failures degrade to a
|
||||
// no-op for that file.
|
||||
fn write_session(&self, key: &str, line: &[u8]) {
|
||||
// Fast path: writer already open. Hold the lock only for the lookup + the
|
||||
// (non-blocking, channel-only) write.
|
||||
{
|
||||
let mut map = self.lock();
|
||||
if let Some(writer) = map.sessions.get_mut(key) {
|
||||
let _ = writer.write_all(line);
|
||||
return;
|
||||
}
|
||||
}
|
||||
// First event for this session: open OUTSIDE the lock.
|
||||
let path = self.dir.join(format!("{key}.txt"));
|
||||
let Ok(mut writer) = crate::appender::non_blocking_file_writer(&path) else {
|
||||
return;
|
||||
};
|
||||
update_latest_symlink(&self.dir, &path);
|
||||
let _ = writer.write_all(line);
|
||||
let mut map = self.lock();
|
||||
// If a concurrent event opened it first, keep that one and drop ours (the
|
||||
// line we wrote already reached the file via our worker).
|
||||
map.sessions.entry(key.to_owned()).or_insert(writer);
|
||||
}
|
||||
|
||||
// Append `line` to the `<role>-<pid>.txt` catch-all, opening it on first use.
|
||||
fn write_fallback(&self, line: &[u8]) {
|
||||
{
|
||||
let mut map = self.lock();
|
||||
if let Some(writer) = map.fallback.as_mut() {
|
||||
let _ = writer.write_all(line);
|
||||
return;
|
||||
}
|
||||
}
|
||||
let path = self.dir.join(format!("{}-{}.txt", self.role, self.pid));
|
||||
let Ok(mut writer) = crate::appender::non_blocking_file_writer(&path) else {
|
||||
return;
|
||||
};
|
||||
let _ = writer.write_all(line);
|
||||
let mut map = self.lock();
|
||||
if map.fallback.is_none() {
|
||||
map.fallback = Some(writer);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl<S> Layer<S> for RoutingLayer
|
||||
where
|
||||
S: Subscriber + for<'span> LookupSpan<'span>,
|
||||
{
|
||||
fn on_new_span(
|
||||
&self,
|
||||
attrs: &tracing::span::Attributes<'_>,
|
||||
id: &tracing::span::Id,
|
||||
ctx: Context<'_, S>,
|
||||
) {
|
||||
let mut visitor = SessionIdVisitor::default();
|
||||
attrs.record(&mut visitor);
|
||||
if let Some(sid) = visitor.0
|
||||
&& let Some(span) = ctx.span(id)
|
||||
{
|
||||
// Sanitize once at capture so the stored key is always filesystem-safe
|
||||
// and `on_event` never re-sanitizes on the hot path.
|
||||
span.extensions_mut().insert(SessionId(sanitize_key(&sid)));
|
||||
}
|
||||
}
|
||||
|
||||
fn on_event(&self, event: &tracing::Event<'_>, ctx: Context<'_, S>) {
|
||||
// Nearest enclosing span (leaf→root) carrying a session id wins. The key
|
||||
// is already sanitized (stored at `on_new_span`).
|
||||
let session_key = ctx.event_scope(event).and_then(|scope| {
|
||||
scope
|
||||
.into_iter()
|
||||
.find_map(|span| span.extensions().get::<SessionId>().map(|s| s.0.clone()))
|
||||
});
|
||||
let line = format_event(event);
|
||||
match session_key {
|
||||
Some(key) => self.write_session(&key, line.as_bytes()),
|
||||
None => self.write_fallback(line.as_bytes()),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// ── Install + lifecycle ──────────────────────────────────────────────────────
|
||||
|
||||
/// Resolve the requested debug target and install the matching firehose layer on
|
||||
/// `registry`, then init the subscriber.
|
||||
///
|
||||
/// PerSession installs the routing layer (firehose filter, RUST_LOG-immune) and
|
||||
/// prunes old session logs; SingleFile installs a flat `fmt` file picking the
|
||||
/// filter by source (GROK_LOG_FILE respects RUST_LOG). Open failures warn AFTER
|
||||
/// init in the single-file case; routing open failures are per-file at write
|
||||
/// time and degrade gracefully. `role` names the per-pid fallback file.
|
||||
pub fn install_firehose<S>(registry: S, role: &str)
|
||||
where
|
||||
S: Subscriber + for<'span> LookupSpan<'span> + Send + Sync + 'static,
|
||||
{
|
||||
use tracing_subscriber::layer::SubscriberExt as _;
|
||||
use tracing_subscriber::util::SubscriberInitExt as _;
|
||||
|
||||
match resolve_debug_target() {
|
||||
Some(DebugTarget::PerSession { dir }) => {
|
||||
let layer = RoutingLayer::new(dir, role.to_owned(), std::process::id())
|
||||
.with_filter(firehose_filter());
|
||||
registry.with(layer).init();
|
||||
// Tie pruning to actually routing a firehose, not to the flag.
|
||||
sweep_old_logs();
|
||||
}
|
||||
Some(DebugTarget::SingleFile { path, src }) => {
|
||||
let filter = match src {
|
||||
DebugSource::GrokLogFile => default_file_filter(),
|
||||
DebugSource::GrokDebugLog => firehose_filter(),
|
||||
};
|
||||
match build_file_layer::<S>(&path, filter) {
|
||||
Ok(layer) => registry.with(layer).init(),
|
||||
Err(e) => {
|
||||
registry.init();
|
||||
tracing::warn!("failed to open {} {path:?}: {e}", src.label());
|
||||
}
|
||||
}
|
||||
}
|
||||
None => registry.init(),
|
||||
}
|
||||
}
|
||||
|
||||
/// Flush parked firehose writers at process exit (no-op when none installed).
|
||||
pub fn flush() {
|
||||
crate::appender::flush_file_log_guards();
|
||||
}
|
||||
|
||||
/// Where the firehose should go, if anywhere.
|
||||
#[derive(Debug, PartialEq, Eq)]
|
||||
pub(crate) enum DebugTarget {
|
||||
/// `GROK_DEBUG_LOG=1` → route per session into `<dir>` (`~/.grok/debug`).
|
||||
PerSession { dir: PathBuf },
|
||||
/// An explicit path → one flat `fmt` file, routing bypassed.
|
||||
SingleFile { path: PathBuf, src: DebugSource },
|
||||
}
|
||||
|
||||
/// Resolve the debug target, honoring precedence: explicit GROK_LOG_FILE wins
|
||||
/// (single file, RUST_LOG filter); else GROK_DEBUG_LOG — a truthy bool routes
|
||||
/// per session into `~/.grok/debug`, an explicit path writes a single file.
|
||||
///
|
||||
/// Read via `var_os` (not `var`) so a non-UTF-8 path isn't silently dropped.
|
||||
pub(crate) fn resolve_debug_target() -> Option<DebugTarget> {
|
||||
let grok_log_file = std::env::var_os("GROK_LOG_FILE");
|
||||
let grok_debug_log = std::env::var_os("GROK_DEBUG_LOG");
|
||||
resolve_debug_target_inner(
|
||||
grok_log_file.as_deref(),
|
||||
grok_debug_log.as_deref(),
|
||||
&grok_home().join("debug"),
|
||||
)
|
||||
}
|
||||
|
||||
// Empty / whitespace (when valid UTF-8) counts as unset; a non-UTF-8 value is
|
||||
// never blank.
|
||||
fn is_blank(v: &OsStr) -> bool {
|
||||
v.to_str().is_some_and(|s| s.trim().is_empty())
|
||||
}
|
||||
|
||||
// Build a path from an env value: trim surrounding whitespace when it is valid
|
||||
// UTF-8, and preserve the raw bytes otherwise (non-UTF-8 paths must survive).
|
||||
fn os_path(v: &OsStr) -> PathBuf {
|
||||
match v.to_str() {
|
||||
Some(s) => PathBuf::from(s.trim()),
|
||||
None => PathBuf::from(v),
|
||||
}
|
||||
}
|
||||
|
||||
// Env-free precedence core so the resolution rules are unit-testable. The role
|
||||
// and pid are no longer part of resolution: the routing layer owns fallback
|
||||
// naming, so resolution only decides routing-dir vs single-file-path. Takes
|
||||
// `OsStr` so non-UTF-8 paths round-trip; only the bool-vs-path discrimination
|
||||
// needs UTF-8 (a non-UTF-8 value can't be a bool keyword, so it's a path).
|
||||
fn resolve_debug_target_inner(
|
||||
grok_log_file: Option<&OsStr>,
|
||||
grok_debug_log: Option<&OsStr>,
|
||||
debug_dir: &Path,
|
||||
) -> Option<DebugTarget> {
|
||||
if let Some(raw) = grok_log_file
|
||||
&& !is_blank(raw)
|
||||
{
|
||||
return Some(DebugTarget::SingleFile {
|
||||
path: os_path(raw),
|
||||
src: DebugSource::GrokLogFile,
|
||||
});
|
||||
}
|
||||
let raw = grok_debug_log?;
|
||||
match raw.to_str().map(str::trim) {
|
||||
Some("" | "0" | "false" | "off" | "no") => None,
|
||||
Some("1" | "true" | "on" | "yes") => Some(DebugTarget::PerSession {
|
||||
dir: debug_dir.to_path_buf(),
|
||||
}),
|
||||
// Any other UTF-8 value, or a non-UTF-8 value (`None`), is an explicit path.
|
||||
_ => Some(DebugTarget::SingleFile {
|
||||
path: os_path(raw),
|
||||
src: DebugSource::GrokDebugLog,
|
||||
}),
|
||||
}
|
||||
}
|
||||
|
||||
/// Retention window for firehose debug logs: files older than this are pruned.
|
||||
const LOG_RETENTION: std::time::Duration = std::time::Duration::from_secs(7 * 24 * 60 * 60);
|
||||
|
||||
/// Prune `*.txt` firehose files (and orphaned `latest.txt` swap temps) under
|
||||
/// `~/.grok/debug` older than [`LOG_RETENTION`] so the dir doesn't grow
|
||||
/// unbounded. Age-based (not count-based) so a still-open log from a concurrent
|
||||
/// process is never unlinked mid-write; best-effort, ignore errors.
|
||||
pub(crate) fn sweep_old_logs() {
|
||||
prune_old_logs(&grok_home().join("debug"), LOG_RETENTION);
|
||||
}
|
||||
|
||||
// Pure prune core: remove `*.txt` files and orphaned `latest.txt` swap temps in
|
||||
// `dir` older than `max_age`. Age-based so a recently-written (active) log is
|
||||
// never deleted; spares the `latest.txt` symlink (a stale link is harmless and
|
||||
// never an active file); best-effort so cleanup never fails logging setup;
|
||||
// testable against a tempdir.
|
||||
fn prune_old_logs(dir: &Path, max_age: std::time::Duration) {
|
||||
let Ok(entries) = std::fs::read_dir(dir) else {
|
||||
return;
|
||||
};
|
||||
let now = std::time::SystemTime::now();
|
||||
for entry in entries.flatten() {
|
||||
let path = entry.path();
|
||||
let Some(name) = path.file_name().and_then(|n| n.to_str()) else {
|
||||
continue;
|
||||
};
|
||||
let is_log = name.ends_with(".txt") && name != LATEST_LINK_NAME;
|
||||
// Swap temps matching this shape that survive the age gate below are
|
||||
// orphans of a crash between `update_latest_symlink`'s create and rename.
|
||||
let is_latest_swap_tmp =
|
||||
name.starts_with(LATEST_TMP_PREFIX) && name.ends_with(LATEST_TMP_SUFFIX);
|
||||
if !is_log && !is_latest_swap_tmp {
|
||||
continue;
|
||||
}
|
||||
// `DirEntry::metadata` does not follow symlinks, so a dangling orphaned
|
||||
// temp still yields its own mtime here.
|
||||
let Ok(modified) = entry.metadata().and_then(|m| m.modified()) else {
|
||||
continue;
|
||||
};
|
||||
if now.duration_since(modified).is_ok_and(|age| age > max_age) {
|
||||
let _ = std::fs::remove_file(&path);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
// Routing tests drive real non-blocking writers whose worker guards are
|
||||
// parked in a process-lifetime static; flushing drains ALL of them. Serialize
|
||||
// such tests so a concurrent `cargo test` thread can't clear another's guards
|
||||
// before it reads. (nextest already isolates each test in its own process.)
|
||||
fn flush_test_lock() -> std::sync::MutexGuard<'static, ()> {
|
||||
static LOCK: Mutex<()> = Mutex::new(());
|
||||
LOCK.lock().unwrap_or_else(|p| p.into_inner())
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn build_file_layer_creates_parent_dir_and_file() {
|
||||
let dir = tempfile::tempdir().unwrap();
|
||||
let path = dir.path().join("nested").join("debug.log");
|
||||
assert!(!path.parent().unwrap().exists());
|
||||
|
||||
let layer = build_file_layer::<tracing_subscriber::Registry>(&path, default_file_filter());
|
||||
|
||||
assert!(layer.is_ok());
|
||||
assert!(path.parent().unwrap().exists());
|
||||
assert!(path.exists());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn build_file_layer_errors_when_open_fails() {
|
||||
// Opening an existing directory in append mode fails, exercising the Err path.
|
||||
let dir = tempfile::tempdir().unwrap();
|
||||
let layer =
|
||||
build_file_layer::<tracing_subscriber::Registry>(dir.path(), default_file_filter());
|
||||
assert!(layer.is_err());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn resolve_target_unset_is_none() {
|
||||
assert!(resolve_debug_target_inner(None, None, Path::new("/debug")).is_none());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn resolve_target_debug_log_disabled_is_none() {
|
||||
for v in ["0", "false", "off", "no", "", " "] {
|
||||
assert!(
|
||||
resolve_debug_target_inner(None, Some(OsStr::new(v)), Path::new("/debug"))
|
||||
.is_none(),
|
||||
"expected None for GROK_DEBUG_LOG={v:?}"
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn resolve_target_debug_log_enabled_is_per_session_dir() {
|
||||
for v in ["1", "true", "on", "yes"] {
|
||||
let target =
|
||||
resolve_debug_target_inner(None, Some(OsStr::new(v)), Path::new("/debug")).unwrap();
|
||||
assert_eq!(
|
||||
target,
|
||||
DebugTarget::PerSession {
|
||||
dir: PathBuf::from("/debug")
|
||||
},
|
||||
"expected PerSession for GROK_DEBUG_LOG={v:?}"
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn resolve_target_debug_log_custom_path_is_single_file() {
|
||||
let target = resolve_debug_target_inner(
|
||||
None,
|
||||
Some(OsStr::new("/tmp/custom.log")),
|
||||
Path::new("/debug"),
|
||||
)
|
||||
.unwrap();
|
||||
assert_eq!(
|
||||
target,
|
||||
DebugTarget::SingleFile {
|
||||
path: PathBuf::from("/tmp/custom.log"),
|
||||
src: DebugSource::GrokDebugLog,
|
||||
}
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn resolve_target_log_file_wins_over_debug_log() {
|
||||
let target = resolve_debug_target_inner(
|
||||
Some(OsStr::new("/tmp/explicit.log")),
|
||||
Some(OsStr::new("1")),
|
||||
Path::new("/debug"),
|
||||
)
|
||||
.unwrap();
|
||||
assert_eq!(
|
||||
target,
|
||||
DebugTarget::SingleFile {
|
||||
path: PathBuf::from("/tmp/explicit.log"),
|
||||
src: DebugSource::GrokLogFile,
|
||||
}
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn resolve_target_empty_log_file_falls_through_to_debug_log() {
|
||||
// Empty / whitespace GROK_LOG_FILE is treated as unset (mirrors GROK_DEBUG_LOG).
|
||||
for blank in ["", " "] {
|
||||
let target = resolve_debug_target_inner(
|
||||
Some(OsStr::new(blank)),
|
||||
Some(OsStr::new("1")),
|
||||
Path::new("/debug"),
|
||||
)
|
||||
.unwrap();
|
||||
assert_eq!(
|
||||
target,
|
||||
DebugTarget::PerSession {
|
||||
dir: PathBuf::from("/debug")
|
||||
}
|
||||
);
|
||||
}
|
||||
assert!(
|
||||
resolve_debug_target_inner(Some(OsStr::new("")), None, Path::new("/debug")).is_none()
|
||||
);
|
||||
}
|
||||
|
||||
#[cfg(unix)]
|
||||
#[test]
|
||||
fn resolve_target_non_utf8_debug_log_path_is_single_file() {
|
||||
use std::os::unix::ffi::OsStrExt;
|
||||
|
||||
// A non-UTF-8 GROK_DEBUG_LOG value is a path, not a bool keyword, and its
|
||||
// bytes must round-trip (not be silently dropped).
|
||||
let raw = OsStr::from_bytes(b"/tmp/\xff/fire.txt");
|
||||
let target = resolve_debug_target_inner(None, Some(raw), Path::new("/debug")).unwrap();
|
||||
match target {
|
||||
DebugTarget::SingleFile { path, src } => {
|
||||
assert_eq!(src, DebugSource::GrokDebugLog);
|
||||
assert_eq!(path.as_os_str(), raw);
|
||||
}
|
||||
other => panic!("expected SingleFile for non-UTF-8 path, got {other:?}"),
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn firehose_directives_parse() {
|
||||
// Guard against the const rotting: every directive must parse strictly.
|
||||
for d in firehose_directives().split(',') {
|
||||
d.parse::<tracing_subscriber::filter::Directive>()
|
||||
.unwrap_or_else(|e| panic!("invalid directive {d:?}: {e}"));
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn firehose_directives_include_acp_update_targets() {
|
||||
let directives = firehose_directives();
|
||||
assert!(directives.contains(&format!("{ACP_UPDATE_TARGET}=debug")));
|
||||
assert!(!directives.contains(&format!("{ACP_UPDATE_PAYLOAD_TARGET}=debug")));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn sanitize_key_replaces_path_separators_and_dot_only() {
|
||||
assert_eq!(sanitize_key("01923-abcd-EF"), "01923-abcd-EF");
|
||||
assert_eq!(sanitize_key("../escape"), ".._escape");
|
||||
assert_eq!(sanitize_key("a/b\\c"), "a_b_c");
|
||||
// Dot-only / empty keys collapse to a safe constant.
|
||||
for dotty in ["", ".", "..", "..."] {
|
||||
assert_eq!(sanitize_key(dotty), "_", "expected '_' for {dotty:?}");
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn layer_routes_event_by_session_span() {
|
||||
use tracing_subscriber::layer::SubscriberExt as _;
|
||||
|
||||
let _lock = flush_test_lock();
|
||||
let dir = tempfile::tempdir().unwrap();
|
||||
let layer = RoutingLayer::new(dir.path().to_path_buf(), "agent".to_owned(), 4242);
|
||||
let subscriber = tracing_subscriber::registry().with(layer);
|
||||
|
||||
tracing::subscriber::with_default(subscriber, || {
|
||||
// `%` mirrors production's `info_span!("session", session_id = %...)`.
|
||||
tracing::info_span!("session", session_id = %"sess-xyz").in_scope(|| {
|
||||
tracing::info!(target: "xai_grok_shell", "inside session");
|
||||
});
|
||||
tracing::info!(target: "xai_grok_shell", "outside session");
|
||||
});
|
||||
crate::appender::flush_file_log_guards();
|
||||
|
||||
let session_file = std::fs::read_to_string(dir.path().join("sess-xyz.txt")).unwrap();
|
||||
assert!(
|
||||
session_file.contains("inside session"),
|
||||
"session file: {session_file:?}"
|
||||
);
|
||||
assert!(!session_file.contains("outside session"));
|
||||
|
||||
let fallback = std::fs::read_to_string(dir.path().join("agent-4242.txt")).unwrap();
|
||||
assert!(
|
||||
fallback.contains("outside session"),
|
||||
"fallback file: {fallback:?}"
|
||||
);
|
||||
assert!(!fallback.contains("inside session"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn layer_routes_under_real_firehose_filter() {
|
||||
use tracing_subscriber::layer::SubscriberExt as _;
|
||||
|
||||
let _lock = flush_test_lock();
|
||||
let dir = tempfile::tempdir().unwrap();
|
||||
// Exactly the production wrapper: routing layer behind FIREHOSE_DIRECTIVES.
|
||||
// Pins the linchpin invariant — the `session` span (INFO, target
|
||||
// `xai_grok_telemetry::session_ctx`) survives the real filter so
|
||||
// `event_scope` still finds it — at the unit level.
|
||||
let layer = RoutingLayer::new(dir.path().to_path_buf(), "agent".to_owned(), 7)
|
||||
.with_filter(firehose_filter());
|
||||
let subscriber = tracing_subscriber::registry().with(layer);
|
||||
|
||||
tracing::subscriber::with_default(subscriber, || {
|
||||
tracing::info_span!(
|
||||
target: "xai_grok_telemetry::session_ctx",
|
||||
"session",
|
||||
session_id = %"sid-real"
|
||||
)
|
||||
.in_scope(|| {
|
||||
tracing::debug!(target: "xai_grok_shell", "filtered routing works");
|
||||
});
|
||||
});
|
||||
crate::appender::flush_file_log_guards();
|
||||
|
||||
let session_file = std::fs::read_to_string(dir.path().join("sid-real.txt")).unwrap();
|
||||
assert!(
|
||||
session_file.contains("filtered routing works"),
|
||||
"session file under real filter: {session_file:?}"
|
||||
);
|
||||
// Must route to the session file, NOT silently fall back to per-pid.
|
||||
assert!(
|
||||
!dir.path().join("agent-7.txt").exists(),
|
||||
"event must route to the session file, not the fallback"
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn layer_routes_two_sessions_to_distinct_files() {
|
||||
use tracing_subscriber::layer::SubscriberExt as _;
|
||||
|
||||
let _lock = flush_test_lock();
|
||||
let dir = tempfile::tempdir().unwrap();
|
||||
let layer = RoutingLayer::new(dir.path().to_path_buf(), "agent".to_owned(), 1);
|
||||
let subscriber = tracing_subscriber::registry().with(layer);
|
||||
|
||||
tracing::subscriber::with_default(subscriber, || {
|
||||
tracing::info_span!("session", session_id = %"sid-one").in_scope(|| {
|
||||
// Two events in one session also prove within-session accumulation.
|
||||
tracing::info!(target: "xai_grok_shell", "one first");
|
||||
tracing::info!(target: "xai_grok_shell", "one second");
|
||||
});
|
||||
tracing::info_span!("session", session_id = %"sid-two").in_scope(|| {
|
||||
tracing::info!(target: "xai_grok_shell", "two only");
|
||||
});
|
||||
});
|
||||
crate::appender::flush_file_log_guards();
|
||||
|
||||
let one = std::fs::read_to_string(dir.path().join("sid-one.txt")).unwrap();
|
||||
let two = std::fs::read_to_string(dir.path().join("sid-two.txt")).unwrap();
|
||||
assert!(
|
||||
one.contains("one first") && one.contains("one second") && !one.contains("two only"),
|
||||
"sid-one.txt: {one:?}"
|
||||
);
|
||||
assert!(
|
||||
two.contains("two only") && !two.contains("one first"),
|
||||
"sid-two.txt: {two:?}"
|
||||
);
|
||||
}
|
||||
|
||||
#[cfg(unix)]
|
||||
#[test]
|
||||
fn opening_session_sink_points_latest_symlink() {
|
||||
let _lock = flush_test_lock();
|
||||
let dir = tempfile::tempdir().unwrap();
|
||||
let layer = RoutingLayer::new(dir.path().to_path_buf(), "agent".to_owned(), 1);
|
||||
|
||||
layer.write_session("sess-1", b"x\n");
|
||||
let link = dir.path().join("latest.txt");
|
||||
assert_eq!(std::fs::read_link(&link).unwrap(), Path::new("sess-1.txt"));
|
||||
|
||||
// Opening a second session repoints latest.txt at it.
|
||||
layer.write_session("sess-2", b"y\n");
|
||||
assert_eq!(std::fs::read_link(&link).unwrap(), Path::new("sess-2.txt"));
|
||||
|
||||
crate::appender::flush_file_log_guards();
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn prune_old_logs_removes_old_keeps_recent_and_spares_nonmatching_and_latest() {
|
||||
use std::time::{Duration, SystemTime};
|
||||
|
||||
let dir = tempfile::tempdir().unwrap();
|
||||
let now = SystemTime::now();
|
||||
let max_age = Duration::from_secs(7 * 24 * 60 * 60);
|
||||
|
||||
let old = std::fs::File::create(dir.path().join("old-session.txt")).unwrap();
|
||||
old.set_modified(now - Duration::from_secs(8 * 24 * 60 * 60))
|
||||
.unwrap();
|
||||
let recent = std::fs::File::create(dir.path().join("recent-session.txt")).unwrap();
|
||||
recent.set_modified(now).unwrap();
|
||||
// A non-.txt file must be left untouched even if it is old.
|
||||
let other = std::fs::File::create(dir.path().join("unified.jsonl")).unwrap();
|
||||
other
|
||||
.set_modified(now - Duration::from_secs(30 * 24 * 60 * 60))
|
||||
.unwrap();
|
||||
// An old `latest.txt` must be spared (harmless stale link / sentinel).
|
||||
let latest = std::fs::File::create(dir.path().join("latest.txt")).unwrap();
|
||||
latest
|
||||
.set_modified(now - Duration::from_secs(30 * 24 * 60 * 60))
|
||||
.unwrap();
|
||||
// Orphaned `latest.txt` swap temps follow the same age rule: old reaped,
|
||||
// recent spared. Regular files here so mtimes are settable cross-platform;
|
||||
// the symlink-specific path is covered by the Unix-gated test below.
|
||||
let old_tmp =
|
||||
std::fs::File::create(dir.path().join(".latest.old-session.txt.tmp")).unwrap();
|
||||
old_tmp
|
||||
.set_modified(now - Duration::from_secs(8 * 24 * 60 * 60))
|
||||
.unwrap();
|
||||
let recent_tmp =
|
||||
std::fs::File::create(dir.path().join(".latest.recent-session.txt.tmp")).unwrap();
|
||||
recent_tmp.set_modified(now).unwrap();
|
||||
|
||||
prune_old_logs(dir.path(), max_age);
|
||||
|
||||
assert!(!dir.path().join("old-session.txt").exists());
|
||||
assert!(dir.path().join("recent-session.txt").exists());
|
||||
assert!(dir.path().join("unified.jsonl").exists());
|
||||
assert!(dir.path().join("latest.txt").exists());
|
||||
assert!(!dir.path().join(".latest.old-session.txt.tmp").exists());
|
||||
assert!(dir.path().join(".latest.recent-session.txt.tmp").exists());
|
||||
}
|
||||
|
||||
#[cfg(unix)]
|
||||
#[test]
|
||||
fn prune_old_logs_reaps_dangling_orphaned_latest_tmp_symlink() {
|
||||
use std::time::{Duration, SystemTime};
|
||||
|
||||
// Models the real orphan: a crash between `update_latest_symlink`'s
|
||||
// create and rename leaves the temp symlink, and its target session file
|
||||
// may itself be pruned later — so the link is dangling. Literal name (not
|
||||
// the consts) so renaming the scheme can't silently strand orphans
|
||||
// created by already-shipped binaries.
|
||||
let dir = tempfile::tempdir().unwrap();
|
||||
let max_age = Duration::from_secs(7 * 24 * 60 * 60);
|
||||
// `filetime` ages the link itself; std's `set_modified` follows it (and a
|
||||
// dangling link can't even be opened).
|
||||
let old = filetime::FileTime::from_system_time(
|
||||
SystemTime::now() - Duration::from_secs(8 * 24 * 60 * 60),
|
||||
);
|
||||
|
||||
let tmp = dir.path().join(".latest.gone-session.txt.tmp");
|
||||
std::os::unix::fs::symlink("gone-session.txt", &tmp).unwrap();
|
||||
filetime::set_symlink_file_times(&tmp, old, old).unwrap();
|
||||
// A just-created (mid-swap) temp must be spared by age.
|
||||
let fresh_tmp = dir.path().join(".latest.live-session.txt.tmp");
|
||||
std::os::unix::fs::symlink("live-session.txt", &fresh_tmp).unwrap();
|
||||
// `latest.txt` must stay spared by name even as an old dangling symlink.
|
||||
let latest = dir.path().join("latest.txt");
|
||||
std::os::unix::fs::symlink("gone-session.txt", &latest).unwrap();
|
||||
filetime::set_symlink_file_times(&latest, old, old).unwrap();
|
||||
|
||||
prune_old_logs(dir.path(), max_age);
|
||||
|
||||
// `Path::exists` follows symlinks (false for dangling links either way),
|
||||
// so assert on the links themselves via `symlink_metadata`.
|
||||
assert!(
|
||||
std::fs::symlink_metadata(&tmp).is_err(),
|
||||
"old orphaned dangling temp symlink must be pruned"
|
||||
);
|
||||
assert!(
|
||||
std::fs::symlink_metadata(&fresh_tmp).is_ok(),
|
||||
"fresh mid-swap temp must be spared by age"
|
||||
);
|
||||
assert!(
|
||||
std::fs::symlink_metadata(&latest).is_ok(),
|
||||
"latest.txt must be spared"
|
||||
);
|
||||
}
|
||||
|
||||
#[cfg(unix)]
|
||||
#[test]
|
||||
fn update_latest_symlink_failed_rename_removes_temp() {
|
||||
// Sanity: prove the temp symlink is creatable here, so the helper's
|
||||
// symlink step must succeed and the post-call absence below can only
|
||||
// come from the rename-failure cleanup branch.
|
||||
let dir = tempfile::tempdir().unwrap();
|
||||
let tmp = dir.path().join(".latest.sess.txt.tmp");
|
||||
std::os::unix::fs::symlink("sess.txt", &tmp).unwrap();
|
||||
std::fs::remove_file(&tmp).unwrap();
|
||||
// Force the rename to fail: a non-empty directory at `latest.txt` makes
|
||||
// rename(2) of a non-directory over it error (EISDIR/ENOTEMPTY).
|
||||
let blocker = dir.path().join("latest.txt");
|
||||
std::fs::create_dir(&blocker).unwrap();
|
||||
std::fs::File::create(blocker.join("occupant.txt")).unwrap();
|
||||
|
||||
update_latest_symlink(dir.path(), &dir.path().join("sess.txt"));
|
||||
|
||||
assert!(
|
||||
std::fs::symlink_metadata(&tmp).is_err(),
|
||||
"failed swap must remove the temp symlink, not orphan it"
|
||||
);
|
||||
assert!(
|
||||
blocker.join("occupant.txt").exists(),
|
||||
"rename must have failed, leaving the blocker dir untouched"
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn prune_old_logs_spares_active_logs_regardless_of_count() {
|
||||
use std::time::{Duration, SystemTime};
|
||||
|
||||
// Guards the reported bug: a concurrent process's still-open (recently
|
||||
// written) log must never be unlinked, however many newer logs exist.
|
||||
let dir = tempfile::tempdir().unwrap();
|
||||
let now = SystemTime::now();
|
||||
for i in 0..25 {
|
||||
let f = std::fs::File::create(dir.path().join(format!("cli-{i}.txt"))).unwrap();
|
||||
f.set_modified(now).unwrap();
|
||||
}
|
||||
|
||||
prune_old_logs(dir.path(), Duration::from_secs(7 * 24 * 60 * 60));
|
||||
|
||||
let count = std::fs::read_dir(dir.path())
|
||||
.unwrap()
|
||||
.flatten()
|
||||
.filter(|e| e.file_name().to_str().is_some_and(|n| n.ends_with(".txt")))
|
||||
.count();
|
||||
assert_eq!(count, 25);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn prune_old_logs_missing_dir_is_noop() {
|
||||
// Best-effort: a nonexistent debug dir must not panic.
|
||||
prune_old_logs(
|
||||
Path::new("/no/such/grok/debug/dir"),
|
||||
std::time::Duration::from_secs(1),
|
||||
);
|
||||
}
|
||||
}
|
||||
64
crates/codegen/xai-grok-telemetry/src/enums.rs
Normal file
64
crates/codegen/xai-grok-telemetry/src/enums.rs
Normal file
|
|
@ -0,0 +1,64 @@
|
|||
//! Shared telemetry/config enums extracted from shell.
|
||||
//!
|
||||
//! These were originally defined inside `xai-grok-shell` (in
|
||||
//! `session::mcp_servers` and `util::config`) but are referenced by
|
||||
//! telemetry payload structs in this crate, so they live here and shell
|
||||
//! re-exports them from their original paths to keep callers unchanged.
|
||||
|
||||
/// MCP initialization strategy
|
||||
#[derive(Debug, Clone, Copy, PartialEq, serde::Serialize, Default)]
|
||||
#[serde(rename_all = "snake_case")]
|
||||
pub enum McpInitStrategy {
|
||||
/// Wait for MCP initialization before first LLM call
|
||||
#[default]
|
||||
Blocking,
|
||||
/// Start immediately, advertise tools as they become available
|
||||
Progressive,
|
||||
}
|
||||
|
||||
impl<S: AsRef<str>> From<S> for McpInitStrategy {
|
||||
fn from(s: S) -> Self {
|
||||
match s.as_ref() {
|
||||
"progressive" => McpInitStrategy::Progressive,
|
||||
_ => McpInitStrategy::Blocking,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// How a PR creation was performed. Shared between the shell's session
|
||||
/// signals (`turn_result.json`) and the `pr_created` telemetry event.
|
||||
#[derive(Debug, Clone, Copy, PartialEq, Eq, serde::Serialize, serde::Deserialize)]
|
||||
#[serde(rename_all = "snake_case")]
|
||||
pub enum PrCreationSource {
|
||||
/// `gh pr create` via the bash tool.
|
||||
Bash,
|
||||
/// An MCP `create_pull_request` tool.
|
||||
Mcp,
|
||||
}
|
||||
|
||||
/// How the agent handles tool execution permissions.
|
||||
#[derive(Debug, Clone, Copy, PartialEq, Eq, Default, serde::Serialize)]
|
||||
#[serde(rename_all = "snake_case")]
|
||||
pub enum PermissionMode {
|
||||
/// Prompt the user for each tool call (default).
|
||||
#[default]
|
||||
Ask,
|
||||
/// Approve everything without prompting.
|
||||
AlwaysApprove,
|
||||
/// LLM transcript classifier reviews non-fast-path tool calls.
|
||||
Auto,
|
||||
}
|
||||
|
||||
impl PermissionMode {
|
||||
pub fn is_always_approve(self) -> bool {
|
||||
matches!(self, Self::AlwaysApprove)
|
||||
}
|
||||
|
||||
pub fn is_auto(self) -> bool {
|
||||
matches!(self, Self::Auto)
|
||||
}
|
||||
|
||||
pub fn from_yolo(yolo: bool) -> Self {
|
||||
if yolo { Self::AlwaysApprove } else { Self::Ask }
|
||||
}
|
||||
}
|
||||
2096
crates/codegen/xai-grok-telemetry/src/events.rs
Normal file
2096
crates/codegen/xai-grok-telemetry/src/events.rs
Normal file
File diff suppressed because it is too large
Load diff
743
crates/codegen/xai-grok-telemetry/src/external/config.rs
vendored
Normal file
743
crates/codegen/xai-grok-telemetry/src/external/config.rs
vendored
Normal file
|
|
@ -0,0 +1,743 @@
|
|||
//! Configuration resolution for the external OTEL stream.
|
||||
//!
|
||||
//! Pure resolution — no I/O besides reading env vars. The shell resolves the
|
||||
//! startup value once (layering the `[telemetry]` `otel_*` config keys under
|
||||
//! the env vars) and passes the resolved struct to [`crate::external::init`].
|
||||
//!
|
||||
//! Activation requires a **double opt-in** (user-confirmed, RQ7):
|
||||
//! `GROK_EXTERNAL_OTEL=1` *and* at least one of `OTEL_METRICS_EXPORTER` /
|
||||
//! `OTEL_LOGS_EXPORTER` set to a real exporter. The master switch alone
|
||||
//! enables nothing; the exporter vars alone enable nothing.
|
||||
|
||||
use std::time::Duration;
|
||||
|
||||
/// OTLP transport/protocol for external exporters.
|
||||
#[derive(Debug, Clone, Copy, PartialEq, Eq, Default)]
|
||||
pub enum OtlpTransport {
|
||||
/// OTLP over HTTP with protobuf bodies.
|
||||
#[default]
|
||||
HttpProtobuf,
|
||||
/// OTLP over gRPC/protobuf.
|
||||
Grpc,
|
||||
}
|
||||
|
||||
impl OtlpTransport {
|
||||
fn parse(raw: &str) -> Option<Self> {
|
||||
match raw.trim().to_ascii_lowercase().as_str() {
|
||||
"http/protobuf" | "http-protobuf" | "http" => Some(Self::HttpProtobuf),
|
||||
"grpc" => Some(Self::Grpc),
|
||||
"" => Some(Self::HttpProtobuf),
|
||||
_ => None,
|
||||
}
|
||||
}
|
||||
|
||||
pub fn as_protocol_str(self) -> &'static str {
|
||||
match self {
|
||||
Self::HttpProtobuf => "http/protobuf",
|
||||
Self::Grpc => "grpc",
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Master switch env var. Deliberately *not* `GROK_ENABLE_TELEMETRY`: that
|
||||
/// would be a word-order typo away from the long-standing
|
||||
/// `GROK_TELEMETRY_ENABLED` (product events/Mixpanel mode), and the two control
|
||||
/// opposite-pointing data flows (to xAI vs. to the customer's collector).
|
||||
pub const ENV_MASTER_SWITCH: &str = "GROK_EXTERNAL_OTEL";
|
||||
|
||||
/// Exporter selection for one signal (`OTEL_METRICS_EXPORTER` /
|
||||
/// `OTEL_LOGS_EXPORTER`).
|
||||
#[derive(Debug, Clone, Copy, PartialEq, Eq, Default)]
|
||||
pub enum ExporterSelection {
|
||||
/// No exporter — the signal is not produced.
|
||||
#[default]
|
||||
None,
|
||||
/// OTLP to the configured endpoint using [`OtlpTransport`].
|
||||
Otlp,
|
||||
/// Redacted records printed to **stderr** (debugging). Stdout protocol
|
||||
/// channels (headless/stream-JSON) are never touched.
|
||||
Console,
|
||||
}
|
||||
|
||||
impl ExporterSelection {
|
||||
fn parse(raw: &str) -> Option<Self> {
|
||||
match raw.trim().to_ascii_lowercase().as_str() {
|
||||
"otlp" => Some(Self::Otlp),
|
||||
"console" => Some(Self::Console),
|
||||
"none" | "" => Some(Self::None),
|
||||
_ => None,
|
||||
}
|
||||
}
|
||||
|
||||
/// `true` for any selection that produces output.
|
||||
pub fn is_active(self) -> bool {
|
||||
!matches!(self, Self::None)
|
||||
}
|
||||
}
|
||||
|
||||
/// Content gates (additive opt-ins; default off). May only **tighten**
|
||||
/// post-init — a remote policy can force them off, never on.
|
||||
#[derive(Debug, Clone, Copy, Default, PartialEq, Eq)]
|
||||
pub struct ContentGates {
|
||||
/// `OTEL_LOG_USER_PROMPTS=1`: prompt text on `grok_code.user_prompt`
|
||||
/// (60 KB cap, secret-scrubbed).
|
||||
pub log_user_prompts: bool,
|
||||
/// `OTEL_LOG_TOOL_DETAILS=1`: gated tool params / full paths / verbatim
|
||||
/// MCP, skill, and plugin names.
|
||||
pub log_tool_details: bool,
|
||||
}
|
||||
|
||||
/// Delta vs. cumulative metric temporality
|
||||
/// (`OTEL_EXPORTER_OTLP_METRICS_TEMPORALITY_PREFERENCE`). Default **Delta**.
|
||||
#[derive(Debug, Clone, Copy, PartialEq, Eq, Default)]
|
||||
pub enum TemporalityPreference {
|
||||
#[default]
|
||||
Delta,
|
||||
Cumulative,
|
||||
}
|
||||
|
||||
/// Identity of the binary emitting external telemetry; becomes resource
|
||||
/// attributes. Filled by the caller (pager/shell) at init.
|
||||
#[derive(Debug, Clone, Default)]
|
||||
pub struct ExternalClientInfo {
|
||||
/// Engine build (version + commit) → `service.version`.
|
||||
pub service_version: String,
|
||||
/// Front-end client version → `client.version`.
|
||||
pub client_version: String,
|
||||
/// How the session was launched (`cli`/`headless`/`agent`) →
|
||||
/// `app.entrypoint`.
|
||||
pub app_entrypoint: String,
|
||||
}
|
||||
|
||||
/// Config-file layer for the external stream, built by the shell from the
|
||||
/// `otel_*` keys of the `[telemetry]` table and layered *under* env vars
|
||||
/// during resolution. (Field names here are the internal carrier; the
|
||||
/// user-facing keys are `otel_enabled`, `otel_metrics_exporter`, … — see
|
||||
/// [`crate::config::TelemetryConfig`].)
|
||||
///
|
||||
/// There is deliberately **no `headers` key** (user decision, RQ4): collector
|
||||
/// auth is supplied via the `OTEL_EXPORTER_OTLP_HEADERS` env var only, so
|
||||
/// collector tokens are never stored on disk.
|
||||
#[derive(Debug, Clone, Default, serde::Deserialize, serde::Serialize)]
|
||||
#[serde(default)]
|
||||
pub struct ExternalOtelFileConfig {
|
||||
/// `= GROK_EXTERNAL_OTEL` (env wins).
|
||||
pub enabled: Option<bool>,
|
||||
/// `otlp` | `console` | `none`.
|
||||
pub metrics_exporter: Option<String>,
|
||||
/// `otlp` | `console` | `none`.
|
||||
pub logs_exporter: Option<String>,
|
||||
/// OTLP base endpoint (`/v1/logs`, `/v1/metrics` appended per spec for HTTP).
|
||||
pub endpoint: Option<String>,
|
||||
/// `http/protobuf` | `grpc`.
|
||||
pub protocol: Option<String>,
|
||||
/// Content gate (admins can pin this to `false` via requirements).
|
||||
pub log_user_prompts: Option<bool>,
|
||||
/// Content gate (admins can pin this to `false` via requirements).
|
||||
pub log_tool_details: Option<bool>,
|
||||
}
|
||||
|
||||
/// Fully resolved configuration for the external stream. Returned by
|
||||
/// [`ExternalOtelConfig::resolve`] only when the double opt-in is satisfied;
|
||||
/// `None` means the module is never constructed (zero allocation, zero
|
||||
/// threads, zero sockets).
|
||||
#[derive(Debug, Clone)]
|
||||
pub struct ExternalOtelConfig {
|
||||
pub metrics_exporter: ExporterSelection,
|
||||
pub logs_exporter: ExporterSelection,
|
||||
pub transport: OtlpTransport,
|
||||
/// Resolved logs endpoint (full `…/v1/logs` for HTTP; collector origin for gRPC).
|
||||
pub logs_endpoint: String,
|
||||
/// Resolved metrics endpoint (full `…/v1/metrics` for HTTP; collector origin for gRPC).
|
||||
pub metrics_endpoint: String,
|
||||
/// Customer collector headers for log exports, parsed from
|
||||
/// `OTEL_EXPORTER_OTLP_HEADERS` plus `OTEL_EXPORTER_OTLP_LOGS_HEADERS`.
|
||||
/// The **only** headers the external log exporter ever sends.
|
||||
pub logs_headers: Vec<(String, String)>,
|
||||
/// Customer collector headers for metric exports, parsed from
|
||||
/// `OTEL_EXPORTER_OTLP_HEADERS` plus `OTEL_EXPORTER_OTLP_METRICS_HEADERS`.
|
||||
/// The **only** headers the external metric exporter ever sends.
|
||||
pub metrics_headers: Vec<(String, String)>,
|
||||
/// `OTEL_EXPORTER_OTLP_TIMEOUT` (ms). Default 10 s.
|
||||
pub timeout: Duration,
|
||||
/// `OTEL_METRIC_EXPORT_INTERVAL` (ms). Default 60 s.
|
||||
pub metric_export_interval: Duration,
|
||||
/// `OTEL_BLRP_SCHEDULE_DELAY` (spec name, wins) /
|
||||
/// `OTEL_LOGS_EXPORT_INTERVAL` (compatibility alias). Default 5 s.
|
||||
pub logs_export_interval: Duration,
|
||||
pub gates: ContentGates,
|
||||
pub temporality: TemporalityPreference,
|
||||
/// `OTEL_METRICS_INCLUDE_SESSION_ID` (default on): `session.id` on
|
||||
/// metrics (cardinality opt-out).
|
||||
pub include_session_id_on_metrics: bool,
|
||||
/// `OTEL_METRICS_INCLUDE_VERSION` (default off): `app.version` on
|
||||
/// metrics.
|
||||
pub include_version_on_metrics: bool,
|
||||
/// Resource identity, filled by the caller at init.
|
||||
pub client: ExternalClientInfo,
|
||||
/// Set by the shell when the **internal** firehose resolved its
|
||||
/// endpoint/headers from `OTEL_EXPORTER_OTLP_*` (the deprecated
|
||||
/// fallback). [`crate::external::init`] refuses to activate when true —
|
||||
/// the no-double-send invariant is enforced in code, not release
|
||||
/// discipline.
|
||||
pub internal_pipeline_consumed_otel_vars: bool,
|
||||
/// Which layer supplied the master switch (`"env"` | `"config"`), for the
|
||||
/// internal adoption meta-event. `remote` is not a possible startup
|
||||
/// source (init reads env + local config only).
|
||||
pub enabled_source: &'static str,
|
||||
}
|
||||
|
||||
fn env_bool(raw: &str) -> Option<bool> {
|
||||
match raw.trim().to_ascii_lowercase().as_str() {
|
||||
"1" | "true" | "yes" | "on" => Some(true),
|
||||
"0" | "false" | "no" | "off" | "" => Some(false),
|
||||
_ => None,
|
||||
}
|
||||
}
|
||||
|
||||
fn parse_ms(raw: Option<String>, default: Duration) -> Duration {
|
||||
raw.and_then(|s| s.trim().parse::<u64>().ok())
|
||||
.map(Duration::from_millis)
|
||||
.unwrap_or(default)
|
||||
}
|
||||
|
||||
/// Parse `k=v,k2=v2` header lists (OTLP env spec); blank keys skipped.
|
||||
pub fn parse_header_list(raw: &str) -> Vec<(String, String)> {
|
||||
raw.split(',')
|
||||
.filter_map(|kv| {
|
||||
let (k, v) = kv.split_once('=')?;
|
||||
let k = k.trim();
|
||||
(!k.is_empty()).then(|| (k.to_string(), v.trim().to_string()))
|
||||
})
|
||||
.collect()
|
||||
}
|
||||
|
||||
/// OTLP HTTP default base endpoint per spec.
|
||||
const DEFAULT_OTLP_HTTP_BASE: &str = "http://localhost:4318";
|
||||
/// OTLP gRPC default endpoint per spec.
|
||||
const DEFAULT_OTLP_GRPC_ENDPOINT: &str = "http://localhost:4317";
|
||||
|
||||
fn resolve_signal_endpoint(
|
||||
signal_specific: Option<String>,
|
||||
base: Option<&str>,
|
||||
path: &str,
|
||||
transport: OtlpTransport,
|
||||
) -> String {
|
||||
if let Some(full) = signal_specific.filter(|s| !s.trim().is_empty()) {
|
||||
return full.trim().trim_end_matches('/').to_string();
|
||||
}
|
||||
let default_base = match transport {
|
||||
OtlpTransport::HttpProtobuf => DEFAULT_OTLP_HTTP_BASE,
|
||||
OtlpTransport::Grpc => DEFAULT_OTLP_GRPC_ENDPOINT,
|
||||
};
|
||||
let base = base
|
||||
.filter(|s| !s.trim().is_empty())
|
||||
.unwrap_or(default_base)
|
||||
.trim()
|
||||
.trim_end_matches('/');
|
||||
match transport {
|
||||
OtlpTransport::HttpProtobuf => format!("{base}/{path}"),
|
||||
OtlpTransport::Grpc => base.to_owned(),
|
||||
}
|
||||
}
|
||||
|
||||
impl ExternalOtelConfig {
|
||||
/// Resolve from process env layered over the optional `[telemetry]`
|
||||
/// `otel_*` config-file layer. Returns `None` unless the double opt-in is
|
||||
/// satisfied (master switch + at least one real exporter) and the
|
||||
/// transport is supported.
|
||||
pub fn resolve(file: Option<&ExternalOtelFileConfig>) -> Option<Self> {
|
||||
Self::resolve_with(|name| std::env::var(name).ok(), file)
|
||||
}
|
||||
|
||||
/// Testable resolution core: `getenv` abstracts `std::env::var` so tests
|
||||
/// don't race on process-global env state.
|
||||
pub fn resolve_with(
|
||||
getenv: impl Fn(&str) -> Option<String>,
|
||||
file: Option<&ExternalOtelFileConfig>,
|
||||
) -> Option<Self> {
|
||||
// Master switch: env > config file > default off.
|
||||
let (enabled, enabled_source) =
|
||||
match getenv(ENV_MASTER_SWITCH).as_deref().and_then(env_bool) {
|
||||
Some(v) => (v, "env"),
|
||||
None => match file.and_then(|f| f.enabled) {
|
||||
Some(v) => (v, "config"),
|
||||
None => (false, "env"),
|
||||
},
|
||||
};
|
||||
if !enabled {
|
||||
return None;
|
||||
}
|
||||
|
||||
let select = |env_name: &str, file_value: Option<&str>| -> ExporterSelection {
|
||||
let raw = getenv(env_name).or_else(|| file_value.map(str::to_owned));
|
||||
match raw.as_deref().map(ExporterSelection::parse) {
|
||||
Some(Some(sel)) => sel,
|
||||
Some(None) => {
|
||||
tracing::warn!(
|
||||
var = env_name,
|
||||
"external otel: unrecognized exporter selection; treating as `none`"
|
||||
);
|
||||
ExporterSelection::None
|
||||
}
|
||||
None => ExporterSelection::None,
|
||||
}
|
||||
};
|
||||
let metrics_exporter = select(
|
||||
"OTEL_METRICS_EXPORTER",
|
||||
file.and_then(|f| f.metrics_exporter.as_deref()),
|
||||
);
|
||||
let logs_exporter = select(
|
||||
"OTEL_LOGS_EXPORTER",
|
||||
file.and_then(|f| f.logs_exporter.as_deref()),
|
||||
);
|
||||
// Double opt-in (RQ7): the master switch alone enables nothing.
|
||||
if !metrics_exporter.is_active() && !logs_exporter.is_active() {
|
||||
return None;
|
||||
}
|
||||
|
||||
let raw_protocol =
|
||||
getenv("OTEL_EXPORTER_OTLP_PROTOCOL").or_else(|| file.and_then(|f| f.protocol.clone()));
|
||||
let transport = match raw_protocol.as_deref().map(OtlpTransport::parse) {
|
||||
Some(Some(transport)) => transport,
|
||||
Some(None) => {
|
||||
tracing::warn!(
|
||||
protocol = raw_protocol.as_deref().unwrap_or_default(),
|
||||
"external otel: unrecognized OTLP protocol; stream disabled"
|
||||
);
|
||||
return None;
|
||||
}
|
||||
None => OtlpTransport::HttpProtobuf,
|
||||
};
|
||||
|
||||
let base_endpoint = getenv("OTEL_EXPORTER_OTLP_ENDPOINT")
|
||||
.filter(|s| !s.trim().is_empty())
|
||||
.or_else(|| file.and_then(|f| f.endpoint.clone()))
|
||||
.filter(|s| !s.trim().is_empty());
|
||||
let logs_endpoint = resolve_signal_endpoint(
|
||||
getenv("OTEL_EXPORTER_OTLP_LOGS_ENDPOINT"),
|
||||
base_endpoint.as_deref(),
|
||||
"v1/logs",
|
||||
transport,
|
||||
);
|
||||
let metrics_endpoint = resolve_signal_endpoint(
|
||||
getenv("OTEL_EXPORTER_OTLP_METRICS_ENDPOINT"),
|
||||
base_endpoint.as_deref(),
|
||||
"v1/metrics",
|
||||
transport,
|
||||
);
|
||||
|
||||
// Headers: env only (RQ4) — never from the config file. Resolve them
|
||||
// per signal so signal-specific overrides never bleed across streams.
|
||||
let base_headers = parse_header_list(
|
||||
getenv("OTEL_EXPORTER_OTLP_HEADERS")
|
||||
.as_deref()
|
||||
.unwrap_or(""),
|
||||
);
|
||||
let resolve_signal_headers = |signal_var: &str| {
|
||||
let mut headers = base_headers.clone();
|
||||
if let Some(extra) = getenv(signal_var) {
|
||||
for (k, v) in parse_header_list(&extra) {
|
||||
if let Some(existing) = headers.iter_mut().find(|(ek, _)| *ek == k) {
|
||||
existing.1 = v;
|
||||
} else {
|
||||
headers.push((k, v));
|
||||
}
|
||||
}
|
||||
}
|
||||
headers
|
||||
};
|
||||
let logs_headers = resolve_signal_headers("OTEL_EXPORTER_OTLP_LOGS_HEADERS");
|
||||
let metrics_headers = resolve_signal_headers("OTEL_EXPORTER_OTLP_METRICS_HEADERS");
|
||||
|
||||
let gates = ContentGates {
|
||||
log_user_prompts: getenv("OTEL_LOG_USER_PROMPTS")
|
||||
.as_deref()
|
||||
.and_then(env_bool)
|
||||
.or_else(|| file.and_then(|f| f.log_user_prompts))
|
||||
.unwrap_or(false),
|
||||
log_tool_details: getenv("OTEL_LOG_TOOL_DETAILS")
|
||||
.as_deref()
|
||||
.and_then(env_bool)
|
||||
.or_else(|| file.and_then(|f| f.log_tool_details))
|
||||
.unwrap_or(false),
|
||||
};
|
||||
|
||||
let temporality = match getenv("OTEL_EXPORTER_OTLP_METRICS_TEMPORALITY_PREFERENCE")
|
||||
.map(|s| s.trim().to_ascii_lowercase())
|
||||
.as_deref()
|
||||
{
|
||||
Some("cumulative") => TemporalityPreference::Cumulative,
|
||||
// `delta`, `lowmemory`, unset, or unrecognized → Delta default.
|
||||
_ => TemporalityPreference::Delta,
|
||||
};
|
||||
|
||||
Some(Self {
|
||||
metrics_exporter,
|
||||
logs_exporter,
|
||||
transport,
|
||||
logs_endpoint,
|
||||
metrics_endpoint,
|
||||
logs_headers,
|
||||
metrics_headers,
|
||||
timeout: parse_ms(
|
||||
getenv("OTEL_EXPORTER_OTLP_TIMEOUT"),
|
||||
Duration::from_millis(10_000),
|
||||
),
|
||||
metric_export_interval: parse_ms(
|
||||
getenv("OTEL_METRIC_EXPORT_INTERVAL"),
|
||||
Duration::from_millis(60_000),
|
||||
),
|
||||
logs_export_interval: parse_ms(
|
||||
getenv("OTEL_BLRP_SCHEDULE_DELAY").or_else(|| getenv("OTEL_LOGS_EXPORT_INTERVAL")),
|
||||
Duration::from_millis(5_000),
|
||||
),
|
||||
gates,
|
||||
temporality,
|
||||
include_session_id_on_metrics: getenv("OTEL_METRICS_INCLUDE_SESSION_ID")
|
||||
.as_deref()
|
||||
.and_then(env_bool)
|
||||
.unwrap_or(true),
|
||||
include_version_on_metrics: getenv("OTEL_METRICS_INCLUDE_VERSION")
|
||||
.as_deref()
|
||||
.and_then(env_bool)
|
||||
.unwrap_or(false),
|
||||
client: ExternalClientInfo::default(),
|
||||
internal_pipeline_consumed_otel_vars: false,
|
||||
enabled_source,
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
use std::collections::HashMap;
|
||||
|
||||
fn env(pairs: &[(&str, &str)]) -> impl Fn(&str) -> Option<String> + use<> {
|
||||
let map: HashMap<String, String> = pairs
|
||||
.iter()
|
||||
.map(|(k, v)| (k.to_string(), v.to_string()))
|
||||
.collect();
|
||||
move |name| map.get(name).cloned()
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn default_is_off() {
|
||||
assert!(ExternalOtelConfig::resolve_with(env(&[]), None).is_none());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn master_switch_alone_enables_nothing() {
|
||||
// RQ7: GROK_EXTERNAL_OTEL=1 without an explicit exporter is inert.
|
||||
assert!(
|
||||
ExternalOtelConfig::resolve_with(env(&[("GROK_EXTERNAL_OTEL", "1")]), None).is_none()
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn exporters_alone_enable_nothing() {
|
||||
assert!(
|
||||
ExternalOtelConfig::resolve_with(env(&[("OTEL_METRICS_EXPORTER", "otlp")]), None)
|
||||
.is_none()
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn double_opt_in_activates() {
|
||||
let cfg = ExternalOtelConfig::resolve_with(
|
||||
env(&[
|
||||
("GROK_EXTERNAL_OTEL", "1"),
|
||||
("OTEL_METRICS_EXPORTER", "otlp"),
|
||||
]),
|
||||
None,
|
||||
)
|
||||
.expect("must activate");
|
||||
assert_eq!(cfg.metrics_exporter, ExporterSelection::Otlp);
|
||||
assert_eq!(cfg.logs_exporter, ExporterSelection::None);
|
||||
assert_eq!(cfg.metrics_endpoint, "http://localhost:4318/v1/metrics");
|
||||
assert_eq!(cfg.logs_endpoint, "http://localhost:4318/v1/logs");
|
||||
assert!(!cfg.gates.log_user_prompts);
|
||||
assert!(!cfg.gates.log_tool_details);
|
||||
assert!(cfg.include_session_id_on_metrics);
|
||||
assert!(!cfg.include_version_on_metrics);
|
||||
assert_eq!(cfg.temporality, TemporalityPreference::Delta);
|
||||
assert_eq!(cfg.transport, OtlpTransport::HttpProtobuf);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn grpc_protocol_accepted() {
|
||||
let cfg = ExternalOtelConfig::resolve_with(
|
||||
env(&[
|
||||
("GROK_EXTERNAL_OTEL", "1"),
|
||||
("OTEL_LOGS_EXPORTER", "otlp"),
|
||||
("OTEL_EXPORTER_OTLP_PROTOCOL", "grpc"),
|
||||
]),
|
||||
None,
|
||||
)
|
||||
.expect("grpc must activate");
|
||||
assert_eq!(cfg.transport, OtlpTransport::Grpc);
|
||||
assert_eq!(cfg.logs_endpoint, "http://localhost:4317");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn http_protobuf_protocol_accepted() {
|
||||
let cfg = ExternalOtelConfig::resolve_with(
|
||||
env(&[
|
||||
("GROK_EXTERNAL_OTEL", "1"),
|
||||
("OTEL_LOGS_EXPORTER", "otlp"),
|
||||
("OTEL_EXPORTER_OTLP_PROTOCOL", "http/protobuf"),
|
||||
]),
|
||||
None,
|
||||
);
|
||||
let cfg = cfg.unwrap();
|
||||
assert_eq!(cfg.transport, OtlpTransport::HttpProtobuf);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn unknown_protocol_disables() {
|
||||
let cfg = ExternalOtelConfig::resolve_with(
|
||||
env(&[
|
||||
("GROK_EXTERNAL_OTEL", "1"),
|
||||
("OTEL_LOGS_EXPORTER", "otlp"),
|
||||
("OTEL_EXPORTER_OTLP_PROTOCOL", "http/json"),
|
||||
]),
|
||||
None,
|
||||
);
|
||||
assert!(cfg.is_none(), "unknown protocols must disable the stream");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn endpoint_resolution_follows_otlp_http_spec() {
|
||||
let cfg = ExternalOtelConfig::resolve_with(
|
||||
env(&[
|
||||
("GROK_EXTERNAL_OTEL", "1"),
|
||||
("OTEL_LOGS_EXPORTER", "otlp"),
|
||||
("OTEL_METRICS_EXPORTER", "otlp"),
|
||||
(
|
||||
"OTEL_EXPORTER_OTLP_ENDPOINT",
|
||||
"https://collector.corp.example:4318/",
|
||||
),
|
||||
(
|
||||
"OTEL_EXPORTER_OTLP_LOGS_ENDPOINT",
|
||||
"https://logs.corp.example/custom",
|
||||
),
|
||||
]),
|
||||
None,
|
||||
)
|
||||
.unwrap();
|
||||
// Signal-specific endpoint used verbatim; base + spec path otherwise.
|
||||
assert_eq!(cfg.logs_endpoint, "https://logs.corp.example/custom");
|
||||
assert_eq!(
|
||||
cfg.metrics_endpoint,
|
||||
"https://collector.corp.example:4318/v1/metrics"
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn grpc_endpoint_resolution_uses_collector_endpoint_without_http_paths() {
|
||||
let cfg = ExternalOtelConfig::resolve_with(
|
||||
env(&[
|
||||
("GROK_EXTERNAL_OTEL", "1"),
|
||||
("OTEL_LOGS_EXPORTER", "otlp"),
|
||||
("OTEL_METRICS_EXPORTER", "otlp"),
|
||||
("OTEL_EXPORTER_OTLP_PROTOCOL", "grpc"),
|
||||
(
|
||||
"OTEL_EXPORTER_OTLP_ENDPOINT",
|
||||
"https://collector.corp.example:4317/",
|
||||
),
|
||||
]),
|
||||
None,
|
||||
)
|
||||
.unwrap();
|
||||
assert_eq!(cfg.logs_endpoint, "https://collector.corp.example:4317");
|
||||
assert_eq!(cfg.metrics_endpoint, "https://collector.corp.example:4317");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn file_protocol_layered_under_env() {
|
||||
let file = ExternalOtelFileConfig {
|
||||
enabled: Some(true),
|
||||
metrics_exporter: None,
|
||||
logs_exporter: Some("otlp".into()),
|
||||
endpoint: None,
|
||||
log_user_prompts: None,
|
||||
log_tool_details: None,
|
||||
protocol: Some("grpc".into()),
|
||||
};
|
||||
let cfg = ExternalOtelConfig::resolve_with(env(&[]), Some(&file)).unwrap();
|
||||
assert_eq!(cfg.transport, OtlpTransport::Grpc);
|
||||
|
||||
let cfg = ExternalOtelConfig::resolve_with(
|
||||
env(&[("OTEL_EXPORTER_OTLP_PROTOCOL", "http/protobuf")]),
|
||||
Some(&file),
|
||||
)
|
||||
.unwrap();
|
||||
assert_eq!(cfg.transport, OtlpTransport::HttpProtobuf);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn headers_parsed_and_signal_specific_scoped() {
|
||||
let cfg = ExternalOtelConfig::resolve_with(
|
||||
env(&[
|
||||
("GROK_EXTERNAL_OTEL", "1"),
|
||||
("OTEL_LOGS_EXPORTER", "otlp"),
|
||||
("OTEL_EXPORTER_OTLP_HEADERS", "x-token=abc, x-org=corp"),
|
||||
("OTEL_EXPORTER_OTLP_LOGS_HEADERS", "x-token=override"),
|
||||
]),
|
||||
None,
|
||||
)
|
||||
.unwrap();
|
||||
assert_eq!(
|
||||
cfg.logs_headers,
|
||||
vec![
|
||||
("x-token".to_string(), "override".to_string()),
|
||||
("x-org".to_string(), "corp".to_string()),
|
||||
]
|
||||
);
|
||||
assert_eq!(
|
||||
cfg.metrics_headers,
|
||||
vec![
|
||||
("x-token".to_string(), "abc".to_string()),
|
||||
("x-org".to_string(), "corp".to_string()),
|
||||
]
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn logs_and_metrics_headers_stay_isolated() {
|
||||
let cfg = ExternalOtelConfig::resolve_with(
|
||||
env(&[
|
||||
("GROK_EXTERNAL_OTEL", "1"),
|
||||
("OTEL_LOGS_EXPORTER", "otlp"),
|
||||
("OTEL_METRICS_EXPORTER", "otlp"),
|
||||
("OTEL_EXPORTER_OTLP_HEADERS", "authorization=Bearer base"),
|
||||
(
|
||||
"OTEL_EXPORTER_OTLP_LOGS_HEADERS",
|
||||
"authorization=Bearer logs",
|
||||
),
|
||||
(
|
||||
"OTEL_EXPORTER_OTLP_METRICS_HEADERS",
|
||||
"authorization=Bearer metrics",
|
||||
),
|
||||
]),
|
||||
None,
|
||||
)
|
||||
.unwrap();
|
||||
assert_eq!(
|
||||
cfg.logs_headers,
|
||||
vec![("authorization".to_string(), "Bearer logs".to_string())]
|
||||
);
|
||||
assert_eq!(
|
||||
cfg.metrics_headers,
|
||||
vec![("authorization".to_string(), "Bearer metrics".to_string())]
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn content_gates_default_off_env_enables() {
|
||||
let cfg = ExternalOtelConfig::resolve_with(
|
||||
env(&[
|
||||
("GROK_EXTERNAL_OTEL", "1"),
|
||||
("OTEL_LOGS_EXPORTER", "otlp"),
|
||||
("OTEL_LOG_USER_PROMPTS", "1"),
|
||||
("OTEL_LOG_TOOL_DETAILS", "true"),
|
||||
]),
|
||||
None,
|
||||
)
|
||||
.unwrap();
|
||||
assert!(cfg.gates.log_user_prompts);
|
||||
assert!(cfg.gates.log_tool_details);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn intervals_and_timeout_parsed_with_blrp_precedence() {
|
||||
let cfg = ExternalOtelConfig::resolve_with(
|
||||
env(&[
|
||||
("GROK_EXTERNAL_OTEL", "1"),
|
||||
("OTEL_METRICS_EXPORTER", "otlp"),
|
||||
("OTEL_EXPORTER_OTLP_TIMEOUT", "2500"),
|
||||
("OTEL_METRIC_EXPORT_INTERVAL", "30000"),
|
||||
("OTEL_BLRP_SCHEDULE_DELAY", "1000"),
|
||||
("OTEL_LOGS_EXPORT_INTERVAL", "9999"),
|
||||
]),
|
||||
None,
|
||||
)
|
||||
.unwrap();
|
||||
assert_eq!(cfg.timeout, Duration::from_millis(2500));
|
||||
assert_eq!(cfg.metric_export_interval, Duration::from_millis(30_000));
|
||||
// Spec name wins over the compatibility alias.
|
||||
assert_eq!(cfg.logs_export_interval, Duration::from_millis(1000));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn logs_export_interval_alias_honored_when_spec_name_absent() {
|
||||
let cfg = ExternalOtelConfig::resolve_with(
|
||||
env(&[
|
||||
("GROK_EXTERNAL_OTEL", "1"),
|
||||
("OTEL_METRICS_EXPORTER", "otlp"),
|
||||
("OTEL_LOGS_EXPORT_INTERVAL", "9999"),
|
||||
]),
|
||||
None,
|
||||
)
|
||||
.unwrap();
|
||||
assert_eq!(cfg.logs_export_interval, Duration::from_millis(9999));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn file_config_layered_under_env() {
|
||||
let file = ExternalOtelFileConfig {
|
||||
enabled: Some(true),
|
||||
metrics_exporter: Some("otlp".into()),
|
||||
logs_exporter: Some("console".into()),
|
||||
endpoint: Some("https://file.example:4318".into()),
|
||||
log_user_prompts: Some(true),
|
||||
log_tool_details: None,
|
||||
protocol: None,
|
||||
};
|
||||
// No env at all: file config alone activates.
|
||||
let cfg = ExternalOtelConfig::resolve_with(env(&[]), Some(&file)).unwrap();
|
||||
assert_eq!(cfg.metrics_exporter, ExporterSelection::Otlp);
|
||||
assert_eq!(cfg.logs_exporter, ExporterSelection::Console);
|
||||
assert_eq!(cfg.metrics_endpoint, "https://file.example:4318/v1/metrics");
|
||||
assert!(cfg.gates.log_user_prompts);
|
||||
|
||||
// Env wins over file on every layered key.
|
||||
let cfg = ExternalOtelConfig::resolve_with(
|
||||
env(&[
|
||||
("OTEL_METRICS_EXPORTER", "none"),
|
||||
("OTEL_LOGS_EXPORTER", "otlp"),
|
||||
("OTEL_LOG_USER_PROMPTS", "0"),
|
||||
("OTEL_EXPORTER_OTLP_ENDPOINT", "https://env.example:4318"),
|
||||
]),
|
||||
Some(&file),
|
||||
)
|
||||
.unwrap();
|
||||
assert_eq!(cfg.metrics_exporter, ExporterSelection::None);
|
||||
assert_eq!(cfg.logs_exporter, ExporterSelection::Otlp);
|
||||
assert_eq!(cfg.logs_endpoint, "https://env.example:4318/v1/logs");
|
||||
assert!(!cfg.gates.log_user_prompts);
|
||||
|
||||
// Env master switch off wins over file `enabled = true`.
|
||||
let cfg =
|
||||
ExternalOtelConfig::resolve_with(env(&[("GROK_EXTERNAL_OTEL", "0")]), Some(&file));
|
||||
assert!(cfg.is_none());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn cumulative_temporality_honored() {
|
||||
let cfg = ExternalOtelConfig::resolve_with(
|
||||
env(&[
|
||||
("GROK_EXTERNAL_OTEL", "1"),
|
||||
("OTEL_METRICS_EXPORTER", "otlp"),
|
||||
(
|
||||
"OTEL_EXPORTER_OTLP_METRICS_TEMPORALITY_PREFERENCE",
|
||||
"cumulative",
|
||||
),
|
||||
]),
|
||||
None,
|
||||
)
|
||||
.unwrap();
|
||||
assert_eq!(cfg.temporality, TemporalityPreference::Cumulative);
|
||||
}
|
||||
}
|
||||
269
crates/codegen/xai-grok-telemetry/src/external/emit.rs
vendored
Normal file
269
crates/codegen/xai-grok-telemetry/src/external/emit.rs
vendored
Normal file
|
|
@ -0,0 +1,269 @@
|
|||
//! The *how* of external emission: content-gate application, secret scrub +
|
||||
//! truncation, ctx (`session.id`/`turn_number`/`prompt.id`/`event.sequence`)
|
||||
//! injection, metric-increment conversion, and provider hand-off.
|
||||
//!
|
||||
//! The per-event *what* (field → attribute mapping) lives in
|
||||
//! [`super::schema`], wired via the `telemetry_event!` macro's
|
||||
//! `external = …` arm.
|
||||
|
||||
use opentelemetry::KeyValue;
|
||||
use opentelemetry::logs::{AnyValue, LogRecord as _, Logger as _, Severity};
|
||||
use opentelemetry::metrics::{Counter, Meter};
|
||||
|
||||
use super::ExternalTelemetry;
|
||||
use super::config::ContentGates;
|
||||
use super::schema::{
|
||||
AttrValue, ExternalKey, ExternalRecord, Gate, METRIC_ERROR_COUNT, METRIC_SESSION_COUNT,
|
||||
METRIC_TOKEN_USAGE, METRIC_TOOL_DECISION, METRIC_TOOL_USAGE, METRIC_TURN_COUNT,
|
||||
MetricIncrement,
|
||||
};
|
||||
|
||||
/// Pre-created counters (schema pinned by test: names, units, attr keys).
|
||||
pub(crate) struct Instruments {
|
||||
session_count: Counter<u64>,
|
||||
token_usage: Counter<u64>,
|
||||
turn_count: Counter<u64>,
|
||||
tool_decision: Counter<u64>,
|
||||
tool_usage: Counter<u64>,
|
||||
error_count: Counter<u64>,
|
||||
}
|
||||
|
||||
impl Instruments {
|
||||
pub(crate) fn new(meter: &Meter) -> Self {
|
||||
Self {
|
||||
session_count: meter
|
||||
.u64_counter(METRIC_SESSION_COUNT)
|
||||
.with_unit("{session}")
|
||||
.build(),
|
||||
token_usage: meter
|
||||
.u64_counter(METRIC_TOKEN_USAGE)
|
||||
.with_unit("{token}")
|
||||
.build(),
|
||||
turn_count: meter
|
||||
.u64_counter(METRIC_TURN_COUNT)
|
||||
.with_unit("{turn}")
|
||||
.build(),
|
||||
tool_decision: meter
|
||||
.u64_counter(METRIC_TOOL_DECISION)
|
||||
.with_unit("{decision}")
|
||||
.build(),
|
||||
tool_usage: meter
|
||||
.u64_counter(METRIC_TOOL_USAGE)
|
||||
.with_unit("{call}")
|
||||
.build(),
|
||||
error_count: meter
|
||||
.u64_counter(METRIC_ERROR_COUNT)
|
||||
.with_unit("{error}")
|
||||
.build(),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
fn gate_open(gates: ContentGates, gate: Gate) -> bool {
|
||||
match gate {
|
||||
Gate::UserPrompts => gates.log_user_prompts,
|
||||
Gate::ToolDetails => gates.log_tool_details,
|
||||
}
|
||||
}
|
||||
|
||||
/// Scrub + truncate one string attribute value. Every string passes the
|
||||
/// secret/path scrub; the prompt key gets the 60 KB content cap, everything
|
||||
/// else the standard 512→128 value truncation. Defense-in-depth only — the
|
||||
/// export-time validators in [`super::redact`] enforce the result.
|
||||
fn scrub_string(key: ExternalKey, s: String) -> String {
|
||||
let scrubbed = crate::redact_common::redact_to_owned(&s);
|
||||
match key {
|
||||
ExternalKey::Prompt => super::truncate::truncate_content(&scrubbed).unwrap_or(scrubbed),
|
||||
_ => super::truncate::truncate_value_owned(scrubbed),
|
||||
}
|
||||
}
|
||||
|
||||
fn to_any_value(v: AttrValue) -> AnyValue {
|
||||
match v {
|
||||
AttrValue::Str(s) => AnyValue::String(s.into()),
|
||||
AttrValue::I64(i) => AnyValue::Int(i),
|
||||
AttrValue::Bool(b) => AnyValue::Boolean(b),
|
||||
}
|
||||
}
|
||||
|
||||
/// Convert one mapped [`ExternalRecord`] into a log record and metric
|
||||
/// increments. Synchronous and cheap: the `BatchLogProcessor` queues the
|
||||
/// record; no `tokio::spawn` (contrast with the product-events path).
|
||||
pub(crate) fn emit_record(ext: &ExternalTelemetry, mut record: ExternalRecord) {
|
||||
let gates = *ext.gates.read();
|
||||
|
||||
// Gated attributes: emitted only when the matching gate is on. A gated
|
||||
// value sharing a key with a default attr (verbatim vs. sanitized
|
||||
// `tool_name`) replaces the default.
|
||||
for gated in std::mem::take(&mut record.gated) {
|
||||
if !gate_open(gates, gated.gate) {
|
||||
continue;
|
||||
}
|
||||
if let Some(existing) = record.attrs.iter_mut().find(|(k, _)| *k == gated.key) {
|
||||
existing.1 = gated.value;
|
||||
} else {
|
||||
record.attrs.push((gated.key, gated.value));
|
||||
}
|
||||
}
|
||||
|
||||
// Ambient ctx: a mapping-supplied `session.id` wins; the ctx is a
|
||||
// fallback for in-session events (the session-start sites are spawned
|
||||
// outside the ctx scope and carry their own ids).
|
||||
let ctx = crate::session_ctx::external_ctx_snapshot();
|
||||
let mapped_session_id = record
|
||||
.attrs
|
||||
.iter()
|
||||
.find(|(k, _)| *k == ExternalKey::SessionId)
|
||||
.and_then(|(_, v)| match v {
|
||||
AttrValue::Str(s) => Some(s.clone()),
|
||||
_ => None,
|
||||
});
|
||||
let session_id = mapped_session_id.or_else(|| ctx.as_ref().map(|c| c.session_id.clone()));
|
||||
|
||||
for (key, value) in record.attrs.iter_mut() {
|
||||
if let AttrValue::Str(s) = value {
|
||||
*value = AttrValue::Str(scrub_string(*key, std::mem::take(s)));
|
||||
}
|
||||
}
|
||||
|
||||
let identity = ext.identity.read().clone();
|
||||
|
||||
if let (Some(event), Some(logger)) = (record.event, ext.logger.as_ref()) {
|
||||
let mut log_record = logger.create_log_record();
|
||||
log_record.set_event_name(event.as_str());
|
||||
log_record.set_severity_number(Severity::Info);
|
||||
let now = std::time::SystemTime::now();
|
||||
log_record.set_timestamp(now);
|
||||
log_record.set_observed_timestamp(now);
|
||||
log_record.add_attribute(
|
||||
ExternalKey::EventSequence.as_str(),
|
||||
ext.next_sequence() as i64,
|
||||
);
|
||||
if record
|
||||
.attrs
|
||||
.iter()
|
||||
.all(|(k, _)| *k != ExternalKey::SessionId)
|
||||
&& let Some(sid) = session_id.as_deref()
|
||||
{
|
||||
log_record.add_attribute(ExternalKey::SessionId.as_str(), sid.to_owned());
|
||||
}
|
||||
if let Some(ctx) = ctx.as_ref() {
|
||||
if let Some(turn) = ctx.turn_number {
|
||||
log_record.add_attribute(ExternalKey::TurnNumber.as_str(), turn as i64);
|
||||
}
|
||||
// prompt.id: events only, never metrics (unbounded cardinality).
|
||||
if let Some(prompt_id) = ctx.prompt_id.as_deref() {
|
||||
log_record.add_attribute(ExternalKey::PromptId.as_str(), prompt_id.to_owned());
|
||||
}
|
||||
}
|
||||
for (key, value) in &record.attrs {
|
||||
log_record.add_attribute(key.as_str(), to_any_value(value.clone()));
|
||||
}
|
||||
for (key, value) in [
|
||||
(ExternalKey::UserId, identity.user_id.as_deref()),
|
||||
(
|
||||
ExternalKey::OrganizationId,
|
||||
identity.organization_id.as_deref(),
|
||||
),
|
||||
(ExternalKey::TeamId, identity.team_id.as_deref()),
|
||||
(ExternalKey::DeploymentId, identity.deployment_id.as_deref()),
|
||||
] {
|
||||
if let Some(v) = value.filter(|v| !v.is_empty()) {
|
||||
log_record.add_attribute(key.as_str(), v.to_owned());
|
||||
}
|
||||
}
|
||||
logger.emit(log_record);
|
||||
}
|
||||
|
||||
if let Some(instruments) = ext.instruments.as_ref() {
|
||||
for increment in record.metrics {
|
||||
add_increment(
|
||||
ext,
|
||||
instruments,
|
||||
increment,
|
||||
session_id.as_deref(),
|
||||
&identity,
|
||||
);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
fn add_increment(
|
||||
ext: &ExternalTelemetry,
|
||||
instruments: &Instruments,
|
||||
increment: MetricIncrement,
|
||||
session_id: Option<&str>,
|
||||
identity: &super::IdentityAttrs,
|
||||
) {
|
||||
// Identity/cardinality attrs shared by every instrument. `prompt.id` is
|
||||
// deliberately never attached to metrics.
|
||||
let mut attrs: Vec<KeyValue> = Vec::with_capacity(8);
|
||||
if ext.include_session_id_on_metrics
|
||||
&& let Some(sid) = session_id.filter(|s| !s.is_empty())
|
||||
{
|
||||
attrs.push(KeyValue::new("session.id", sid.to_owned()));
|
||||
}
|
||||
if ext.include_version_on_metrics && !ext.app_version.is_empty() {
|
||||
attrs.push(KeyValue::new("app.version", ext.app_version.clone()));
|
||||
}
|
||||
for (key, value) in [
|
||||
("user.id", identity.user_id.as_deref()),
|
||||
("organization.id", identity.organization_id.as_deref()),
|
||||
("team.id", identity.team_id.as_deref()),
|
||||
("deployment.id", identity.deployment_id.as_deref()),
|
||||
] {
|
||||
if let Some(v) = value.filter(|v| !v.is_empty()) {
|
||||
attrs.push(KeyValue::new(key, v.to_owned()));
|
||||
}
|
||||
}
|
||||
|
||||
// `model` is the one non-enum metric attribute value: scrub it at
|
||||
// increment time (call-site discipline is never the guarantee on its own
|
||||
// — the PR 6 collector fixture pins this with a wire-payload canary).
|
||||
let scrub = |s: &str| crate::redact_common::redact_to_owned(s);
|
||||
|
||||
match increment {
|
||||
MetricIncrement::SessionCount => {
|
||||
instruments.session_count.add(1, &attrs);
|
||||
}
|
||||
MetricIncrement::TokenUsage {
|
||||
token_type,
|
||||
model,
|
||||
count,
|
||||
} => {
|
||||
attrs.push(KeyValue::new("type", token_type));
|
||||
attrs.push(KeyValue::new("model", scrub(&model)));
|
||||
instruments.token_usage.add(count, &attrs);
|
||||
}
|
||||
MetricIncrement::TurnCount { outcome, model } => {
|
||||
attrs.push(KeyValue::new("outcome", outcome));
|
||||
attrs.push(KeyValue::new("model", scrub(&model)));
|
||||
instruments.turn_count.add(1, &attrs);
|
||||
}
|
||||
MetricIncrement::ToolDecision {
|
||||
tool_name,
|
||||
decision,
|
||||
access_kind,
|
||||
permission_mode,
|
||||
} => {
|
||||
attrs.push(KeyValue::new("tool_name", scrub(&tool_name)));
|
||||
attrs.push(KeyValue::new("decision", decision));
|
||||
attrs.push(KeyValue::new("access_kind", access_kind));
|
||||
attrs.push(KeyValue::new("permission_mode", permission_mode));
|
||||
instruments.tool_decision.add(1, &attrs);
|
||||
}
|
||||
MetricIncrement::ToolUsage { tool_name, outcome } => {
|
||||
attrs.push(KeyValue::new("tool_name", scrub(&tool_name)));
|
||||
attrs.push(KeyValue::new("outcome", outcome));
|
||||
instruments.tool_usage.add(1, &attrs);
|
||||
}
|
||||
MetricIncrement::ErrorCount {
|
||||
error_category,
|
||||
model,
|
||||
} => {
|
||||
attrs.push(KeyValue::new("error_category", scrub(&error_category)));
|
||||
attrs.push(KeyValue::new("model", scrub(&model)));
|
||||
instruments.error_count.add(1, &attrs);
|
||||
}
|
||||
}
|
||||
}
|
||||
480
crates/codegen/xai-grok-telemetry/src/external/mod.rs
vendored
Normal file
480
crates/codegen/xai-grok-telemetry/src/external/mod.rs
vendored
Normal file
|
|
@ -0,0 +1,480 @@
|
|||
//! Opt-in, content-redacted **external OTEL** telemetry stream.
|
||||
//!
|
||||
//! Enterprise customers point the Grok CLI at *their own* OpenTelemetry
|
||||
//! collector (standard `OTEL_*` env vars + the `GROK_EXTERNAL_OTEL` master
|
||||
//! switch) and receive a curated, ZDR-safe schema: ~6 counters and ~17
|
||||
//! log-record events fanned out from the same typed call sites that emit the
|
||||
//! product events ([`crate::session_ctx::log_event`]).
|
||||
//!
|
||||
//! Structural invariants (enforced by construction and tests):
|
||||
//! - The providers here are **never** registered with `opentelemetry::global`
|
||||
//! (the internal tracer provider owns the global slot); everything is
|
||||
//! handle-based through the [`EXTERNAL`] registry.
|
||||
//! - The exporters carry **only** customer headers/metadata from
|
||||
//! `OTEL_EXPORTER_OTLP_HEADERS` — this module has no dependency on
|
||||
//! `AuthCredentialProvider` and no code path that can attach internal auth
|
||||
//! headers.
|
||||
//! - Default **off**: with `GROK_EXTERNAL_OTEL` unset (or no exporter
|
||||
//! selected) nothing is constructed — zero allocation, zero threads, zero
|
||||
//! sockets.
|
||||
//! - Independent of `TelemetryMode`, GCS trace upload, and the data-collection /
|
||||
//! data-retention opt-outs (user-confirmed): those govern xAI-side
|
||||
//! retention; this stream ships only to the customer's own collector under
|
||||
//! the customer's own explicit double opt-in.
|
||||
//!
|
||||
//! This module is the second authoritative privacy boundary in this crate
|
||||
//! (alongside `otel_layer::redact`).
|
||||
|
||||
pub mod config;
|
||||
mod emit;
|
||||
mod providers;
|
||||
mod redact;
|
||||
pub mod schema;
|
||||
pub mod truncate;
|
||||
|
||||
use std::sync::atomic::{AtomicBool, AtomicU64, Ordering};
|
||||
use std::sync::{Arc, OnceLock};
|
||||
|
||||
use opentelemetry::logs::LoggerProvider as _;
|
||||
use opentelemetry::metrics::MeterProvider as _;
|
||||
use opentelemetry_sdk::logs::{SdkLogger, SdkLoggerProvider};
|
||||
use opentelemetry_sdk::metrics::SdkMeterProvider;
|
||||
|
||||
pub use config::{ContentGates, ExternalOtelConfig, ExternalOtelFileConfig};
|
||||
|
||||
static EXTERNAL: OnceLock<Option<Arc<ExternalTelemetry>>> = OnceLock::new();
|
||||
|
||||
/// Identity *attributes* (plain id strings — never tokens). Derived from a
|
||||
/// `CredentialSnapshot` at the telemetry-client init sites; updated post-auth
|
||||
/// and on logout.
|
||||
#[derive(Debug, Clone, Default)]
|
||||
pub struct IdentityAttrs {
|
||||
pub user_id: Option<String>,
|
||||
pub organization_id: Option<String>,
|
||||
pub team_id: Option<String>,
|
||||
pub deployment_id: Option<String>,
|
||||
}
|
||||
|
||||
impl IdentityAttrs {
|
||||
pub fn from_snapshot(snapshot: &xai_grok_auth::CredentialSnapshot) -> Self {
|
||||
Self {
|
||||
user_id: snapshot.user_id.clone(),
|
||||
organization_id: snapshot.organization_id.clone(),
|
||||
team_id: snapshot.team_id.clone(),
|
||||
deployment_id: snapshot.deployment_id.clone(),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Remote-settings policy for the external stream. **Restrictive-only by
|
||||
/// construction**: there is deliberately no enable direction (remote settings
|
||||
/// are fetched per-run and never persisted, so a remote "enable" could never
|
||||
/// reach init).
|
||||
#[derive(Debug, Clone, Copy, Default)]
|
||||
pub struct ExternalOtelRemotePolicy {
|
||||
/// Fleet kill switch: flush, then drop subsequent emissions in-process.
|
||||
pub force_disable: bool,
|
||||
/// Force the content gates off regardless of local env/config.
|
||||
pub lock_content_gates: bool,
|
||||
}
|
||||
|
||||
/// The handle owning both providers. Never global; reached only through the
|
||||
/// [`EXTERNAL`] registry.
|
||||
pub struct ExternalTelemetry {
|
||||
logger_provider: Option<SdkLoggerProvider>,
|
||||
meter_provider: Option<SdkMeterProvider>,
|
||||
logger: Option<SdkLogger>,
|
||||
instruments: Option<emit::Instruments>,
|
||||
/// Emission gate; cleared by the remote force-disable. The single
|
||||
/// authority for "emitting right now".
|
||||
active: AtomicBool,
|
||||
/// Content gates; may only TIGHTEN post-init.
|
||||
gates: redact::SharedGates,
|
||||
identity: parking_lot::RwLock<IdentityAttrs>,
|
||||
/// `event.sequence` (monotonic, per-process).
|
||||
sequence: AtomicU64,
|
||||
shutdown_once: std::sync::Once,
|
||||
include_session_id_on_metrics: bool,
|
||||
include_version_on_metrics: bool,
|
||||
app_version: String,
|
||||
health: Arc<redact::ExportHealth>,
|
||||
/// Init summary for the adoption meta-event (emitted once, post-auth).
|
||||
configured_meta: ConfiguredMeta,
|
||||
meta_event_once: std::sync::Once,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone)]
|
||||
struct ConfiguredMeta {
|
||||
metrics_exporter: &'static str,
|
||||
logs_exporter: &'static str,
|
||||
logs_endpoint_origin: String,
|
||||
metrics_endpoint_origin: String,
|
||||
protocol: &'static str,
|
||||
prompts_gate: bool,
|
||||
details_gate: bool,
|
||||
source: &'static str,
|
||||
}
|
||||
|
||||
fn exporter_label(sel: config::ExporterSelection) -> &'static str {
|
||||
match sel {
|
||||
config::ExporterSelection::None => "none",
|
||||
config::ExporterSelection::Otlp => "otlp",
|
||||
config::ExporterSelection::Console => "console",
|
||||
}
|
||||
}
|
||||
|
||||
impl ExternalTelemetry {
|
||||
pub(crate) fn next_sequence(&self) -> u64 {
|
||||
self.sequence.fetch_add(1, Ordering::Relaxed)
|
||||
}
|
||||
}
|
||||
|
||||
/// Initialize the external stream. Called once from binary startup after
|
||||
/// config resolution, **before auth** (no credentials needed). `None` records
|
||||
/// the dormant state — the default path allocates nothing.
|
||||
pub fn init(cfg: Option<ExternalOtelConfig>) {
|
||||
let value = cfg.and_then(build_handle);
|
||||
if EXTERNAL.set(value).is_err() {
|
||||
tracing::debug!("external otel: init called more than once; keeping first registration");
|
||||
}
|
||||
}
|
||||
|
||||
fn build_handle(cfg: ExternalOtelConfig) -> Option<Arc<ExternalTelemetry>> {
|
||||
// No-double-send invariant, enforced in code (not release discipline):
|
||||
// if the internal firehose resolved its endpoint/headers from
|
||||
// `OTEL_EXPORTER_OTLP_*` (the deprecated fallback), refuse to activate.
|
||||
if cfg.internal_pipeline_consumed_otel_vars {
|
||||
tracing::warn!(
|
||||
"external otel: refusing to activate — the internal trace pipeline consumed \
|
||||
OTEL_EXPORTER_OTLP_* (deprecated fallback). Migrate internal repointing to \
|
||||
GROK_INTERNAL_OTLP_* to use the external stream."
|
||||
);
|
||||
return None;
|
||||
}
|
||||
|
||||
let gates: redact::SharedGates = Arc::new(parking_lot::RwLock::new(cfg.gates));
|
||||
let health = Arc::new(redact::ExportHealth::default());
|
||||
let built = match providers::build(&cfg, gates.clone(), health.clone()) {
|
||||
Ok(built) => built,
|
||||
Err(e) => {
|
||||
tracing::warn!(error = %e, "external otel: exporter construction failed; stream disabled");
|
||||
return None;
|
||||
}
|
||||
};
|
||||
if built.logger_provider.is_none() && built.meter_provider.is_none() {
|
||||
return None;
|
||||
}
|
||||
|
||||
let logger = built
|
||||
.logger_provider
|
||||
.as_ref()
|
||||
.map(|p| p.logger(schema::SCOPE_NAME));
|
||||
let instruments = built
|
||||
.meter_provider
|
||||
.as_ref()
|
||||
.map(|p| emit::Instruments::new(&p.meter(schema::SCOPE_NAME)));
|
||||
|
||||
let configured_meta = ConfiguredMeta {
|
||||
metrics_exporter: exporter_label(cfg.metrics_exporter),
|
||||
logs_exporter: exporter_label(cfg.logs_exporter),
|
||||
logs_endpoint_origin: crate::redact_common::url_origin(&cfg.logs_endpoint).into_owned(),
|
||||
metrics_endpoint_origin: crate::redact_common::url_origin(&cfg.metrics_endpoint)
|
||||
.into_owned(),
|
||||
protocol: cfg.transport.as_protocol_str(),
|
||||
prompts_gate: cfg.gates.log_user_prompts,
|
||||
details_gate: cfg.gates.log_tool_details,
|
||||
source: cfg.enabled_source,
|
||||
};
|
||||
|
||||
tracing::debug!(
|
||||
metrics_exporter = configured_meta.metrics_exporter,
|
||||
logs_exporter = configured_meta.logs_exporter,
|
||||
"external otel: stream active"
|
||||
);
|
||||
|
||||
Some(Arc::new(ExternalTelemetry {
|
||||
logger_provider: built.logger_provider,
|
||||
meter_provider: built.meter_provider,
|
||||
logger,
|
||||
instruments,
|
||||
active: AtomicBool::new(true),
|
||||
gates,
|
||||
identity: parking_lot::RwLock::new(IdentityAttrs::default()),
|
||||
sequence: AtomicU64::new(0),
|
||||
shutdown_once: std::sync::Once::new(),
|
||||
include_session_id_on_metrics: cfg.include_session_id_on_metrics,
|
||||
include_version_on_metrics: cfg.include_version_on_metrics,
|
||||
app_version: cfg.client.client_version.clone(),
|
||||
health,
|
||||
configured_meta,
|
||||
meta_event_once: std::sync::Once::new(),
|
||||
}))
|
||||
}
|
||||
|
||||
fn handle() -> Option<Arc<ExternalTelemetry>> {
|
||||
EXTERNAL.get().and_then(|opt| opt.clone())
|
||||
}
|
||||
|
||||
fn active_handle() -> Option<Arc<ExternalTelemetry>> {
|
||||
handle().filter(|ext| ext.active.load(Ordering::Relaxed))
|
||||
}
|
||||
|
||||
/// Cheap check used by the fan-out hook and the split-sink call sites:
|
||||
/// registry present AND the runtime emission gate set. A stale `true` read
|
||||
/// only costs a wasted mapping, never an export ([`emit`] re-checks).
|
||||
pub fn is_active() -> bool {
|
||||
matches!(EXTERNAL.get(), Some(Some(ext)) if ext.active.load(Ordering::Relaxed))
|
||||
}
|
||||
|
||||
/// Map and emit one typed telemetry event. No-op unless the stream is active
|
||||
/// and the event has an `external = …` mapping. Synchronous and cheap (the
|
||||
/// batch processor queues; nothing blocks on I/O).
|
||||
pub fn emit<T: crate::events::TelemetryEvent>(data: &T) {
|
||||
let Some(ext) = active_handle() else {
|
||||
return;
|
||||
};
|
||||
let Some(record) = data.external_record() else {
|
||||
return;
|
||||
};
|
||||
emit::emit_record(&ext, record);
|
||||
}
|
||||
|
||||
/// Update identity attrs when auth completes (called alongside the
|
||||
/// telemetry-client init sites). Also emits the one-shot internal adoption
|
||||
/// meta-event — post-auth, when the product events client is live.
|
||||
pub fn set_identity(attrs: IdentityAttrs) {
|
||||
let Some(ext) = handle() else {
|
||||
return;
|
||||
};
|
||||
set_identity_on(&ext, attrs);
|
||||
}
|
||||
|
||||
pub(crate) fn set_identity_on(ext: &ExternalTelemetry, attrs: IdentityAttrs) {
|
||||
*ext.identity.write() = attrs;
|
||||
ext.meta_event_once.call_once(|| {
|
||||
let meta = &ext.configured_meta;
|
||||
crate::session_ctx::log_session_event(crate::events::ExternalOtelConfigured {
|
||||
metrics_exporter: meta.metrics_exporter.to_owned(),
|
||||
logs_exporter: meta.logs_exporter.to_owned(),
|
||||
protocol: meta.protocol.to_owned(),
|
||||
logs_endpoint_origin: meta.logs_endpoint_origin.clone(),
|
||||
metrics_endpoint_origin: meta.metrics_endpoint_origin.clone(),
|
||||
prompts_gate: meta.prompts_gate,
|
||||
details_gate: meta.details_gate,
|
||||
source: meta.source.to_owned(),
|
||||
});
|
||||
});
|
||||
}
|
||||
|
||||
/// Apply remote policy when `RemoteSettings` arrive (post-auth, alongside
|
||||
/// [`set_identity`]). **TIGHTEN-ONLY**: may clear `active` (fleet kill switch
|
||||
/// — flushes, then drops subsequent emissions) and may force content gates
|
||||
/// off; it can never enable a stream that env/config left off, and never
|
||||
/// loosens gates mid-run.
|
||||
pub fn apply_remote_policy(policy: ExternalOtelRemotePolicy) {
|
||||
let Some(ext) = handle() else {
|
||||
return;
|
||||
};
|
||||
apply_remote_policy_on(&ext, policy);
|
||||
}
|
||||
|
||||
pub(crate) fn apply_remote_policy_on(ext: &ExternalTelemetry, policy: ExternalOtelRemotePolicy) {
|
||||
if policy.lock_content_gates {
|
||||
let mut gates = ext.gates.write();
|
||||
if gates.log_user_prompts || gates.log_tool_details {
|
||||
*gates = ContentGates::default();
|
||||
drop(gates);
|
||||
crate::session_ctx::log_session_event(crate::events::ExternalOtelRemotePolicyApplied {
|
||||
action: "gates_locked".to_owned(),
|
||||
});
|
||||
}
|
||||
}
|
||||
if policy.force_disable && ext.active.swap(false, Ordering::Relaxed) {
|
||||
flush_on(ext);
|
||||
crate::session_ctx::log_session_event(crate::events::ExternalOtelRemotePolicyApplied {
|
||||
action: "force_disable".to_owned(),
|
||||
});
|
||||
tracing::debug!("external otel: force-disabled by remote policy");
|
||||
}
|
||||
}
|
||||
|
||||
/// Flush both providers (logout path: called *before* credentials are
|
||||
/// cleared, so post-logout records cannot carry the prior user's ids —
|
||||
/// follow with [`set_identity`] carrying the new/empty identity).
|
||||
pub fn flush() {
|
||||
let Some(ext) = handle() else {
|
||||
return;
|
||||
};
|
||||
flush_on(&ext);
|
||||
}
|
||||
|
||||
pub(crate) fn flush_on(ext: &ExternalTelemetry) {
|
||||
if let Some(p) = ext.logger_provider.as_ref()
|
||||
&& let Err(e) = p.force_flush()
|
||||
{
|
||||
tracing::debug!(error = %e, "external otel: logger flush failed");
|
||||
}
|
||||
if let Some(p) = ext.meter_provider.as_ref()
|
||||
&& let Err(e) = p.force_flush()
|
||||
{
|
||||
tracing::debug!(error = %e, "external otel: meter flush failed");
|
||||
}
|
||||
}
|
||||
|
||||
/// Flush + shutdown both providers with a 2-second watchdog. Idempotent —
|
||||
/// reachable from every `shutdown_otel()` exit path (16 `OtelGuard` sites,
|
||||
/// the direct call, and the signal handler); subsequent calls are no-ops.
|
||||
pub fn shutdown() {
|
||||
let Some(ext) = handle() else {
|
||||
return;
|
||||
};
|
||||
ext.shutdown_once.call_once(|| {
|
||||
ext.active.store(false, Ordering::Relaxed);
|
||||
emit_export_health(&ext);
|
||||
let logger_provider = ext.logger_provider.clone();
|
||||
let meter_provider = ext.meter_provider.clone();
|
||||
let (tx, rx) = std::sync::mpsc::channel::<()>();
|
||||
// Detached thread + timed wait: a hung provider must not hang exit
|
||||
// (`std::thread::scope` is unusable here — it joins unconditionally).
|
||||
std::thread::spawn(move || {
|
||||
if let Some(p) = logger_provider
|
||||
&& let Err(e) = p.shutdown()
|
||||
{
|
||||
tracing::debug!(error = %e, "external otel: logger shutdown failed");
|
||||
}
|
||||
if let Some(p) = meter_provider
|
||||
&& let Err(e) = p.shutdown()
|
||||
{
|
||||
tracing::debug!(error = %e, "external otel: meter shutdown failed");
|
||||
}
|
||||
let _ = tx.send(());
|
||||
});
|
||||
if rx.recv_timeout(std::time::Duration::from_secs(2)).is_err() {
|
||||
tracing::debug!("external otel: shutdown watchdog expired; abandoning flush thread");
|
||||
}
|
||||
});
|
||||
}
|
||||
|
||||
/// Best-effort product-events export-health meta-event (never exported
|
||||
/// externally — avoid feedback loops). Emitting needs a Tokio runtime
|
||||
/// (`emit_event` spawns); skip silently when exiting without one.
|
||||
fn emit_export_health(ext: &ExternalTelemetry) {
|
||||
let health = &ext.health;
|
||||
let snapshot = crate::events::ExternalOtelExportHealth {
|
||||
records_dropped: health.records_dropped.load(Ordering::Relaxed),
|
||||
metric_exports_dropped: health.metric_exports_dropped.load(Ordering::Relaxed),
|
||||
export_failures: health.export_failures.load(Ordering::Relaxed),
|
||||
export_successes: health.export_successes.load(Ordering::Relaxed),
|
||||
};
|
||||
tracing::debug!(
|
||||
records_dropped = snapshot.records_dropped,
|
||||
metric_exports_dropped = snapshot.metric_exports_dropped,
|
||||
export_failures = snapshot.export_failures,
|
||||
export_successes = snapshot.export_successes,
|
||||
"external otel: export health"
|
||||
);
|
||||
if tokio::runtime::Handle::try_current().is_ok() {
|
||||
crate::session_ctx::log_session_event(snapshot);
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
pub(crate) mod test_support {
|
||||
//! Build an [`ExternalTelemetry`] over in-memory exporters so unit tests
|
||||
//! can assert exactly what would reach the wire (post-validator).
|
||||
|
||||
use super::*;
|
||||
use opentelemetry_sdk::logs::InMemoryLogExporter;
|
||||
use opentelemetry_sdk::metrics::{InMemoryMetricExporter, PeriodicReader};
|
||||
|
||||
pub(crate) struct TestStream {
|
||||
pub ext: ExternalTelemetry,
|
||||
pub logs: InMemoryLogExporter,
|
||||
pub metrics: InMemoryMetricExporter,
|
||||
}
|
||||
|
||||
pub(crate) fn build(gates: ContentGates) -> TestStream {
|
||||
let shared_gates: redact::SharedGates = Arc::new(parking_lot::RwLock::new(gates));
|
||||
let health = Arc::new(redact::ExportHealth::default());
|
||||
let logs = InMemoryLogExporter::default();
|
||||
let metrics = InMemoryMetricExporter::default();
|
||||
|
||||
let logger_provider = SdkLoggerProvider::builder()
|
||||
.with_simple_exporter(redact::RedactingLogExporter::new(
|
||||
logs.clone(),
|
||||
shared_gates.clone(),
|
||||
health.clone(),
|
||||
))
|
||||
.build();
|
||||
let meter_provider = SdkMeterProvider::builder()
|
||||
.with_reader(
|
||||
PeriodicReader::builder(redact::ValidatingMetricExporter::new(
|
||||
metrics.clone(),
|
||||
health.clone(),
|
||||
))
|
||||
.build(),
|
||||
)
|
||||
.build();
|
||||
|
||||
let logger = logger_provider.logger(schema::SCOPE_NAME);
|
||||
let instruments = emit::Instruments::new(&meter_provider.meter(schema::SCOPE_NAME));
|
||||
|
||||
let ext = ExternalTelemetry {
|
||||
logger_provider: Some(logger_provider),
|
||||
meter_provider: Some(meter_provider),
|
||||
logger: Some(logger),
|
||||
instruments: Some(instruments),
|
||||
active: AtomicBool::new(true),
|
||||
gates: shared_gates,
|
||||
identity: parking_lot::RwLock::new(IdentityAttrs::default()),
|
||||
sequence: AtomicU64::new(0),
|
||||
shutdown_once: std::sync::Once::new(),
|
||||
include_session_id_on_metrics: true,
|
||||
include_version_on_metrics: false,
|
||||
app_version: String::new(),
|
||||
health,
|
||||
configured_meta: ConfiguredMeta {
|
||||
metrics_exporter: "test",
|
||||
logs_exporter: "test",
|
||||
logs_endpoint_origin: String::new(),
|
||||
metrics_endpoint_origin: String::new(),
|
||||
protocol: "test",
|
||||
prompts_gate: gates.log_user_prompts,
|
||||
details_gate: gates.log_tool_details,
|
||||
source: "env",
|
||||
},
|
||||
meta_event_once: std::sync::Once::new(),
|
||||
};
|
||||
TestStream { ext, logs, metrics }
|
||||
}
|
||||
|
||||
pub(crate) fn emit_into(stream: &TestStream, record: schema::ExternalRecord) {
|
||||
emit::emit_record(&stream.ext, record);
|
||||
stream
|
||||
.ext
|
||||
.logger_provider
|
||||
.as_ref()
|
||||
.expect("test logger provider")
|
||||
.force_flush()
|
||||
.expect("flush logs");
|
||||
stream
|
||||
.ext
|
||||
.meter_provider
|
||||
.as_ref()
|
||||
.expect("test meter provider")
|
||||
.force_flush()
|
||||
.expect("flush metrics");
|
||||
}
|
||||
|
||||
pub(crate) fn emit_event_into<T: crate::events::TelemetryEvent>(
|
||||
stream: &TestStream,
|
||||
event: &T,
|
||||
) {
|
||||
if let Some(record) = event.external_record() {
|
||||
emit_into(stream, record);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests;
|
||||
585
crates/codegen/xai-grok-telemetry/src/external/providers.rs
vendored
Normal file
585
crates/codegen/xai-grok-telemetry/src/external/providers.rs
vendored
Normal file
|
|
@ -0,0 +1,585 @@
|
|||
//! Provider construction for the external stream: `SdkLoggerProvider` +
|
||||
//! `SdkMeterProvider`, never registered globally, never sharing anything with
|
||||
//! the internal `RefreshableSpanExporter` pipeline.
|
||||
//!
|
||||
//! The exporters are plain `opentelemetry_otlp` http/protobuf or gRPC/protobuf
|
||||
//! exporters built with **only** the customer headers from
|
||||
//! `OTEL_EXPORTER_OTLP_HEADERS` — no code path here can attach
|
||||
//! `Authorization`/`X-XAI-Token-Auth`/`x-userid`;
|
||||
//! those constants live in `otel_layer` and are not referenced by this
|
||||
//! module. No `AuthCredentialProvider` is ever read.
|
||||
|
||||
use std::sync::Arc;
|
||||
use std::time::Duration;
|
||||
|
||||
use http::{HeaderMap, HeaderName, HeaderValue};
|
||||
use opentelemetry_otlp::{
|
||||
Protocol, WithExportConfig, WithHttpConfig, WithTonicConfig, tonic_types::metadata::MetadataMap,
|
||||
};
|
||||
use opentelemetry_sdk::logs::{
|
||||
BatchConfig, BatchConfigBuilder, BatchLogProcessor as ThreadBatchLogProcessor,
|
||||
LoggerProviderBuilder, SdkLoggerProvider,
|
||||
log_processor_with_async_runtime::BatchLogProcessor as RuntimeBatchLogProcessor,
|
||||
};
|
||||
use opentelemetry_sdk::metrics::{
|
||||
MeterProviderBuilder, PeriodicReader as ThreadPeriodicReader, SdkMeterProvider, Temporality,
|
||||
periodic_reader_with_async_runtime::PeriodicReader as RuntimePeriodicReader,
|
||||
};
|
||||
type BuildResult<T> = Result<T, opentelemetry_otlp::ExporterBuildError>;
|
||||
|
||||
type RuntimeCommand = std::pin::Pin<Box<dyn std::future::Future<Output = ()> + Send>>;
|
||||
|
||||
#[derive(Clone)]
|
||||
struct DedicatedRuntime {
|
||||
tx: tokio::sync::mpsc::UnboundedSender<RuntimeCommand>,
|
||||
}
|
||||
|
||||
impl std::fmt::Debug for DedicatedRuntime {
|
||||
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
|
||||
f.write_str("DedicatedRuntime")
|
||||
}
|
||||
}
|
||||
|
||||
impl DedicatedRuntime {
|
||||
fn new() -> Self {
|
||||
let (tx, mut rx) = tokio::sync::mpsc::unbounded_channel::<RuntimeCommand>();
|
||||
std::thread::spawn(move || {
|
||||
let rt = tokio::runtime::Builder::new_current_thread()
|
||||
.enable_all()
|
||||
.build()
|
||||
.expect("external OTEL gRPC runtime");
|
||||
rt.block_on(async move {
|
||||
while let Some(future) = rx.recv().await {
|
||||
tokio::spawn(future);
|
||||
}
|
||||
});
|
||||
});
|
||||
Self { tx }
|
||||
}
|
||||
|
||||
fn run<T: Send + 'static>(
|
||||
&self,
|
||||
f: impl FnOnce() -> BuildResult<T> + Send + 'static,
|
||||
) -> BuildResult<T> {
|
||||
let (tx, rx) = std::sync::mpsc::channel();
|
||||
self.tx
|
||||
.send(Box::pin(async move {
|
||||
let _ = tx.send(f());
|
||||
}))
|
||||
.expect("external OTEL gRPC runtime thread must be alive");
|
||||
rx.recv()
|
||||
.expect("external OTEL gRPC runtime build response")
|
||||
}
|
||||
}
|
||||
|
||||
impl opentelemetry_sdk::runtime::Runtime for DedicatedRuntime {
|
||||
fn spawn<F>(&self, future: F)
|
||||
where
|
||||
F: std::future::Future<Output = ()> + Send + 'static,
|
||||
{
|
||||
let _ = self.tx.send(Box::pin(future));
|
||||
}
|
||||
|
||||
fn delay(&self, duration: Duration) -> impl std::future::Future<Output = ()> + Send + 'static {
|
||||
tokio::time::sleep(duration)
|
||||
}
|
||||
}
|
||||
|
||||
impl opentelemetry_sdk::runtime::RuntimeChannel for DedicatedRuntime {
|
||||
type Receiver<T: std::fmt::Debug + Send> = tokio_stream::wrappers::ReceiverStream<T>;
|
||||
type Sender<T: std::fmt::Debug + Send> = tokio::sync::mpsc::Sender<T>;
|
||||
|
||||
fn batch_message_channel<T: std::fmt::Debug + Send>(
|
||||
&self,
|
||||
capacity: usize,
|
||||
) -> (Self::Sender<T>, Self::Receiver<T>) {
|
||||
let (sender, receiver) = tokio::sync::mpsc::channel(capacity);
|
||||
(
|
||||
sender,
|
||||
tokio_stream::wrappers::ReceiverStream::new(receiver),
|
||||
)
|
||||
}
|
||||
}
|
||||
|
||||
use super::config::{ExporterSelection, ExternalOtelConfig, OtlpTransport, TemporalityPreference};
|
||||
use super::redact::{ExportHealth, RedactingLogExporter, SharedGates, ValidatingMetricExporter};
|
||||
|
||||
/// Resource shared by both providers. `builder_empty()` (not `builder()`):
|
||||
/// the default `EnvResourceDetector` would export `OTEL_RESOURCE_ATTRIBUTES`
|
||||
/// env values, bypassing the schema (same rationale as the internal layer).
|
||||
fn build_resource(cfg: &ExternalOtelConfig) -> opentelemetry_sdk::Resource {
|
||||
let mut attrs = vec![
|
||||
opentelemetry::KeyValue::new("service.version", cfg.client.service_version.clone()),
|
||||
opentelemetry::KeyValue::new("client.version", cfg.client.client_version.clone()),
|
||||
opentelemetry::KeyValue::new("app.entrypoint", cfg.client.app_entrypoint.clone()),
|
||||
opentelemetry::KeyValue::new("grok_code.schema.version", super::schema::SCHEMA_VERSION),
|
||||
];
|
||||
// terminal.type: emulator brand (TERM_PROGRAM) or terminfo type (TERM).
|
||||
if let Some(terminal_type) = std::env::var("TERM_PROGRAM")
|
||||
.ok()
|
||||
.or_else(|| std::env::var("TERM").ok())
|
||||
.filter(|v| !v.is_empty())
|
||||
{
|
||||
attrs.push(opentelemetry::KeyValue::new("terminal.type", terminal_type));
|
||||
}
|
||||
opentelemetry_sdk::Resource::builder_empty()
|
||||
// RQ6 (final): `grok-cli`, a wire commitment.
|
||||
.with_service_name("grok-cli")
|
||||
.with_attributes(attrs)
|
||||
.build()
|
||||
}
|
||||
|
||||
fn temporality(pref: TemporalityPreference) -> Temporality {
|
||||
match pref {
|
||||
TemporalityPreference::Delta => Temporality::Delta,
|
||||
TemporalityPreference::Cumulative => Temporality::Cumulative,
|
||||
}
|
||||
}
|
||||
|
||||
/// Console (stderr) log exporter for local debugging
|
||||
/// (`OTEL_LOGS_EXPORTER=console`). Writes to **stderr** so stdout protocol
|
||||
/// channels (headless/stream-JSON) are never corrupted.
|
||||
#[derive(Debug)]
|
||||
struct StderrLogExporter;
|
||||
|
||||
impl opentelemetry_sdk::logs::LogExporter for StderrLogExporter {
|
||||
fn export(
|
||||
&self,
|
||||
batch: opentelemetry_sdk::logs::LogBatch<'_>,
|
||||
) -> impl std::future::Future<Output = opentelemetry_sdk::error::OTelSdkResult> + Send {
|
||||
for (record, _scope) in batch.iter() {
|
||||
let attrs: Vec<String> = record
|
||||
.attributes_iter()
|
||||
.map(|(k, v)| format!("{}={v:?}", k.as_str()))
|
||||
.collect();
|
||||
eprintln!(
|
||||
"[external-otel] event={} {}",
|
||||
record.event_name().unwrap_or("?"),
|
||||
attrs.join(" ")
|
||||
);
|
||||
}
|
||||
std::future::ready(Ok(()))
|
||||
}
|
||||
}
|
||||
|
||||
/// Console (stderr) metric exporter for local debugging.
|
||||
#[derive(Debug)]
|
||||
struct StderrMetricExporter {
|
||||
temporality: Temporality,
|
||||
}
|
||||
|
||||
impl opentelemetry_sdk::metrics::exporter::PushMetricExporter for StderrMetricExporter {
|
||||
fn export(
|
||||
&self,
|
||||
metrics: &opentelemetry_sdk::metrics::data::ResourceMetrics,
|
||||
) -> impl std::future::Future<Output = opentelemetry_sdk::error::OTelSdkResult> + Send {
|
||||
for scope in metrics.scope_metrics() {
|
||||
for metric in scope.metrics() {
|
||||
eprintln!(
|
||||
"[external-otel] metric={} {:?}",
|
||||
metric.name(),
|
||||
metric.data()
|
||||
);
|
||||
}
|
||||
}
|
||||
std::future::ready(Ok(()))
|
||||
}
|
||||
|
||||
fn force_flush(&self) -> opentelemetry_sdk::error::OTelSdkResult {
|
||||
Ok(())
|
||||
}
|
||||
|
||||
fn shutdown_with_timeout(
|
||||
&self,
|
||||
_timeout: std::time::Duration,
|
||||
) -> opentelemetry_sdk::error::OTelSdkResult {
|
||||
Ok(())
|
||||
}
|
||||
|
||||
fn temporality(&self) -> Temporality {
|
||||
self.temporality
|
||||
}
|
||||
}
|
||||
|
||||
/// Customer headers as the HTTP OTLP builder's header map. The **only** headers
|
||||
/// the external HTTP exporters send (pinned by the header-isolation test below).
|
||||
fn customer_headers(headers: &[(String, String)]) -> std::collections::HashMap<String, String> {
|
||||
headers.iter().cloned().collect()
|
||||
}
|
||||
|
||||
/// Customer headers as gRPC metadata. Invalid metadata keys/values are skipped;
|
||||
/// this mirrors the HTTP builder's "only customer-supplied headers" invariant
|
||||
/// without letting one malformed entry disable telemetry entirely.
|
||||
fn customer_metadata(input: &[(String, String)]) -> MetadataMap {
|
||||
let mut headers = HeaderMap::new();
|
||||
for (key, value) in input {
|
||||
let Ok(header_name) = HeaderName::try_from(key.as_str()) else {
|
||||
tracing::warn!(key = %key, "external otel: skipping invalid gRPC metadata key");
|
||||
continue;
|
||||
};
|
||||
let Ok(header_value) = HeaderValue::from_str(value) else {
|
||||
tracing::warn!(key = %key, "external otel: skipping invalid gRPC metadata value");
|
||||
continue;
|
||||
};
|
||||
headers.insert(header_name, header_value);
|
||||
}
|
||||
MetadataMap::from_headers(headers)
|
||||
}
|
||||
|
||||
pub(crate) struct BuiltProviders {
|
||||
pub logger_provider: Option<SdkLoggerProvider>,
|
||||
pub meter_provider: Option<SdkMeterProvider>,
|
||||
}
|
||||
|
||||
enum OtlpExportTransport<'a> {
|
||||
HttpProtobuf(&'a crate::otlp_http::BlockingOtlpClient),
|
||||
Grpc(&'a DedicatedRuntime),
|
||||
}
|
||||
|
||||
trait OtlpExportFactory {
|
||||
type Exporter;
|
||||
|
||||
fn export(&self, transport: OtlpExportTransport<'_>) -> BuildResult<Self::Exporter>;
|
||||
}
|
||||
|
||||
struct OtlpLogExporterBuilder<'a> {
|
||||
cfg: &'a ExternalOtelConfig,
|
||||
}
|
||||
|
||||
impl OtlpExportFactory for OtlpLogExporterBuilder<'_> {
|
||||
type Exporter = opentelemetry_otlp::LogExporter;
|
||||
|
||||
fn export(&self, transport: OtlpExportTransport<'_>) -> BuildResult<Self::Exporter> {
|
||||
match transport {
|
||||
OtlpExportTransport::HttpProtobuf(http_client) => {
|
||||
opentelemetry_otlp::LogExporter::builder()
|
||||
.with_http()
|
||||
// Pin http/protobuf. opentelemetry-otlp's default protocol
|
||||
// is compile-time, feature-gated: `http-json` (if unified
|
||||
// into the build, as it is under Bazel) flips the default to
|
||||
// JSON, while a pure-cargo build of this crate defaults to
|
||||
// protobuf. Pin explicitly so the contract holds on every
|
||||
// build when HTTP transport is selected.
|
||||
.with_protocol(Protocol::HttpBinary)
|
||||
.with_http_client(http_client.clone())
|
||||
.with_endpoint(&self.cfg.logs_endpoint)
|
||||
.with_headers(customer_headers(&self.cfg.logs_headers))
|
||||
.build()
|
||||
}
|
||||
OtlpExportTransport::Grpc(runtime) => {
|
||||
let endpoint = self.cfg.logs_endpoint.clone();
|
||||
let timeout = self.cfg.timeout;
|
||||
let metadata = customer_metadata(&self.cfg.logs_headers);
|
||||
runtime.run(move || {
|
||||
opentelemetry_otlp::LogExporter::builder()
|
||||
.with_tonic()
|
||||
.with_endpoint(endpoint)
|
||||
.with_timeout(timeout)
|
||||
.with_metadata(metadata)
|
||||
.build()
|
||||
})
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
struct OtlpMetricExporterBuilder<'a> {
|
||||
cfg: &'a ExternalOtelConfig,
|
||||
temporality: Temporality,
|
||||
}
|
||||
|
||||
impl OtlpExportFactory for OtlpMetricExporterBuilder<'_> {
|
||||
type Exporter = opentelemetry_otlp::MetricExporter;
|
||||
|
||||
fn export(&self, transport: OtlpExportTransport<'_>) -> BuildResult<Self::Exporter> {
|
||||
match transport {
|
||||
OtlpExportTransport::HttpProtobuf(http_client) => {
|
||||
opentelemetry_otlp::MetricExporter::builder()
|
||||
.with_http()
|
||||
// Pin http/protobuf (see the logs exporter above for the
|
||||
// feature-unification rationale).
|
||||
.with_protocol(Protocol::HttpBinary)
|
||||
.with_http_client(http_client.clone())
|
||||
.with_endpoint(&self.cfg.metrics_endpoint)
|
||||
.with_headers(customer_headers(&self.cfg.metrics_headers))
|
||||
.with_temporality(self.temporality)
|
||||
.build()
|
||||
}
|
||||
OtlpExportTransport::Grpc(runtime) => {
|
||||
let endpoint = self.cfg.metrics_endpoint.clone();
|
||||
let timeout = self.cfg.timeout;
|
||||
let metadata = customer_metadata(&self.cfg.metrics_headers);
|
||||
let temporality = self.temporality;
|
||||
runtime.run(move || {
|
||||
opentelemetry_otlp::MetricExporter::builder()
|
||||
.with_tonic()
|
||||
.with_endpoint(endpoint)
|
||||
.with_timeout(timeout)
|
||||
.with_metadata(metadata)
|
||||
.with_temporality(temporality)
|
||||
.build()
|
||||
})
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
fn build_log_otlp_provider(
|
||||
builder: LoggerProviderBuilder,
|
||||
cfg: &ExternalOtelConfig,
|
||||
batch_config: BatchConfig,
|
||||
http_client: Option<&crate::otlp_http::BlockingOtlpClient>,
|
||||
gates: SharedGates,
|
||||
health: Arc<ExportHealth>,
|
||||
) -> BuildResult<LoggerProviderBuilder> {
|
||||
let exporter_builder = OtlpLogExporterBuilder { cfg };
|
||||
Ok(match cfg.transport {
|
||||
OtlpTransport::HttpProtobuf => {
|
||||
let exporter = exporter_builder.export(OtlpExportTransport::HttpProtobuf(
|
||||
http_client.expect("client built for http/protobuf selection"),
|
||||
))?;
|
||||
builder.with_log_processor(
|
||||
ThreadBatchLogProcessor::builder(RedactingLogExporter::new(
|
||||
exporter, gates, health,
|
||||
))
|
||||
.with_batch_config(batch_config)
|
||||
.build(),
|
||||
)
|
||||
}
|
||||
OtlpTransport::Grpc => {
|
||||
let runtime = DedicatedRuntime::new();
|
||||
let exporter = exporter_builder.export(OtlpExportTransport::Grpc(&runtime))?;
|
||||
builder.with_log_processor(
|
||||
RuntimeBatchLogProcessor::builder(
|
||||
RedactingLogExporter::new(exporter, gates, health),
|
||||
runtime,
|
||||
)
|
||||
.with_batch_config(batch_config)
|
||||
.build(),
|
||||
)
|
||||
}
|
||||
})
|
||||
}
|
||||
|
||||
fn build_metric_otlp_provider(
|
||||
builder: MeterProviderBuilder,
|
||||
cfg: &ExternalOtelConfig,
|
||||
http_client: Option<&crate::otlp_http::BlockingOtlpClient>,
|
||||
health: Arc<ExportHealth>,
|
||||
) -> BuildResult<MeterProviderBuilder> {
|
||||
let exporter_builder = OtlpMetricExporterBuilder {
|
||||
cfg,
|
||||
temporality: temporality(cfg.temporality),
|
||||
};
|
||||
Ok(match cfg.transport {
|
||||
OtlpTransport::HttpProtobuf => {
|
||||
let exporter = exporter_builder.export(OtlpExportTransport::HttpProtobuf(
|
||||
http_client.expect("client built for http/protobuf selection"),
|
||||
))?;
|
||||
builder.with_reader(
|
||||
ThreadPeriodicReader::builder(ValidatingMetricExporter::new(exporter, health))
|
||||
.with_interval(cfg.metric_export_interval)
|
||||
.build(),
|
||||
)
|
||||
}
|
||||
OtlpTransport::Grpc => {
|
||||
let runtime = DedicatedRuntime::new();
|
||||
let exporter = exporter_builder.export(OtlpExportTransport::Grpc(&runtime))?;
|
||||
builder.with_reader(
|
||||
RuntimePeriodicReader::builder(
|
||||
ValidatingMetricExporter::new(exporter, health),
|
||||
runtime,
|
||||
)
|
||||
.with_interval(cfg.metric_export_interval)
|
||||
.build(),
|
||||
)
|
||||
}
|
||||
})
|
||||
}
|
||||
|
||||
fn wrap_console_log_exporter(
|
||||
builder: LoggerProviderBuilder,
|
||||
batch_config: BatchConfig,
|
||||
gates: SharedGates,
|
||||
health: Arc<ExportHealth>,
|
||||
) -> LoggerProviderBuilder {
|
||||
builder.with_log_processor(
|
||||
ThreadBatchLogProcessor::builder(RedactingLogExporter::new(
|
||||
StderrLogExporter,
|
||||
gates,
|
||||
health,
|
||||
))
|
||||
.with_batch_config(batch_config)
|
||||
.build(),
|
||||
)
|
||||
}
|
||||
|
||||
fn wrap_console_metric_exporter(
|
||||
builder: MeterProviderBuilder,
|
||||
cfg: &ExternalOtelConfig,
|
||||
health: Arc<ExportHealth>,
|
||||
) -> MeterProviderBuilder {
|
||||
builder.with_reader(
|
||||
ThreadPeriodicReader::builder(ValidatingMetricExporter::new(
|
||||
StderrMetricExporter {
|
||||
temporality: temporality(cfg.temporality),
|
||||
},
|
||||
health,
|
||||
))
|
||||
.with_interval(cfg.metric_export_interval)
|
||||
.build(),
|
||||
)
|
||||
}
|
||||
|
||||
/// Build the providers per the resolved config. Returns `None` providers for
|
||||
/// signals whose exporter selection is `none`.
|
||||
pub(crate) fn build(
|
||||
cfg: &ExternalOtelConfig,
|
||||
gates: SharedGates,
|
||||
health: Arc<ExportHealth>,
|
||||
) -> Result<BuiltProviders, opentelemetry_otlp::ExporterBuildError> {
|
||||
// Build the shared blocking client only when the HTTP transport is
|
||||
// selected (`otlp_http` handles the dedicated-thread construction). A
|
||||
// build failure disables the external stream (caller warns) — it must
|
||||
// never panic the process.
|
||||
let needs_http_client = cfg.transport == OtlpTransport::HttpProtobuf
|
||||
&& (cfg.logs_exporter == ExporterSelection::Otlp
|
||||
|| cfg.metrics_exporter == ExporterSelection::Otlp);
|
||||
let http_client = needs_http_client
|
||||
.then(|| crate::otlp_http::build_blocking_client(cfg.timeout))
|
||||
.transpose()
|
||||
.map_err(opentelemetry_otlp::ExporterBuildError::InternalFailure)?;
|
||||
|
||||
// Console output is suppressed in the agent/headless entrypoints:
|
||||
// wrapping harnesses routinely capture stderr for diagnostics, and
|
||||
// interleaving periodic telemetry dumps there degrades those logs.
|
||||
let console_ok = !matches!(cfg.client.app_entrypoint.as_str(), "agent" | "headless");
|
||||
|
||||
let logger_provider = match cfg.logs_exporter {
|
||||
ExporterSelection::None => None,
|
||||
ExporterSelection::Console if !console_ok => {
|
||||
tracing::debug!(
|
||||
"external otel: console logs exporter suppressed in agent/headless entrypoint"
|
||||
);
|
||||
None
|
||||
}
|
||||
selection => {
|
||||
let batch_config = BatchConfigBuilder::default()
|
||||
.with_scheduled_delay(cfg.logs_export_interval)
|
||||
.with_max_export_batch_size(64)
|
||||
.build();
|
||||
let builder = SdkLoggerProvider::builder().with_resource(build_resource(cfg));
|
||||
let provider = match selection {
|
||||
ExporterSelection::Otlp => build_log_otlp_provider(
|
||||
builder,
|
||||
cfg,
|
||||
batch_config,
|
||||
http_client.as_ref(),
|
||||
gates.clone(),
|
||||
health.clone(),
|
||||
)?,
|
||||
_ => {
|
||||
wrap_console_log_exporter(builder, batch_config, gates.clone(), health.clone())
|
||||
}
|
||||
};
|
||||
Some(provider.build())
|
||||
}
|
||||
};
|
||||
|
||||
let meter_provider = match cfg.metrics_exporter {
|
||||
ExporterSelection::None => None,
|
||||
ExporterSelection::Console if !console_ok => {
|
||||
tracing::debug!(
|
||||
"external otel: console metrics exporter suppressed in agent/headless entrypoint"
|
||||
);
|
||||
None
|
||||
}
|
||||
selection => {
|
||||
let builder = SdkMeterProvider::builder().with_resource(build_resource(cfg));
|
||||
let provider = match selection {
|
||||
ExporterSelection::Otlp => {
|
||||
build_metric_otlp_provider(builder, cfg, http_client.as_ref(), health.clone())?
|
||||
}
|
||||
_ => wrap_console_metric_exporter(builder, cfg, health.clone()),
|
||||
};
|
||||
Some(provider.build())
|
||||
}
|
||||
};
|
||||
|
||||
Ok(BuiltProviders {
|
||||
logger_provider,
|
||||
meter_provider,
|
||||
})
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
use crate::external::config::ExternalOtelConfig;
|
||||
use std::collections::HashMap;
|
||||
|
||||
fn cfg_with_headers(headers: Vec<(String, String)>) -> ExternalOtelConfig {
|
||||
let mut cfg = ExternalOtelConfig::resolve_with(
|
||||
|name| match name {
|
||||
"GROK_EXTERNAL_OTEL" => Some("1".into()),
|
||||
"OTEL_LOGS_EXPORTER" => Some("otlp".into()),
|
||||
_ => None,
|
||||
},
|
||||
None,
|
||||
)
|
||||
.expect("test config must resolve");
|
||||
cfg.logs_headers = headers.clone();
|
||||
cfg.metrics_headers = headers;
|
||||
cfg
|
||||
}
|
||||
|
||||
/// Header-isolation invariant (T2): the outgoing header map equals
|
||||
/// exactly the parsed `OTEL_EXPORTER_OTLP_HEADERS` — no `Authorization`,
|
||||
/// `X-XAI-Token-Auth`, `x-userid`, or `x-teamid` unless customer-supplied
|
||||
/// (complement of the internal pipeline's
|
||||
/// `extra_headers_override_bearer_but_keep_static_identity`).
|
||||
#[test]
|
||||
fn exporter_headers_are_exactly_customer_headers() {
|
||||
let cfg = cfg_with_headers(vec![("x-collector-token".into(), "abc".into())]);
|
||||
let headers = customer_headers(&cfg.logs_headers);
|
||||
let expected: HashMap<String, String> =
|
||||
[("x-collector-token".to_string(), "abc".to_string())].into();
|
||||
assert_eq!(headers, expected);
|
||||
for forbidden in ["Authorization", "X-XAI-Token-Auth", "x-userid", "x-teamid"] {
|
||||
assert!(
|
||||
!headers.contains_key(forbidden),
|
||||
"{forbidden} must never be auto-attached to external exports"
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn customer_supplied_authorization_passes_through() {
|
||||
// The customer may auth their own collector however they want.
|
||||
let cfg = cfg_with_headers(vec![("Authorization".into(), "Bearer customer".into())]);
|
||||
assert_eq!(
|
||||
customer_headers(&cfg.logs_headers)
|
||||
.get("Authorization")
|
||||
.map(String::as_str),
|
||||
Some("Bearer customer")
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn exporter_metadata_is_customer_headers_only() {
|
||||
let cfg = cfg_with_headers(vec![
|
||||
("x-collector-token".into(), "abc".into()),
|
||||
("bad header".into(), "skip".into()),
|
||||
]);
|
||||
let metadata = customer_metadata(&cfg.logs_headers);
|
||||
assert_eq!(
|
||||
metadata
|
||||
.get("x-collector-token")
|
||||
.and_then(|v| v.to_str().ok()),
|
||||
Some("abc")
|
||||
);
|
||||
for forbidden in ["x-xai-token-auth", "x-userid", "x-teamid"] {
|
||||
assert!(metadata.get(forbidden).is_none());
|
||||
}
|
||||
}
|
||||
}
|
||||
248
crates/codegen/xai-grok-telemetry/src/external/redact.rs
vendored
Normal file
248
crates/codegen/xai-grok-telemetry/src/external/redact.rs
vendored
Normal file
|
|
@ -0,0 +1,248 @@
|
|||
//! Export-time fail-closed validators for the external stream.
|
||||
//!
|
||||
//! The primary redaction (typed-key schema, gating, secret scrub, truncation)
|
||||
//! happens at **emit time** in [`super::emit`], because `opentelemetry_sdk`
|
||||
//! 0.30 log records and metric data are not mutable from an exporter wrapper.
|
||||
//! These wrappers are the **authoritative chokepoint** anyway: they verify,
|
||||
//! per record/data point, that nothing reaches the wire that the emit path
|
||||
//! shouldn't have produced — and on any violation they *drop* (a record for
|
||||
//! logs, the whole export for metrics) rather than scrub in place. Dropping
|
||||
//! telemetry on a schema bug is acceptable; leaking is not.
|
||||
|
||||
use std::sync::Arc;
|
||||
use std::sync::atomic::{AtomicU64, Ordering};
|
||||
|
||||
use opentelemetry::InstrumentationScope;
|
||||
use opentelemetry::logs::AnyValue;
|
||||
use opentelemetry_sdk::Resource;
|
||||
use opentelemetry_sdk::error::OTelSdkResult;
|
||||
use opentelemetry_sdk::logs::{LogBatch, LogExporter, SdkLogRecord};
|
||||
use opentelemetry_sdk::metrics::Temporality;
|
||||
use opentelemetry_sdk::metrics::data::{AggregatedMetrics, MetricData, ResourceMetrics};
|
||||
use opentelemetry_sdk::metrics::exporter::PushMetricExporter;
|
||||
|
||||
use super::config::ContentGates;
|
||||
use super::schema::{Gate, external_allowed_keys, gate_for_key};
|
||||
|
||||
/// Shared, tighten-only view of the content gates. The remote kill switch may
|
||||
/// force gates off mid-run; the exporters re-read on every export.
|
||||
pub(crate) type SharedGates = Arc<parking_lot::RwLock<ContentGates>>;
|
||||
|
||||
/// Export-health counters (read by the internal `export_health` meta-event).
|
||||
#[derive(Debug, Default)]
|
||||
pub(crate) struct ExportHealth {
|
||||
/// Log records dropped by the validator.
|
||||
pub records_dropped: AtomicU64,
|
||||
/// Whole metric exports dropped by the validator.
|
||||
pub metric_exports_dropped: AtomicU64,
|
||||
/// Failed export attempts (transport errors), both signals.
|
||||
pub export_failures: AtomicU64,
|
||||
/// Successful export attempts, both signals.
|
||||
pub export_successes: AtomicU64,
|
||||
}
|
||||
|
||||
fn gate_open(gates: &ContentGates, gate: Gate) -> bool {
|
||||
match gate {
|
||||
Gate::UserPrompts => gates.log_user_prompts,
|
||||
Gate::ToolDetails => gates.log_tool_details,
|
||||
}
|
||||
}
|
||||
|
||||
/// `true` when this record is clean: every attribute key is schema-named,
|
||||
/// gated keys have their gate open, string values carry no secret shapes the
|
||||
/// emit path should have scrubbed, and the body is empty (`event.name` is the
|
||||
/// structured identity — external records carry no free-text body).
|
||||
fn record_is_clean(record: &SdkLogRecord, gates: &ContentGates) -> bool {
|
||||
if record.body().is_some() {
|
||||
tracing::debug!("external otel: dropping record with non-empty body");
|
||||
return false;
|
||||
}
|
||||
for (key, value) in record.attributes_iter() {
|
||||
let key_str = key.as_str();
|
||||
if !external_allowed_keys().contains(key_str) {
|
||||
tracing::debug!(
|
||||
key = key_str,
|
||||
"external otel: dropping record with non-schema key"
|
||||
);
|
||||
return false;
|
||||
}
|
||||
if let Some(gate) = gate_for_key(key_str)
|
||||
&& !gate_open(gates, gate)
|
||||
{
|
||||
tracing::debug!(
|
||||
key = key_str,
|
||||
"external otel: dropping record with closed-gate key"
|
||||
);
|
||||
return false;
|
||||
}
|
||||
match value {
|
||||
AnyValue::Int(_) | AnyValue::Double(_) | AnyValue::Boolean(_) => {}
|
||||
AnyValue::String(s) => {
|
||||
if crate::redact_common::redact_owned(s.as_str()).is_some() {
|
||||
tracing::debug!(
|
||||
key = key_str,
|
||||
"external otel: dropping record with unscrubbed string value"
|
||||
);
|
||||
return false;
|
||||
}
|
||||
}
|
||||
// Bytes / lists / maps / future variants are content (fail-closed).
|
||||
_ => {
|
||||
tracing::debug!(
|
||||
key = key_str,
|
||||
"external otel: dropping record with non-scalar value"
|
||||
);
|
||||
return false;
|
||||
}
|
||||
}
|
||||
}
|
||||
true
|
||||
}
|
||||
|
||||
/// Wraps the OTLP [`LogExporter`]; drops any record that violates the pinned
|
||||
/// schema before delegating.
|
||||
#[derive(Debug)]
|
||||
pub(crate) struct RedactingLogExporter<E> {
|
||||
inner: E,
|
||||
gates: SharedGates,
|
||||
health: Arc<ExportHealth>,
|
||||
}
|
||||
|
||||
impl<E> RedactingLogExporter<E> {
|
||||
pub(crate) fn new(inner: E, gates: SharedGates, health: Arc<ExportHealth>) -> Self {
|
||||
Self {
|
||||
inner,
|
||||
gates,
|
||||
health,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl<E: LogExporter> LogExporter for RedactingLogExporter<E> {
|
||||
fn export(
|
||||
&self,
|
||||
batch: LogBatch<'_>,
|
||||
) -> impl std::future::Future<Output = OTelSdkResult> + Send {
|
||||
let gates = *self.gates.read();
|
||||
async move {
|
||||
let clean: Vec<(&SdkLogRecord, &InstrumentationScope)> = batch
|
||||
.iter()
|
||||
.filter(|(record, _)| {
|
||||
let ok = record_is_clean(record, &gates);
|
||||
if !ok {
|
||||
self.health.records_dropped.fetch_add(1, Ordering::Relaxed);
|
||||
}
|
||||
ok
|
||||
})
|
||||
.collect();
|
||||
if clean.is_empty() {
|
||||
return Ok(());
|
||||
}
|
||||
let result = self.inner.export(LogBatch::new(&clean)).await;
|
||||
match &result {
|
||||
Ok(()) => self.health.export_successes.fetch_add(1, Ordering::Relaxed),
|
||||
Err(_) => self.health.export_failures.fetch_add(1, Ordering::Relaxed),
|
||||
};
|
||||
result
|
||||
}
|
||||
}
|
||||
|
||||
fn shutdown_with_timeout(&self, timeout: std::time::Duration) -> OTelSdkResult {
|
||||
self.inner.shutdown_with_timeout(timeout)
|
||||
}
|
||||
|
||||
fn set_resource(&mut self, resource: &Resource) {
|
||||
self.inner.set_resource(resource);
|
||||
}
|
||||
}
|
||||
|
||||
/// `true` when every data point's attribute keys are within the pinned
|
||||
/// metric-attribute set.
|
||||
fn metrics_are_clean(metrics: &ResourceMetrics) -> bool {
|
||||
fn keys_ok<'a>(attrs: impl Iterator<Item = &'a opentelemetry::KeyValue>) -> bool {
|
||||
for kv in attrs {
|
||||
let key = kv.key.as_str();
|
||||
if !super::schema::METRIC_ALLOWED_ATTR_KEYS.contains(&key) {
|
||||
tracing::debug!(
|
||||
key,
|
||||
"external otel: metric data point carries a non-schema attribute key"
|
||||
);
|
||||
return false;
|
||||
}
|
||||
}
|
||||
true
|
||||
}
|
||||
|
||||
fn data_ok<T>(data: &MetricData<T>) -> bool {
|
||||
match data {
|
||||
MetricData::Gauge(g) => g.data_points().all(|p| keys_ok(p.attributes())),
|
||||
MetricData::Sum(s) => s.data_points().all(|p| keys_ok(p.attributes())),
|
||||
MetricData::Histogram(h) => h.data_points().all(|p| keys_ok(p.attributes())),
|
||||
MetricData::ExponentialHistogram(h) => h.data_points().all(|p| keys_ok(p.attributes())),
|
||||
}
|
||||
}
|
||||
|
||||
metrics.scope_metrics().all(|scope| {
|
||||
scope.metrics().all(|metric| match metric.data() {
|
||||
AggregatedMetrics::F64(d) => data_ok(d),
|
||||
AggregatedMetrics::U64(d) => data_ok(d),
|
||||
AggregatedMetrics::I64(d) => data_ok(d),
|
||||
})
|
||||
})
|
||||
}
|
||||
|
||||
/// Wraps the OTLP `MetricExporter`. `opentelemetry_sdk` 0.30's
|
||||
/// `ResourceMetrics` read path is iterator-based and cannot be mutated, so on
|
||||
/// any attribute-key violation the wrapper **drops the entire export**
|
||||
/// (returns `Ok`, logs an internal warning, increments the export-health
|
||||
/// counter) rather than scrubbing in place. Coarse, but genuinely
|
||||
/// fail-closed.
|
||||
#[derive(Debug)]
|
||||
pub(crate) struct ValidatingMetricExporter<E> {
|
||||
inner: E,
|
||||
health: Arc<ExportHealth>,
|
||||
}
|
||||
|
||||
impl<E> ValidatingMetricExporter<E> {
|
||||
pub(crate) fn new(inner: E, health: Arc<ExportHealth>) -> Self {
|
||||
Self { inner, health }
|
||||
}
|
||||
}
|
||||
|
||||
impl<E: PushMetricExporter> PushMetricExporter for ValidatingMetricExporter<E> {
|
||||
fn export(
|
||||
&self,
|
||||
metrics: &ResourceMetrics,
|
||||
) -> impl std::future::Future<Output = OTelSdkResult> + Send {
|
||||
let clean = metrics_are_clean(metrics);
|
||||
async move {
|
||||
if !clean {
|
||||
self.health
|
||||
.metric_exports_dropped
|
||||
.fetch_add(1, Ordering::Relaxed);
|
||||
tracing::debug!(
|
||||
"external otel: dropped metric export (schema violation; fail-closed)"
|
||||
);
|
||||
return Ok(());
|
||||
}
|
||||
let result = self.inner.export(metrics).await;
|
||||
match &result {
|
||||
Ok(()) => self.health.export_successes.fetch_add(1, Ordering::Relaxed),
|
||||
Err(_) => self.health.export_failures.fetch_add(1, Ordering::Relaxed),
|
||||
};
|
||||
result
|
||||
}
|
||||
}
|
||||
|
||||
fn force_flush(&self) -> OTelSdkResult {
|
||||
self.inner.force_flush()
|
||||
}
|
||||
|
||||
fn shutdown_with_timeout(&self, timeout: std::time::Duration) -> OTelSdkResult {
|
||||
self.inner.shutdown_with_timeout(timeout)
|
||||
}
|
||||
|
||||
fn temporality(&self) -> Temporality {
|
||||
self.inner.temporality()
|
||||
}
|
||||
}
|
||||
1118
crates/codegen/xai-grok-telemetry/src/external/schema.rs
vendored
Normal file
1118
crates/codegen/xai-grok-telemetry/src/external/schema.rs
vendored
Normal file
File diff suppressed because it is too large
Load diff
991
crates/codegen/xai-grok-telemetry/src/external/tests.rs
vendored
Normal file
991
crates/codegen/xai-grok-telemetry/src/external/tests.rs
vendored
Normal file
|
|
@ -0,0 +1,991 @@
|
|||
//! Unit tests for the external stream: pinned allowlists, per-event schema
|
||||
//! snapshots, canary leak tests, gate enforcement, and the tighten-only
|
||||
//! remote policy. Everything asserting wire shape goes through the
|
||||
//! in-memory exporters *behind the export-time validators*, so the tests pin
|
||||
//! what actually leaves the process.
|
||||
|
||||
use super::config::ContentGates;
|
||||
use super::schema::{self, AttrValue, ExternalKey, ExternalRecord, MetricIncrement};
|
||||
use super::test_support::{TestStream, build, emit_event_into, emit_into};
|
||||
use crate::events;
|
||||
use opentelemetry::logs::AnyValue;
|
||||
|
||||
fn gates_off() -> ContentGates {
|
||||
ContentGates::default()
|
||||
}
|
||||
|
||||
fn gates_all_on() -> ContentGates {
|
||||
ContentGates {
|
||||
log_user_prompts: true,
|
||||
log_tool_details: true,
|
||||
}
|
||||
}
|
||||
|
||||
/// Exported (event_name, sorted attr key/value-debug pairs) for assertions.
|
||||
fn exported_events(stream: &TestStream) -> Vec<(String, Vec<(String, String)>)> {
|
||||
stream
|
||||
.logs
|
||||
.get_emitted_logs()
|
||||
.expect("in-memory logs")
|
||||
.iter()
|
||||
.map(|log| {
|
||||
let record = &log.record;
|
||||
let mut attrs: Vec<(String, String)> = record
|
||||
.attributes_iter()
|
||||
.map(|(k, v)| {
|
||||
let value = match v {
|
||||
AnyValue::String(s) => s.as_str().to_owned(),
|
||||
AnyValue::Int(i) => i.to_string(),
|
||||
AnyValue::Boolean(b) => b.to_string(),
|
||||
other => format!("{other:?}"),
|
||||
};
|
||||
(k.as_str().to_owned(), value)
|
||||
})
|
||||
.collect();
|
||||
attrs.sort();
|
||||
(record.event_name().unwrap_or("?").to_owned(), attrs)
|
||||
})
|
||||
.collect()
|
||||
}
|
||||
|
||||
fn attr_keys(event: &(String, Vec<(String, String)>)) -> Vec<&str> {
|
||||
event.1.iter().map(|(k, _)| k.as_str()).collect()
|
||||
}
|
||||
|
||||
fn attr(event: &(String, Vec<(String, String)>), key: &str) -> Option<String> {
|
||||
event
|
||||
.1
|
||||
.iter()
|
||||
.find(|(k, _)| k == key)
|
||||
.map(|(_, v)| v.clone())
|
||||
}
|
||||
|
||||
/// All exported metric (name, sorted data-point attr keys) pairs.
|
||||
fn exported_metric_names(stream: &TestStream) -> Vec<String> {
|
||||
stream
|
||||
.metrics
|
||||
.get_finished_metrics()
|
||||
.expect("in-memory metrics")
|
||||
.iter()
|
||||
.flat_map(|rm| {
|
||||
rm.scope_metrics()
|
||||
.flat_map(|s| s.metrics().map(|m| m.name().to_owned()))
|
||||
.collect::<Vec<_>>()
|
||||
})
|
||||
.collect()
|
||||
}
|
||||
|
||||
// ─────────────────────────────────────────────────────────────────────────────
|
||||
// Pinned allowlists
|
||||
// ─────────────────────────────────────────────────────────────────────────────
|
||||
|
||||
#[test]
|
||||
fn external_allowed_keys_are_pinned() {
|
||||
// Keep this an independent copy — don't reference ALL_KEYS or
|
||||
// ExternalKey::as_str, or the assert becomes a tautology and stops gating
|
||||
// schema changes. Adding a key exports a new field: confirm it carries no
|
||||
// user content, then update this pin.
|
||||
let expected: &[&str] = &[
|
||||
"session.id",
|
||||
"turn_number",
|
||||
"prompt.id",
|
||||
"event.sequence",
|
||||
"user.id",
|
||||
"organization.id",
|
||||
"team.id",
|
||||
"deployment.id",
|
||||
"model",
|
||||
"permission_mode",
|
||||
"mcp_server_count",
|
||||
"plugin_count",
|
||||
"skill_count",
|
||||
"hook_count",
|
||||
"memory_enabled",
|
||||
"is_git_repo",
|
||||
"client_identifier",
|
||||
"duration_secs",
|
||||
"turn_count",
|
||||
"tool_call_count",
|
||||
"compaction_count",
|
||||
"prompt_length",
|
||||
"prompt",
|
||||
"screen_mode",
|
||||
"outcome",
|
||||
"duration_ms",
|
||||
"error_category",
|
||||
"cancellation_category",
|
||||
"stop_reason",
|
||||
"input_tokens",
|
||||
"output_tokens",
|
||||
"reasoning_tokens",
|
||||
"cache_read_tokens",
|
||||
"status_code",
|
||||
"tool_name",
|
||||
"success",
|
||||
"file_extension",
|
||||
"tool_parameters",
|
||||
"file_path",
|
||||
"decision",
|
||||
"access_kind",
|
||||
"source",
|
||||
"status",
|
||||
"transport_type",
|
||||
"tool_count",
|
||||
"error_type",
|
||||
"mcp_server.name",
|
||||
"to_mode",
|
||||
"trigger",
|
||||
"skill_source",
|
||||
"skill.name",
|
||||
"install_kind",
|
||||
"plugin_scope",
|
||||
"plugin_name",
|
||||
"plugin_version",
|
||||
"compaction_trigger",
|
||||
"compaction_outcome",
|
||||
"tokens_before",
|
||||
"tokens_after",
|
||||
"phase",
|
||||
"subagent_type",
|
||||
"auth_method",
|
||||
"from_model",
|
||||
"to_model",
|
||||
"error_code",
|
||||
"tip",
|
||||
"action",
|
||||
];
|
||||
let actual: Vec<&str> = schema::ALL_KEYS.iter().map(|k| k.as_str()).collect();
|
||||
assert_eq!(
|
||||
actual, expected,
|
||||
"EXTERNAL_ALLOWED_KEYS changed: a new key is a wire-schema change — confirm it carries \
|
||||
no user content, then update this pin."
|
||||
);
|
||||
// Notably absent: the internal allowlist's file-path family.
|
||||
for forbidden in ["path", "cwd", "repo_path", "worktree", "gcs_path"] {
|
||||
assert!(
|
||||
!schema::external_allowed_keys().contains(forbidden),
|
||||
"{forbidden} must never be externally allowlisted"
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn metric_attr_keys_are_pinned() {
|
||||
let expected: &[&str] = &[
|
||||
"type",
|
||||
"model",
|
||||
"outcome",
|
||||
"tool_name",
|
||||
"decision",
|
||||
"access_kind",
|
||||
"permission_mode",
|
||||
"error_category",
|
||||
"session.id",
|
||||
"app.version",
|
||||
"user.id",
|
||||
"organization.id",
|
||||
"team.id",
|
||||
"deployment.id",
|
||||
];
|
||||
assert_eq!(
|
||||
schema::METRIC_ALLOWED_ATTR_KEYS,
|
||||
expected,
|
||||
"metric attribute keys changed — wire-schema change"
|
||||
);
|
||||
assert!(
|
||||
!schema::METRIC_ALLOWED_ATTR_KEYS.contains(&"prompt.id"),
|
||||
"prompt.id is events-only (unbounded cardinality on metrics)"
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn event_names_are_pinned() {
|
||||
use schema::ExternalEventName as E;
|
||||
let expected: &[(E, &str)] = &[
|
||||
(E::SessionStart, "grok_code.session_start"),
|
||||
(E::SessionEnd, "grok_code.session_end"),
|
||||
(E::UserPrompt, "grok_code.user_prompt"),
|
||||
(E::TurnCompleted, "grok_code.turn_completed"),
|
||||
(E::ApiRequest, "grok_code.api_request"),
|
||||
(E::ApiError, "grok_code.api_error"),
|
||||
(E::ToolResult, "grok_code.tool_result"),
|
||||
(E::ToolDecision, "grok_code.tool_decision"),
|
||||
(E::McpServerConnection, "grok_code.mcp_server_connection"),
|
||||
(
|
||||
E::PermissionModeChanged,
|
||||
"grok_code.permission_mode_changed",
|
||||
),
|
||||
(E::SkillActivated, "grok_code.skill_activated"),
|
||||
(E::PluginLoaded, "grok_code.plugin_loaded"),
|
||||
(E::Compaction, "grok_code.compaction"),
|
||||
(E::Subagent, "grok_code.subagent"),
|
||||
(E::Auth, "grok_code.auth"),
|
||||
(E::InternalError, "grok_code.internal_error"),
|
||||
(E::ModelSwitched, "grok_code.model_switched"),
|
||||
(E::ContextualTip, "grok_code.contextual_tip"),
|
||||
];
|
||||
assert_eq!(expected.len(), <E as strum::EnumCount>::COUNT);
|
||||
for (variant, name) in expected {
|
||||
assert_eq!(variant.as_str(), *name, "event name is a wire commitment");
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn client_identifier_allowlist_is_pinned() {
|
||||
let expected: &[&str] = &[
|
||||
"grok-pager",
|
||||
"grok-tui",
|
||||
"grok-shell",
|
||||
"grok-web",
|
||||
"grok-desktop",
|
||||
"grok-code-extension",
|
||||
"nebula",
|
||||
"zed",
|
||||
];
|
||||
assert_eq!(schema::KNOWN_CLIENT_IDENTIFIERS, expected);
|
||||
assert_eq!(
|
||||
schema::sanitize_client_identifier("grok-pager"),
|
||||
"grok-pager"
|
||||
);
|
||||
assert_eq!(
|
||||
schema::sanitize_client_identifier("Evil Corp Internal Tool v2"),
|
||||
"other",
|
||||
"unknown client identifiers are externally controlled free text and must collapse"
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn screen_mode_allowlist_is_pinned() {
|
||||
let expected: &[&str] = &["fullscreen", "inline", "minimal", "headless"];
|
||||
assert_eq!(schema::KNOWN_SCREEN_MODES, expected);
|
||||
assert_eq!(schema::sanitize_screen_mode("minimal"), "minimal");
|
||||
assert_eq!(
|
||||
schema::sanitize_screen_mode("my-custom-fork-mode"),
|
||||
"other",
|
||||
"unknown screen modes are externally controlled free text and must collapse"
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn tool_name_sanitization() {
|
||||
assert_eq!(schema::sanitize_tool_name("read_file"), "read_file");
|
||||
assert_eq!(
|
||||
schema::sanitize_tool_name("nebula__post_message"),
|
||||
"mcp_tool"
|
||||
);
|
||||
assert_eq!(
|
||||
schema::sanitize_tool_name("SuperSecretProjectTool"),
|
||||
"custom_tool",
|
||||
"unknown tool names must not pass verbatim"
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn file_extension_reduction() {
|
||||
assert_eq!(
|
||||
schema::file_extension("/Users/alice/proj/main.rs"),
|
||||
Some("rs".into())
|
||||
);
|
||||
assert_eq!(schema::file_extension("src/App.TSX"), Some("tsx".into()));
|
||||
assert_eq!(schema::file_extension("Makefile"), None);
|
||||
// 10-char cap.
|
||||
assert_eq!(
|
||||
schema::file_extension("x.aaaaaaaaaaaaaaaa"),
|
||||
Some("aaaaaaaaaa".into())
|
||||
);
|
||||
}
|
||||
|
||||
// ─────────────────────────────────────────────────────────────────────────────
|
||||
// Schema snapshots (gates off / on) through the wire-view harness
|
||||
// ─────────────────────────────────────────────────────────────────────────────
|
||||
|
||||
fn sentinel_session_harness() -> events::SessionHarness {
|
||||
events::SessionHarness {
|
||||
session_id: "sess-1".into(),
|
||||
client_identifier: Some("grok-pager".into()),
|
||||
model_id: "grok-4".into(),
|
||||
agent_name: "grok-build-plan".into(),
|
||||
permission_mode: crate::enums::PermissionMode::Ask,
|
||||
mcp_server_names: vec!["secret-server".into(), "other".into()],
|
||||
plugin_names: vec!["p1".into()],
|
||||
skill_names: vec!["s1".into(), "s2".into(), "s3".into()],
|
||||
lsp_server_names: vec![],
|
||||
hook_names: vec!["h1".into()],
|
||||
agents_md_dir_names: vec!["proj".into()],
|
||||
memory_enabled: true,
|
||||
is_git_repo: true,
|
||||
auto_update: None,
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn session_start_snapshot_counts_not_names() {
|
||||
let stream = build(gates_off());
|
||||
emit_event_into(&stream, &sentinel_session_harness());
|
||||
let events = exported_events(&stream);
|
||||
assert_eq!(events.len(), 1);
|
||||
let ev = &events[0];
|
||||
assert_eq!(ev.0, "grok_code.session_start");
|
||||
let mut keys = attr_keys(ev);
|
||||
keys.sort();
|
||||
assert_eq!(
|
||||
keys,
|
||||
vec![
|
||||
"client_identifier",
|
||||
"event.sequence",
|
||||
"hook_count",
|
||||
"is_git_repo",
|
||||
"mcp_server_count",
|
||||
"memory_enabled",
|
||||
"model",
|
||||
"permission_mode",
|
||||
"plugin_count",
|
||||
"session.id",
|
||||
"skill_count",
|
||||
]
|
||||
);
|
||||
assert_eq!(attr(ev, "mcp_server_count").as_deref(), Some("2"));
|
||||
assert_eq!(attr(ev, "skill_count").as_deref(), Some("3"));
|
||||
assert_eq!(attr(ev, "session.id").as_deref(), Some("sess-1"));
|
||||
// Names (MCP/plugin/skill/hook) never appear, only counts.
|
||||
let blob = format!("{events:?}");
|
||||
assert!(!blob.contains("secret-server"), "MCP name leaked: {blob}");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn session_new_increments_session_count_only() {
|
||||
let stream = build(gates_off());
|
||||
emit_event_into(
|
||||
&stream,
|
||||
&events::SessionNew {
|
||||
session_id: "sess-2".into(),
|
||||
client_identifier: None,
|
||||
client_version: None,
|
||||
is_git_repo: false,
|
||||
permission_mode: crate::enums::PermissionMode::Ask,
|
||||
},
|
||||
);
|
||||
assert!(exported_events(&stream).is_empty(), "metric-only mapping");
|
||||
let names = exported_metric_names(&stream);
|
||||
assert_eq!(names, vec!["grok_code.session.count".to_owned()]);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn api_request_snapshot_and_token_usage() {
|
||||
let stream = build(gates_off());
|
||||
emit_event_into(
|
||||
&stream,
|
||||
&events::ModelResponseReceived {
|
||||
model_id: "grok-4".into(),
|
||||
duration_ms: 1200,
|
||||
stop_reason: Some("stop".into()),
|
||||
prompt_tokens: Some(100),
|
||||
completion_tokens: Some(50),
|
||||
reasoning_tokens: Some(25),
|
||||
cached_prompt_tokens: None,
|
||||
},
|
||||
);
|
||||
let events = exported_events(&stream);
|
||||
let ev = &events[0];
|
||||
assert_eq!(ev.0, "grok_code.api_request");
|
||||
assert_eq!(attr(ev, "input_tokens").as_deref(), Some("100"));
|
||||
assert_eq!(attr(ev, "output_tokens").as_deref(), Some("50"));
|
||||
assert_eq!(attr(ev, "reasoning_tokens").as_deref(), Some("25"));
|
||||
assert_eq!(attr(ev, "cache_read_tokens"), None);
|
||||
assert_eq!(
|
||||
exported_metric_names(&stream),
|
||||
vec!["grok_code.token.usage"]
|
||||
);
|
||||
}
|
||||
|
||||
/// One failed turn ⇒ exactly one `error.count` increment, even though the
|
||||
/// failure emits `ApiError` (and possibly `RateLimitHit`) *alongside*
|
||||
/// `TurnCompleted{Error}`. `TurnCompleted{Error}` is the single increment
|
||||
/// source; the api_error log events carry no metric (Bugbot regression:
|
||||
/// double-counted errors at customer collectors).
|
||||
#[test]
|
||||
fn one_failed_turn_increments_error_count_exactly_once() {
|
||||
let stream = build(gates_off());
|
||||
// The turn-error path emits all three for a rate-limited failure.
|
||||
emit_event_into(
|
||||
&stream,
|
||||
&events::RateLimitHit {
|
||||
model_id: "grok-4".into(),
|
||||
attempts: 3,
|
||||
},
|
||||
);
|
||||
emit_event_into(
|
||||
&stream,
|
||||
&events::ApiError {
|
||||
error_category: "rate_limit".into(),
|
||||
model_id: "grok-4".into(),
|
||||
status_code: Some(429),
|
||||
duration_ms: Some(10),
|
||||
},
|
||||
);
|
||||
emit_event_into(
|
||||
&stream,
|
||||
&events::TurnCompleted {
|
||||
outcome: events::Outcome::Error,
|
||||
duration_ms: 10,
|
||||
tool_call_count: 0,
|
||||
model_id: "grok-4".into(),
|
||||
cancellation_category: None,
|
||||
error_category: Some("rate_limit".into()),
|
||||
},
|
||||
);
|
||||
// Both api_error events exported as log records…
|
||||
let names: Vec<String> = exported_events(&stream)
|
||||
.iter()
|
||||
.map(|e| e.0.clone())
|
||||
.collect();
|
||||
assert_eq!(
|
||||
names.iter().filter(|n| *n == "grok_code.api_error").count(),
|
||||
2
|
||||
);
|
||||
// …but error.count incremented exactly once.
|
||||
let total: u64 = stream
|
||||
.metrics
|
||||
.get_finished_metrics()
|
||||
.unwrap()
|
||||
.iter()
|
||||
.flat_map(|rm| rm.scope_metrics())
|
||||
.flat_map(|s| s.metrics())
|
||||
.filter(|m| m.name() == "grok_code.error.count")
|
||||
.map(|m| match m.data() {
|
||||
opentelemetry_sdk::metrics::data::AggregatedMetrics::U64(
|
||||
opentelemetry_sdk::metrics::data::MetricData::Sum(sum),
|
||||
) => sum.data_points().map(|p| p.value()).sum::<u64>(),
|
||||
_ => 0,
|
||||
})
|
||||
.sum();
|
||||
assert_eq!(total, 1, "one failed turn must count exactly one error");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn turn_error_increments_error_count() {
|
||||
let stream = build(gates_off());
|
||||
emit_event_into(
|
||||
&stream,
|
||||
&events::TurnCompleted {
|
||||
outcome: events::Outcome::Error,
|
||||
duration_ms: 10,
|
||||
tool_call_count: 0,
|
||||
model_id: "grok-4".into(),
|
||||
cancellation_category: None,
|
||||
error_category: Some("server_error".into()),
|
||||
},
|
||||
);
|
||||
let mut names = exported_metric_names(&stream);
|
||||
names.sort();
|
||||
assert_eq!(names, vec!["grok_code.error.count", "grok_code.turn.count"]);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn tool_result_gates_off_collapses_and_reduces() {
|
||||
let stream = build(gates_off());
|
||||
emit_event_into(
|
||||
&stream,
|
||||
&events::ToolCallCompleted {
|
||||
tool_name: "nebula__post_message".into(),
|
||||
outcome: xai_file_utils::events::types::ToolOutcome::Success,
|
||||
duration_ms: 42,
|
||||
file_path: Some("/Users/alice/secret-project/main.rs".into()),
|
||||
parameters: Some(serde_json::json!({"text": "CANARY_TOOL_ARGS"})),
|
||||
},
|
||||
);
|
||||
let events = exported_events(&stream);
|
||||
let ev = &events[0];
|
||||
assert_eq!(ev.0, "grok_code.tool_result");
|
||||
assert_eq!(attr(ev, "tool_name").as_deref(), Some("mcp_tool"));
|
||||
assert_eq!(attr(ev, "file_extension").as_deref(), Some("rs"));
|
||||
assert_eq!(attr(ev, "file_path"), None, "full path is details-gated");
|
||||
assert_eq!(
|
||||
attr(ev, "tool_parameters"),
|
||||
None,
|
||||
"params are details-gated"
|
||||
);
|
||||
let blob = format!("{events:?}");
|
||||
assert!(
|
||||
!blob.contains("CANARY_TOOL_ARGS"),
|
||||
"tool args leaked: {blob}"
|
||||
);
|
||||
assert!(!blob.contains("secret-project"), "path leaked: {blob}");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn tool_result_details_gate_exposes_verbatim_scrubbed() {
|
||||
let stream = build(gates_all_on());
|
||||
// Use the *real* home dir: `redact_user_paths` collapses the current
|
||||
// user's home (env-derived), not arbitrary foreign paths.
|
||||
let home = dirs::home_dir()
|
||||
.map(|h| h.to_string_lossy().into_owned())
|
||||
.unwrap_or_else(|| "/home/testuser".into());
|
||||
let path = format!("{home}/proj/main.rs");
|
||||
emit_event_into(
|
||||
&stream,
|
||||
&events::ToolCallCompleted {
|
||||
tool_name: "nebula__post_message".into(),
|
||||
outcome: xai_file_utils::events::types::ToolOutcome::Success,
|
||||
duration_ms: 42,
|
||||
file_path: Some(path.clone()),
|
||||
parameters: Some(serde_json::json!({"key": "sk-CANARYabcdefghij1234567890"})),
|
||||
},
|
||||
);
|
||||
let events = exported_events(&stream);
|
||||
let ev = &events[0];
|
||||
assert_eq!(
|
||||
attr(ev, "tool_name").as_deref(),
|
||||
Some("nebula__post_message"),
|
||||
"details gate exposes the verbatim name"
|
||||
);
|
||||
let exported_path = attr(ev, "file_path").expect("details gate exposes the path");
|
||||
// Secrets are STILL scrubbed inside gated content; the home dir collapses.
|
||||
let blob = format!("{events:?}");
|
||||
assert!(
|
||||
!blob.contains("CANARY"),
|
||||
"secret inside gated params leaked: {blob}"
|
||||
);
|
||||
assert!(
|
||||
!exported_path.contains(&home),
|
||||
"home dir not collapsed in gated path: {exported_path}"
|
||||
);
|
||||
assert!(exported_path.contains("proj/main.rs"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn user_prompt_gates_off_drops_text() {
|
||||
let stream = build(gates_off());
|
||||
emit_event_into(
|
||||
&stream,
|
||||
&events::PromptSubmitted {
|
||||
prompt_length: 26,
|
||||
model_id: "grok-4".into(),
|
||||
client_identifier: None,
|
||||
screen_mode: Some("minimal".into()),
|
||||
prompt_text: Some("CANARY_PROMPT secret user text".into()),
|
||||
},
|
||||
);
|
||||
let events = exported_events(&stream);
|
||||
let ev = &events[0];
|
||||
assert_eq!(ev.0, "grok_code.user_prompt");
|
||||
assert_eq!(attr(ev, "prompt_length").as_deref(), Some("26"));
|
||||
// screen_mode is ungated session metadata, not prompt content.
|
||||
assert_eq!(attr(ev, "screen_mode").as_deref(), Some("minimal"));
|
||||
assert_eq!(attr(ev, "prompt"), None);
|
||||
let blob = format!("{events:?}");
|
||||
assert!(
|
||||
!blob.contains("CANARY_PROMPT"),
|
||||
"prompt text leaked: {blob}"
|
||||
);
|
||||
}
|
||||
|
||||
/// `screen_mode` is externally controlled free text (ACP `_meta.screenMode`);
|
||||
/// unknown values must collapse to `"other"` on the wire, and an absent value
|
||||
/// must emit no attribute at all.
|
||||
#[test]
|
||||
fn user_prompt_screen_mode_sanitized_and_optional() {
|
||||
let stream = build(gates_off());
|
||||
emit_event_into(
|
||||
&stream,
|
||||
&events::PromptSubmitted {
|
||||
prompt_length: 5,
|
||||
model_id: "grok-4".into(),
|
||||
client_identifier: None,
|
||||
screen_mode: Some("Evil Free Text".into()),
|
||||
prompt_text: None,
|
||||
},
|
||||
);
|
||||
emit_event_into(
|
||||
&stream,
|
||||
&events::PromptSubmitted {
|
||||
prompt_length: 5,
|
||||
model_id: "grok-4".into(),
|
||||
client_identifier: None,
|
||||
screen_mode: None,
|
||||
prompt_text: None,
|
||||
},
|
||||
);
|
||||
let events = exported_events(&stream);
|
||||
assert_eq!(attr(&events[0], "screen_mode").as_deref(), Some("other"));
|
||||
assert_eq!(attr(&events[1], "screen_mode"), None);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn user_prompt_gate_on_exports_scrubbed_text() {
|
||||
let stream = build(gates_all_on());
|
||||
emit_event_into(
|
||||
&stream,
|
||||
&events::PromptSubmitted {
|
||||
prompt_length: 10,
|
||||
model_id: "grok-4".into(),
|
||||
client_identifier: None,
|
||||
screen_mode: None,
|
||||
prompt_text: Some("fix the bug; token sk-CANARYabcdefghij1234567890".into()),
|
||||
},
|
||||
);
|
||||
let events = exported_events(&stream);
|
||||
let ev = &events[0];
|
||||
let prompt = attr(ev, "prompt").expect("gate on ⇒ prompt exported");
|
||||
assert!(prompt.contains("fix the bug"));
|
||||
assert!(
|
||||
!prompt.contains("CANARY"),
|
||||
"secret inside prompt not scrubbed"
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn mcp_connection_collapses_server_name_by_default() {
|
||||
let stream = build(gates_off());
|
||||
emit_event_into(
|
||||
&stream,
|
||||
&events::McpServerFailed {
|
||||
server_name: "corp-internal-jira".into(),
|
||||
error_type: events::McpErrorType::Timeout,
|
||||
duration_ms: 1000,
|
||||
timeout_sec: 30,
|
||||
},
|
||||
);
|
||||
let events = exported_events(&stream);
|
||||
let ev = &events[0];
|
||||
assert_eq!(ev.0, "grok_code.mcp_server_connection");
|
||||
assert_eq!(attr(ev, "status").as_deref(), Some("failed"));
|
||||
assert_eq!(attr(ev, "mcp_server.name").as_deref(), Some("mcp_server"));
|
||||
assert_eq!(attr(ev, "error_type").as_deref(), Some("timeout"));
|
||||
assert!(!format!("{events:?}").contains("corp-internal-jira"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn tool_decision_snapshot() {
|
||||
let stream = build(gates_off());
|
||||
emit_event_into(
|
||||
&stream,
|
||||
&events::PermissionDecisionPayload {
|
||||
tool_name: "run_terminal_cmd".into(),
|
||||
access_kind: events::AccessKind::Bash,
|
||||
decision: events::PermissionOutcome::Deny,
|
||||
wait_ms: 1500,
|
||||
permission_mode: crate::enums::PermissionMode::Ask,
|
||||
source: Some("user_reject".into()),
|
||||
subagent_session_id: None,
|
||||
subagent_type: None,
|
||||
},
|
||||
);
|
||||
let events = exported_events(&stream);
|
||||
let ev = &events[0];
|
||||
assert_eq!(ev.0, "grok_code.tool_decision");
|
||||
assert_eq!(attr(ev, "tool_name").as_deref(), Some("run_terminal_cmd"));
|
||||
assert_eq!(attr(ev, "decision").as_deref(), Some("deny"));
|
||||
assert_eq!(attr(ev, "access_kind").as_deref(), Some("bash"));
|
||||
assert_eq!(attr(ev, "permission_mode").as_deref(), Some("ask"));
|
||||
assert_eq!(attr(ev, "source").as_deref(), Some("user_reject"));
|
||||
assert_eq!(
|
||||
exported_metric_names(&stream),
|
||||
vec!["grok_code.tool.decision"]
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn skill_activated_name_gated() {
|
||||
let stream = build(gates_off());
|
||||
emit_event_into(
|
||||
&stream,
|
||||
&events::SkillDispatched {
|
||||
skill_name: "internal-deploy-runbook".into(),
|
||||
plugin_source: None,
|
||||
},
|
||||
);
|
||||
let events = exported_events(&stream);
|
||||
let ev = &events[0];
|
||||
assert_eq!(attr(ev, "skill_source").as_deref(), Some("local"));
|
||||
assert_eq!(attr(ev, "skill.name"), None);
|
||||
assert!(!format!("{events:?}").contains("internal-deploy-runbook"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn contextual_tip_maps_every_tip_and_action() {
|
||||
use events::{ContextualTipAction as A, ContextualTipKind as K};
|
||||
let cases = [
|
||||
(K::Undo, A::Shown, "undo", "shown"),
|
||||
(K::Undo, A::Accepted, "undo", "accepted"),
|
||||
(K::PlanMode, A::Shown, "plan_mode", "shown"),
|
||||
(K::PlanMode, A::Accepted, "plan_mode", "accepted"),
|
||||
(K::ImageInput, A::Shown, "image_input", "shown"),
|
||||
(K::ImageInput, A::Accepted, "image_input", "accepted"),
|
||||
(K::SendNow, A::Shown, "send_now", "shown"),
|
||||
(K::SendNow, A::Accepted, "send_now", "accepted"),
|
||||
(K::SmallScreen, A::Shown, "small_screen", "shown"),
|
||||
(K::SmallScreen, A::Accepted, "small_screen", "accepted"),
|
||||
(K::WordSelect, A::Shown, "word_select", "shown"),
|
||||
(K::WordSelect, A::Accepted, "word_select", "accepted"),
|
||||
];
|
||||
for (tip, action, tip_label, action_label) in cases {
|
||||
let stream = build(gates_off());
|
||||
emit_event_into(&stream, &events::ContextualTip { tip, action });
|
||||
let events = exported_events(&stream);
|
||||
let ev = &events[0];
|
||||
assert_eq!(ev.0, "grok_code.contextual_tip");
|
||||
assert_eq!(attr(ev, "tip").as_deref(), Some(tip_label));
|
||||
assert_eq!(attr(ev, "action").as_deref(), Some(action_label));
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn unmapped_events_produce_nothing() {
|
||||
use crate::events::TelemetryEvent as _;
|
||||
// ~70 events without an `external = …` arm cost nothing and export nothing.
|
||||
let ev = events::SlashCommandUsed {
|
||||
command: "secret command".into(),
|
||||
args_provided: true,
|
||||
};
|
||||
assert!(ev.external_record().is_none());
|
||||
}
|
||||
|
||||
/// Workspace-origin exclusion: events emitted exclusively via
|
||||
/// `EmitterOrigin::Workspace` (`log_session_event_with_origin`) must not carry
|
||||
/// an external mapping — the fan-out hook deliberately lives only in the
|
||||
/// Shell-origin wrappers. The workspace-only surface today is the
|
||||
/// xai-grok-workspace sampler events, which live outside this crate and have
|
||||
/// no `telemetry_event!` binding here; this pin guards the in-crate set.
|
||||
#[test]
|
||||
fn workspace_only_events_have_no_external_mapping() {
|
||||
use crate::events::TelemetryEvent as _;
|
||||
// Trace-upload lifecycle events are session-metrics/internal-only.
|
||||
assert!(
|
||||
crate::session_metrics::TraceUploadAttempted {
|
||||
session_id: String::new(),
|
||||
turn_number: 0,
|
||||
upload_method: "proxy".into(),
|
||||
}
|
||||
.external_record()
|
||||
.is_none()
|
||||
);
|
||||
}
|
||||
|
||||
// ─────────────────────────────────────────────────────────────────────────────
|
||||
// Emit-path behavior: ctx injection, sequence, identity, truncation
|
||||
// ─────────────────────────────────────────────────────────────────────────────
|
||||
|
||||
#[test]
|
||||
fn mapping_supplied_session_id_wins_and_sequence_increments() {
|
||||
let stream = build(gates_off());
|
||||
emit_event_into(&stream, &sentinel_session_harness());
|
||||
emit_event_into(&stream, &sentinel_session_harness());
|
||||
let events = exported_events(&stream);
|
||||
assert_eq!(events.len(), 2);
|
||||
assert_eq!(attr(&events[0], "event.sequence").as_deref(), Some("0"));
|
||||
assert_eq!(attr(&events[1], "event.sequence").as_deref(), Some("1"));
|
||||
assert_eq!(attr(&events[0], "session.id").as_deref(), Some("sess-1"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn identity_attrs_attached_when_set_and_blank_ids_never_export() {
|
||||
let stream = build(gates_off());
|
||||
super::set_identity_on(
|
||||
&stream.ext,
|
||||
super::IdentityAttrs {
|
||||
user_id: Some("user-42".into()),
|
||||
organization_id: Some(String::new()), // blank: must not export
|
||||
team_id: None,
|
||||
deployment_id: Some("dep-7".into()),
|
||||
},
|
||||
);
|
||||
emit_event_into(&stream, &sentinel_session_harness());
|
||||
let events = exported_events(&stream);
|
||||
let ev = &events[0];
|
||||
assert_eq!(attr(ev, "user.id").as_deref(), Some("user-42"));
|
||||
assert_eq!(attr(ev, "deployment.id").as_deref(), Some("dep-7"));
|
||||
assert_eq!(attr(ev, "organization.id"), None, "blank ids never export");
|
||||
assert_eq!(attr(ev, "team.id"), None);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn long_attr_values_truncated() {
|
||||
let stream = build(gates_off());
|
||||
let long_model = "m".repeat(1000);
|
||||
emit_into(
|
||||
&stream,
|
||||
ExternalRecord {
|
||||
event: Some(schema::ExternalEventName::ApiRequest),
|
||||
attrs: vec![(ExternalKey::Model, AttrValue::Str(long_model))],
|
||||
gated: vec![],
|
||||
metrics: vec![],
|
||||
},
|
||||
);
|
||||
let events = exported_events(&stream);
|
||||
let model = attr(&events[0], "model").unwrap();
|
||||
assert!(
|
||||
model.len() < 200,
|
||||
"value not truncated: {} chars",
|
||||
model.len()
|
||||
);
|
||||
assert!(model.ends_with("…[truncated]"));
|
||||
}
|
||||
|
||||
// ─────────────────────────────────────────────────────────────────────────────
|
||||
// Export-time validators (fail-closed)
|
||||
// ─────────────────────────────────────────────────────────────────────────────
|
||||
|
||||
#[test]
|
||||
fn validating_metric_exporter_drops_export_on_bad_attr_key() {
|
||||
use opentelemetry::metrics::MeterProvider as _;
|
||||
let stream = build(gates_off());
|
||||
// Bypass emit.rs: increment with an attribute key outside the pinned set.
|
||||
let meter = stream
|
||||
.ext
|
||||
.meter_provider
|
||||
.as_ref()
|
||||
.unwrap()
|
||||
.meter(schema::SCOPE_NAME);
|
||||
let rogue = meter.u64_counter("grok_code.session.count").build();
|
||||
rogue.add(
|
||||
1,
|
||||
&[opentelemetry::KeyValue::new("prompt", "CANARY_METRIC_LEAK")],
|
||||
);
|
||||
stream
|
||||
.ext
|
||||
.meter_provider
|
||||
.as_ref()
|
||||
.unwrap()
|
||||
.force_flush()
|
||||
.unwrap();
|
||||
let exported = stream.metrics.get_finished_metrics().unwrap();
|
||||
let blob = format!("{exported:?}");
|
||||
assert!(
|
||||
!blob.contains("CANARY_METRIC_LEAK"),
|
||||
"metric export with rogue attr must be dropped entirely: {blob}"
|
||||
);
|
||||
assert!(
|
||||
stream
|
||||
.ext
|
||||
.health
|
||||
.metric_exports_dropped
|
||||
.load(std::sync::atomic::Ordering::Relaxed)
|
||||
>= 1
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn redacting_log_exporter_drops_record_with_closed_gate_key() {
|
||||
// A bug that attaches a gated key with the gate off must be caught at the
|
||||
// exporter even though emit.rs should never produce it.
|
||||
let stream = build(gates_off());
|
||||
use opentelemetry::logs::{LogRecord as _, Logger as _};
|
||||
let logger = stream.ext.logger.as_ref().unwrap();
|
||||
let mut record = logger.create_log_record();
|
||||
record.set_event_name("grok_code.user_prompt");
|
||||
record.add_attribute("prompt", "CANARY_GATED_LEAK");
|
||||
logger.emit(record);
|
||||
stream
|
||||
.ext
|
||||
.logger_provider
|
||||
.as_ref()
|
||||
.unwrap()
|
||||
.force_flush()
|
||||
.unwrap();
|
||||
let blob = format!("{:?}", stream.logs.get_emitted_logs().unwrap());
|
||||
assert!(
|
||||
!blob.contains("CANARY_GATED_LEAK"),
|
||||
"closed-gate key must be dropped by the exporter: {blob}"
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn redacting_log_exporter_drops_record_with_unknown_key() {
|
||||
let stream = build(gates_off());
|
||||
use opentelemetry::logs::{LogRecord as _, Logger as _};
|
||||
let logger = stream.ext.logger.as_ref().unwrap();
|
||||
let mut record = logger.create_log_record();
|
||||
record.set_event_name("grok_code.api_request");
|
||||
record.add_attribute("command", "echo CANARY_UNKNOWN_KEY");
|
||||
logger.emit(record);
|
||||
stream
|
||||
.ext
|
||||
.logger_provider
|
||||
.as_ref()
|
||||
.unwrap()
|
||||
.force_flush()
|
||||
.unwrap();
|
||||
let blob = format!("{:?}", stream.logs.get_emitted_logs().unwrap());
|
||||
assert!(
|
||||
!blob.contains("CANARY_UNKNOWN_KEY"),
|
||||
"unknown key leaked: {blob}"
|
||||
);
|
||||
}
|
||||
|
||||
// ─────────────────────────────────────────────────────────────────────────────
|
||||
// Remote policy: tighten-only
|
||||
// ─────────────────────────────────────────────────────────────────────────────
|
||||
|
||||
#[test]
|
||||
fn remote_force_disable_stops_emission() {
|
||||
let stream = build(gates_off());
|
||||
super::apply_remote_policy_on(
|
||||
&stream.ext,
|
||||
super::ExternalOtelRemotePolicy {
|
||||
force_disable: true,
|
||||
lock_content_gates: false,
|
||||
},
|
||||
);
|
||||
assert!(
|
||||
!stream.ext.active.load(std::sync::atomic::Ordering::Relaxed),
|
||||
"force_disable must clear the emission gate"
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn remote_gate_lock_forces_gates_off_and_never_on() {
|
||||
let stream = build(gates_all_on());
|
||||
super::apply_remote_policy_on(
|
||||
&stream.ext,
|
||||
super::ExternalOtelRemotePolicy {
|
||||
force_disable: false,
|
||||
lock_content_gates: true,
|
||||
},
|
||||
);
|
||||
assert_eq!(*stream.ext.gates.read(), ContentGates::default());
|
||||
// The policy carries no loosen/enable direction by construction: applying
|
||||
// a default policy to an off-gates stream changes nothing.
|
||||
let stream2 = build(gates_off());
|
||||
super::apply_remote_policy_on(&stream2.ext, super::ExternalOtelRemotePolicy::default());
|
||||
assert_eq!(*stream2.ext.gates.read(), ContentGates::default());
|
||||
assert!(
|
||||
stream2
|
||||
.ext
|
||||
.active
|
||||
.load(std::sync::atomic::Ordering::Relaxed)
|
||||
);
|
||||
}
|
||||
|
||||
// ─────────────────────────────────────────────────────────────────────────────
|
||||
// Metric increment derivation
|
||||
// ─────────────────────────────────────────────────────────────────────────────
|
||||
|
||||
#[test]
|
||||
fn metric_increments_pass_model_through_scrub_and_attach_session_id() {
|
||||
let stream = build(gates_off());
|
||||
emit_into(
|
||||
&stream,
|
||||
ExternalRecord {
|
||||
event: None,
|
||||
attrs: vec![(ExternalKey::SessionId, AttrValue::Str("sess-9".into()))],
|
||||
gated: vec![],
|
||||
metrics: vec![MetricIncrement::TokenUsage {
|
||||
token_type: "input",
|
||||
model: "sk-CANARYabcdefghij1234567890".into(),
|
||||
count: 7,
|
||||
}],
|
||||
},
|
||||
);
|
||||
let exported = stream.metrics.get_finished_metrics().unwrap();
|
||||
let blob = format!("{exported:?}");
|
||||
assert!(blob.contains("grok_code.token.usage"));
|
||||
assert!(
|
||||
blob.contains("sess-9"),
|
||||
"session.id missing from metric: {blob}"
|
||||
);
|
||||
assert!(
|
||||
!blob.contains("CANARY"),
|
||||
"model-id metric attribute must pass the secret scrub: {blob}"
|
||||
);
|
||||
}
|
||||
186
crates/codegen/xai-grok-telemetry/src/external/truncate.rs
vendored
Normal file
186
crates/codegen/xai-grok-telemetry/src/external/truncate.rs
vendored
Normal file
|
|
@ -0,0 +1,186 @@
|
|||
//! Truncation helpers for the external OTEL stream.
|
||||
//!
|
||||
//! Constants follow common customer-pipeline parity values: strings
|
||||
//! longer than 512 chars collapse to their first 128 chars plus a marker,
|
||||
//! tool-input JSON is capped at 4 KB / depth 2 / 20 items per collection, and
|
||||
//! gated prompt/content text is capped at 60 KB.
|
||||
|
||||
/// Values longer than this are truncated…
|
||||
pub const MAX_STRING_LEN: usize = 512;
|
||||
/// …to their first 128 chars + [`TRUNCATION_MARKER`].
|
||||
pub const TRUNCATED_PREFIX_LEN: usize = 128;
|
||||
/// Marker appended to truncated strings.
|
||||
pub const TRUNCATION_MARKER: &str = "…[truncated]";
|
||||
/// Total serialized-JSON budget for gated tool parameters.
|
||||
pub const MAX_TOOL_INPUT_JSON_BYTES: usize = 4 * 1024;
|
||||
/// Maximum JSON nesting depth preserved in gated tool parameters.
|
||||
pub const MAX_JSON_DEPTH: usize = 2;
|
||||
/// Maximum items preserved per JSON array/object in gated tool parameters.
|
||||
pub const MAX_COLLECTION_ITEMS: usize = 20;
|
||||
/// Cap for gated prompt/content text.
|
||||
pub const MAX_CONTENT_BYTES: usize = 60 * 1024;
|
||||
/// File-extension attribute cap (`"rs"`, `"tsx"`, …).
|
||||
pub const MAX_FILE_EXTENSION_LEN: usize = 10;
|
||||
|
||||
/// Truncate on a `char` boundary at or before `max_bytes`.
|
||||
fn floor_char_boundary(s: &str, max_bytes: usize) -> usize {
|
||||
if max_bytes >= s.len() {
|
||||
return s.len();
|
||||
}
|
||||
let mut idx = max_bytes;
|
||||
while idx > 0 && !s.is_char_boundary(idx) {
|
||||
idx -= 1;
|
||||
}
|
||||
idx
|
||||
}
|
||||
|
||||
/// Standard attribute-value truncation: strings whose `char` count exceeds
|
||||
/// [`MAX_STRING_LEN`] collapse to their first [`TRUNCATED_PREFIX_LEN`] chars
|
||||
/// plus [`TRUNCATION_MARKER`]. Returns `None` when unchanged.
|
||||
pub fn truncate_value(s: &str) -> Option<String> {
|
||||
// Counting chars (not bytes) keeps the limit stable for non-ASCII text;
|
||||
// chars_count > MAX_STRING_LEN implies the string is "long" regardless of
|
||||
// encoding width.
|
||||
if s.chars().count() <= MAX_STRING_LEN {
|
||||
return None;
|
||||
}
|
||||
let truncated: String = s.chars().take(TRUNCATED_PREFIX_LEN).collect();
|
||||
Some(format!("{truncated}{TRUNCATION_MARKER}"))
|
||||
}
|
||||
|
||||
/// Apply [`truncate_value`], returning an owned string either way.
|
||||
pub fn truncate_value_owned(s: String) -> String {
|
||||
truncate_value(&s).unwrap_or(s)
|
||||
}
|
||||
|
||||
/// Cap gated prompt/content text at [`MAX_CONTENT_BYTES`] (UTF-8-safe).
|
||||
pub fn truncate_content(s: &str) -> Option<String> {
|
||||
if s.len() <= MAX_CONTENT_BYTES {
|
||||
return None;
|
||||
}
|
||||
let idx = floor_char_boundary(s, MAX_CONTENT_BYTES);
|
||||
Some(format!("{}{TRUNCATION_MARKER}", &s[..idx]))
|
||||
}
|
||||
|
||||
/// Reduce a JSON value for the gated `tool_parameters` attribute: depth
|
||||
/// capped at [`MAX_JSON_DEPTH`], collections capped at
|
||||
/// [`MAX_COLLECTION_ITEMS`] entries, strings truncated per [`truncate_value`].
|
||||
/// The serialized result is finally clamped to [`MAX_TOOL_INPUT_JSON_BYTES`].
|
||||
pub fn reduce_tool_input(value: &serde_json::Value) -> String {
|
||||
let reduced = reduce_json(value, 0);
|
||||
let serialized = reduced.to_string();
|
||||
if serialized.len() <= MAX_TOOL_INPUT_JSON_BYTES {
|
||||
return serialized;
|
||||
}
|
||||
// Over budget even after structural reduction: clamp the serialized text.
|
||||
// The result may not be valid JSON, but it is bounded and marked.
|
||||
let idx = floor_char_boundary(&serialized, MAX_TOOL_INPUT_JSON_BYTES);
|
||||
format!("{}{TRUNCATION_MARKER}", &serialized[..idx])
|
||||
}
|
||||
|
||||
fn reduce_json(value: &serde_json::Value, depth: usize) -> serde_json::Value {
|
||||
use serde_json::Value;
|
||||
match value {
|
||||
Value::String(s) => Value::String(truncate_value(s).unwrap_or_else(|| s.clone())),
|
||||
Value::Array(items) => {
|
||||
if depth >= MAX_JSON_DEPTH {
|
||||
return Value::String(format!("[array:{}]", items.len()));
|
||||
}
|
||||
items
|
||||
.iter()
|
||||
.take(MAX_COLLECTION_ITEMS)
|
||||
.map(|v| reduce_json(v, depth + 1))
|
||||
.collect()
|
||||
}
|
||||
Value::Object(map) => {
|
||||
if depth >= MAX_JSON_DEPTH {
|
||||
return Value::String(format!("{{object:{}}}", map.len()));
|
||||
}
|
||||
map.iter()
|
||||
.take(MAX_COLLECTION_ITEMS)
|
||||
.map(|(k, v)| (k.clone(), reduce_json(v, depth + 1)))
|
||||
.collect()
|
||||
}
|
||||
other => other.clone(),
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
#[test]
|
||||
fn short_strings_pass_unchanged() {
|
||||
assert_eq!(truncate_value("hello"), None);
|
||||
let exactly_max: String = "a".repeat(MAX_STRING_LEN);
|
||||
assert_eq!(truncate_value(&exactly_max), None);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn long_strings_collapse_to_prefix_plus_marker() {
|
||||
let long: String = "x".repeat(MAX_STRING_LEN + 1);
|
||||
let out = truncate_value(&long).expect("must truncate");
|
||||
assert!(out.starts_with(&"x".repeat(TRUNCATED_PREFIX_LEN)));
|
||||
assert!(out.ends_with(TRUNCATION_MARKER));
|
||||
assert_eq!(
|
||||
out.chars().count(),
|
||||
TRUNCATED_PREFIX_LEN + TRUNCATION_MARKER.chars().count()
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn truncation_is_utf8_boundary_safe() {
|
||||
// Multi-byte chars around both limits must not split a char.
|
||||
let long: String = "é".repeat(MAX_STRING_LEN + 5);
|
||||
let out = truncate_value(&long).expect("must truncate");
|
||||
assert!(out.starts_with(&"é".repeat(TRUNCATED_PREFIX_LEN)));
|
||||
|
||||
let content: String = "🎉".repeat(MAX_CONTENT_BYTES / 4 + 10);
|
||||
let out = truncate_content(&content).expect("must truncate");
|
||||
assert!(out.len() <= MAX_CONTENT_BYTES + TRUNCATION_MARKER.len());
|
||||
// Round-trip as a str: would panic at construction if we split a char.
|
||||
assert!(out.ends_with(TRUNCATION_MARKER));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn content_under_cap_passes() {
|
||||
assert_eq!(truncate_content("short prompt"), None);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn tool_input_depth_capped() {
|
||||
let v = serde_json::json!({"a": {"b": {"c": {"d": 1}}}});
|
||||
let out = reduce_tool_input(&v);
|
||||
let parsed: serde_json::Value = serde_json::from_str(&out).unwrap();
|
||||
// Depth 0 = root object, depth 1 = a's object; b's value is at depth 2 → collapsed.
|
||||
assert_eq!(parsed["a"]["b"], serde_json::json!("{object:1}"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn tool_input_collection_items_capped() {
|
||||
let items: Vec<serde_json::Value> = (0..50).map(|i| serde_json::json!(i)).collect();
|
||||
let v = serde_json::Value::Array(items);
|
||||
let out = reduce_tool_input(&v);
|
||||
let parsed: serde_json::Value = serde_json::from_str(&out).unwrap();
|
||||
assert_eq!(parsed.as_array().unwrap().len(), MAX_COLLECTION_ITEMS);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn tool_input_total_budget_enforced() {
|
||||
let big: String = "y".repeat(300);
|
||||
let map: serde_json::Map<String, serde_json::Value> = (0..MAX_COLLECTION_ITEMS)
|
||||
.map(|i| (format!("key_{i:02}"), serde_json::json!(big.clone())))
|
||||
.collect();
|
||||
let v = serde_json::json!({"a": map.clone(), "b": map});
|
||||
let out = reduce_tool_input(&v);
|
||||
assert!(out.len() <= MAX_TOOL_INPUT_JSON_BYTES + TRUNCATION_MARKER.len());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn tool_input_strings_truncated_inside_json() {
|
||||
let v = serde_json::json!({"text": "z".repeat(MAX_STRING_LEN + 1)});
|
||||
let out = reduce_tool_input(&v);
|
||||
assert!(out.contains(TRUNCATION_MARKER));
|
||||
assert!(!out.contains(&"z".repeat(MAX_STRING_LEN + 1)));
|
||||
}
|
||||
}
|
||||
124
crates/codegen/xai-grok-telemetry/src/hooks_log.rs
Normal file
124
crates/codegen/xai-grok-telemetry/src/hooks_log.rs
Normal file
|
|
@ -0,0 +1,124 @@
|
|||
//! Hooks and plugins tracing target and optional file-based logging layer.
|
||||
//!
|
||||
//! A dedicated tracing target for hooks and plugins subsystems with an optional
|
||||
//! file logger that writes to `~/.grok/logs/hooks.log`.
|
||||
//!
|
||||
//! ## When to use
|
||||
//!
|
||||
//! Use regular `tracing::info!` / `tracing::debug!` / `tracing::warn!` with
|
||||
//! targets `xai_grok_hooks` or `xai_grok_agent::plugins` at key lifecycle
|
||||
//! points — discovery, dispatch, execution, errors.
|
||||
//!
|
||||
//! ## Enabling
|
||||
//!
|
||||
//! ```bash
|
||||
//! GROK_HOOKS_LOG=1 grok # enable, write to ~/.grok/logs/hooks.log
|
||||
//! GROK_HOOKS_LOG=/tmp/h.log grok # write to custom path
|
||||
//! GROK_HOOKS_LOG=0 grok # explicitly disable
|
||||
//! tail -f ~/.grok/logs/hooks.log # watch in another terminal
|
||||
//! ```
|
||||
|
||||
use std::fmt;
|
||||
use std::path::PathBuf;
|
||||
use std::sync::Mutex;
|
||||
use std::time::Instant;
|
||||
|
||||
use tracing::Subscriber;
|
||||
use tracing_subscriber::fmt::format::Writer;
|
||||
use tracing_subscriber::fmt::time::FormatTime;
|
||||
use tracing_subscriber::fmt::writer::BoxMakeWriter;
|
||||
use tracing_subscriber::layer::Layer;
|
||||
use tracing_subscriber::registry::LookupSpan;
|
||||
|
||||
use xai_grok_config::grok_home;
|
||||
|
||||
const ENV_HOOKS_LOG: &str = "GROK_HOOKS_LOG";
|
||||
|
||||
static LOG_GUARD: std::sync::OnceLock<Mutex<Option<tracing_appender::non_blocking::WorkerGuard>>> =
|
||||
std::sync::OnceLock::new();
|
||||
|
||||
#[derive(Clone)]
|
||||
struct UptimeTimer {
|
||||
epoch: Instant,
|
||||
}
|
||||
|
||||
impl UptimeTimer {
|
||||
fn new() -> Self {
|
||||
Self {
|
||||
epoch: Instant::now(),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl FormatTime for UptimeTimer {
|
||||
fn format_time(&self, w: &mut Writer<'_>) -> fmt::Result {
|
||||
let elapsed = self.epoch.elapsed();
|
||||
write!(w, "+{}.{:03}s", elapsed.as_secs(), elapsed.subsec_millis())
|
||||
}
|
||||
}
|
||||
|
||||
/// Build the hooks/plugins log layer.
|
||||
///
|
||||
/// Writes to `~/.grok/logs/hooks.log` (or custom path via `GROK_HOOKS_LOG`).
|
||||
/// Filters to hooks (`xai_grok_hooks`) and plugins (`xai_grok_agent::plugins`) targets.
|
||||
/// Set `GROK_HOOKS_LOG=0` to disable, `GROK_HOOKS_LOG=/path` to redirect.
|
||||
pub fn layer<S>() -> Option<impl Layer<S>>
|
||||
where
|
||||
S: Subscriber + for<'span> LookupSpan<'span>,
|
||||
{
|
||||
let path = resolve_log_path()?;
|
||||
|
||||
if let Some(parent) = path.parent() {
|
||||
let _ = std::fs::create_dir_all(parent);
|
||||
}
|
||||
|
||||
let file = match std::fs::OpenOptions::new()
|
||||
.create(true)
|
||||
.append(true)
|
||||
.open(&path)
|
||||
{
|
||||
Ok(f) => f,
|
||||
Err(e) => {
|
||||
tracing::warn!("[hooks-log] Failed to open {:?}: {}", path, e);
|
||||
return None;
|
||||
}
|
||||
};
|
||||
|
||||
let (non_blocking, guard) = tracing_appender::non_blocking(file);
|
||||
let guard_slot = LOG_GUARD.get_or_init(|| Mutex::new(None));
|
||||
if let Ok(mut slot) = guard_slot.lock() {
|
||||
*slot = Some(guard);
|
||||
}
|
||||
|
||||
// Filter for both hooks and plugins targets at debug level
|
||||
let filter = tracing_subscriber::filter::EnvFilter::new(
|
||||
"xai_grok_hooks=debug,xai_grok_agent::plugins=debug",
|
||||
);
|
||||
let fmt_layer = tracing_subscriber::fmt::layer()
|
||||
.with_target(true)
|
||||
.with_ansi(false)
|
||||
.with_thread_ids(true)
|
||||
.with_timer(UptimeTimer::new())
|
||||
.with_writer(BoxMakeWriter::new(non_blocking))
|
||||
.with_filter(filter);
|
||||
|
||||
tracing::info!(
|
||||
"[hooks-log] Hooks/plugins logging enabled: {}",
|
||||
path.display()
|
||||
);
|
||||
Some(fmt_layer)
|
||||
}
|
||||
|
||||
fn resolve_log_path() -> Option<PathBuf> {
|
||||
let default_path = || grok_home().join("logs").join("hooks.log");
|
||||
let raw = match std::env::var(ENV_HOOKS_LOG) {
|
||||
Ok(val) => val,
|
||||
Err(_) => return None, // opt-in only
|
||||
};
|
||||
let raw = raw.trim();
|
||||
match raw {
|
||||
"" | "0" | "false" | "off" | "no" => None,
|
||||
"1" | "true" | "on" | "yes" => Some(default_path()),
|
||||
other => Some(PathBuf::from(other)),
|
||||
}
|
||||
}
|
||||
21
crates/codegen/xai-grok-telemetry/src/http.rs
Normal file
21
crates/codegen/xai-grok-telemetry/src/http.rs
Normal file
|
|
@ -0,0 +1,21 @@
|
|||
//! Origin/client identification used by the telemetry engine.
|
||||
//!
|
||||
//! [`OriginClientInfo`] is owned by `xai-grok-sampler` (so `SamplerConfig`
|
||||
//! can use it without depending on shell). Re-exported here so the telemetry
|
||||
//! engine can label events without depending on shell or sampler internals
|
||||
//! beyond the type itself.
|
||||
|
||||
pub use xai_grok_sampler::OriginClientInfo;
|
||||
|
||||
/// Construct an [`OriginClientInfo`] from `GROK_CLIENT_NAME` /
|
||||
/// `GROK_CLIENT_VERSION` env vars. Returns `None` when `GROK_CLIENT_NAME`
|
||||
/// is unset. Free function (not an inherent method) because the type lives
|
||||
/// in another crate.
|
||||
pub fn origin_client_info_from_env() -> Option<OriginClientInfo> {
|
||||
std::env::var("GROK_CLIENT_NAME")
|
||||
.ok()
|
||||
.map(|product| OriginClientInfo {
|
||||
product,
|
||||
version: std::env::var("GROK_CLIENT_VERSION").ok(),
|
||||
})
|
||||
}
|
||||
86
crates/codegen/xai-grok-telemetry/src/id.rs
Normal file
86
crates/codegen/xai-grok-telemetry/src/id.rs
Normal file
|
|
@ -0,0 +1,86 @@
|
|||
//! Stable agent identifier.
|
||||
//!
|
||||
//! Extracted from `xai-grok-shell::agent::unique_identifier` so the
|
||||
//! telemetry engine can stamp events without depending on shell internals.
|
||||
//! `$GROK_HOME` is resolved through `xai-grok-config::grok_home`.
|
||||
|
||||
use std::sync::OnceLock;
|
||||
|
||||
/// Cached agent ID - stored in memory after first load.
|
||||
static AGENT_ID: OnceLock<String> = OnceLock::new();
|
||||
/// Cached agent instance ID - per-process lifetime.
|
||||
static AGENT_INSTANCE_ID: OnceLock<String> = OnceLock::new();
|
||||
|
||||
/// Returns the agent ID, using a file-based cache to avoid expensive system calls.
|
||||
///
|
||||
/// On macOS, `mid::get()` calls `system_profiler` which takes ~1-3 seconds.
|
||||
/// This function caches the result in `$GROK_HOME/agent_id` so subsequent calls
|
||||
/// (even across process restarts) are instant file reads.
|
||||
///
|
||||
/// The in-memory `OnceLock` ensures we only read the file once per process.
|
||||
pub fn agent_id() -> String {
|
||||
AGENT_ID.get_or_init(load_or_compute_agent_id).clone()
|
||||
}
|
||||
|
||||
/// Returns a per-process agent instance ID.
|
||||
/// This is stable across WebSocket reconnects within the same process,
|
||||
/// but changes on process restart.
|
||||
pub fn agent_instance_id() -> String {
|
||||
AGENT_INSTANCE_ID
|
||||
.get_or_init(|| uuid::Uuid::new_v4().to_string())
|
||||
.clone()
|
||||
}
|
||||
|
||||
fn load_or_compute_agent_id() -> String {
|
||||
let cache_path = xai_grok_config::grok_home().join("agent_id");
|
||||
|
||||
// Try to read from cache file first (fast path)
|
||||
if let Ok(cached) = std::fs::read_to_string(&cache_path) {
|
||||
let cached = cached.trim();
|
||||
if !cached.is_empty() {
|
||||
return cached.to_string();
|
||||
}
|
||||
}
|
||||
|
||||
// Compute a unique machine hash:
|
||||
// - macOS: mid uses unique hardware IDs (serial, UUID, SEID).
|
||||
// - Linux: /etc/machine-id is shared across containers from the same base
|
||||
// image, so include $HOSTNAME (container/host name) for uniqueness.
|
||||
// - Fallback: random UUIDv4 if mid or hostname are unavailable.
|
||||
let machine_hash = if cfg!(target_os = "linux") {
|
||||
match std::env::var("HOSTNAME") {
|
||||
Ok(hostname) if !hostname.is_empty() => {
|
||||
let key = format!("agent_id:{hostname}");
|
||||
mid::get(&key).unwrap_or_else(|_| uuid::Uuid::new_v4().to_string())
|
||||
}
|
||||
_ => uuid::Uuid::new_v4().to_string(),
|
||||
}
|
||||
} else {
|
||||
mid::get("agent_id").unwrap_or_else(|_| uuid::Uuid::new_v4().to_string())
|
||||
};
|
||||
let id = uuid::Uuid::new_v5(&uuid::Uuid::NAMESPACE_OID, machine_hash.as_bytes()).to_string();
|
||||
|
||||
// Save to cache file (best effort, ignore errors)
|
||||
let _ = std::fs::write(&cache_path, &id);
|
||||
|
||||
id
|
||||
}
|
||||
|
||||
/// Returns true when workspace marker env vars (`XAI_ROOT` and `XAI_USER`) are set.
|
||||
///
|
||||
/// Used as a coarse local gate for features that require a full workspace
|
||||
/// checkout. External installs typically leave both unset.
|
||||
pub fn has_workspace_env_markers() -> bool {
|
||||
std::env::var("XAI_ROOT").is_ok() && std::env::var("XAI_USER").is_ok()
|
||||
}
|
||||
|
||||
/// Opt-in special-user gate for telemetry.
|
||||
///
|
||||
/// Enabled only when `GROK_TELEMETRY_SPECIAL_USER=1` (or `true`). There is no
|
||||
/// hardcoded username allowlist.
|
||||
pub fn is_special_user() -> bool {
|
||||
matches!(
|
||||
std::env::var("GROK_TELEMETRY_SPECIAL_USER").as_deref(),
|
||||
Ok("1") | Ok("true") | Ok("TRUE")
|
||||
)
|
||||
}
|
||||
634
crates/codegen/xai-grok-telemetry/src/instrumentation.rs
Normal file
634
crates/codegen/xai-grok-telemetry/src/instrumentation.rs
Normal file
|
|
@ -0,0 +1,634 @@
|
|||
use std::io::{self, BufRead};
|
||||
use std::marker::PhantomData;
|
||||
use std::path::PathBuf;
|
||||
use std::sync::{Mutex, OnceLock};
|
||||
use std::time::Instant;
|
||||
|
||||
use anyhow::{Result, anyhow};
|
||||
use chrono::{DateTime, FixedOffset};
|
||||
use serde_json::Value;
|
||||
use tracing::Subscriber;
|
||||
use tracing_chrome::{ChromeLayerBuilder, FlushGuard, TraceStyle};
|
||||
use tracing_subscriber::fmt::writer::BoxMakeWriter;
|
||||
use tracing_subscriber::layer::{Context, Layer};
|
||||
use tracing_subscriber::registry::LookupSpan;
|
||||
|
||||
use xai_grok_config::grok_home;
|
||||
|
||||
const ENV_ENABLED: &str = "GROK_INSTRUMENTATION";
|
||||
const ENV_LOG_PATH: &str = "GROK_INSTRUMENTATION_LOG";
|
||||
const DEFAULT_LOG_DIR: &str = "logs";
|
||||
const DEFAULT_LOG_FILE: &str = "instrumentation.log";
|
||||
const DEFAULT_TRACE_FILE: &str = "instrumentation.trace.json";
|
||||
|
||||
pub const TARGET: &str = "xai_grok_instrumentation";
|
||||
|
||||
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
|
||||
pub enum InstrumentationMode {
|
||||
Disabled,
|
||||
Log,
|
||||
Chrome,
|
||||
Server,
|
||||
}
|
||||
|
||||
static INSTRUMENTATION_MODE: OnceLock<InstrumentationMode> = OnceLock::new();
|
||||
static LOG_GUARD: OnceLock<Mutex<Option<tracing_appender::non_blocking::WorkerGuard>>> =
|
||||
OnceLock::new();
|
||||
static CHROME_GUARD: OnceLock<Mutex<Option<FlushGuard>>> = OnceLock::new();
|
||||
|
||||
fn mode() -> InstrumentationMode {
|
||||
*INSTRUMENTATION_MODE.get_or_init(|| {
|
||||
let env_mode = match std::env::var(ENV_ENABLED) {
|
||||
Ok(v) => match v.trim().to_ascii_lowercase().as_str() {
|
||||
"1" | "true" | "on" | "enabled" | "log" | "json" | "jsonl" => {
|
||||
Some(InstrumentationMode::Log)
|
||||
}
|
||||
"chrome" | "trace" | "trace.json" => Some(InstrumentationMode::Chrome),
|
||||
"server" => Some(InstrumentationMode::Server),
|
||||
"" | "0" | "false" | "off" | "disabled" | "none" => {
|
||||
Some(InstrumentationMode::Disabled)
|
||||
}
|
||||
_ => Some(InstrumentationMode::Log),
|
||||
},
|
||||
Err(_) => None,
|
||||
};
|
||||
|
||||
if let Some(mode) = env_mode {
|
||||
return mode;
|
||||
}
|
||||
|
||||
// Send instrumentation to the configured OpenTelemetry endpoint by default
|
||||
InstrumentationMode::Server
|
||||
})
|
||||
}
|
||||
|
||||
pub fn current_mode() -> InstrumentationMode {
|
||||
mode()
|
||||
}
|
||||
|
||||
fn default_log_path() -> PathBuf {
|
||||
grok_home().join(DEFAULT_LOG_DIR).join(DEFAULT_LOG_FILE)
|
||||
}
|
||||
|
||||
fn log_path_from_env() -> Option<PathBuf> {
|
||||
std::env::var(ENV_LOG_PATH)
|
||||
.ok()
|
||||
.map(|path| path.trim().to_string())
|
||||
.filter(|path| !path.is_empty())
|
||||
.map(PathBuf::from)
|
||||
}
|
||||
|
||||
fn default_output_path(mode: InstrumentationMode) -> PathBuf {
|
||||
let root = grok_home().join(DEFAULT_LOG_DIR);
|
||||
match mode {
|
||||
InstrumentationMode::Chrome => root.join(DEFAULT_TRACE_FILE),
|
||||
// Server uses OTLP export, not file output
|
||||
InstrumentationMode::Log | InstrumentationMode::Disabled | InstrumentationMode::Server => {
|
||||
root.join(DEFAULT_LOG_FILE)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// A wrapper layer that filters events by target name.
|
||||
///
|
||||
/// This is used instead of `.with_filter()` because `Filtered<L, F, S>` layers
|
||||
/// require `FilterId` registration with the subscriber. When boxed as
|
||||
/// `Box<dyn Layer<S>>`, the type information needed for registration is lost,
|
||||
/// causing a panic: "a Filtered layer was used, but it had no FilterId".
|
||||
///
|
||||
/// This wrapper avoids that issue by implementing filtering in the `enabled()`
|
||||
/// method directly, without using the per-layer filter mechanism.
|
||||
pub struct TargetFilterLayer<L, S> {
|
||||
inner: L,
|
||||
target: &'static str,
|
||||
_subscriber: PhantomData<fn(S)>,
|
||||
}
|
||||
|
||||
impl<L, S> TargetFilterLayer<L, S> {
|
||||
pub fn new(inner: L, target: &'static str) -> Self {
|
||||
Self {
|
||||
inner,
|
||||
target,
|
||||
_subscriber: PhantomData,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl<L, S> Layer<S> for TargetFilterLayer<L, S>
|
||||
where
|
||||
L: Layer<S>,
|
||||
S: Subscriber + for<'span> LookupSpan<'span>,
|
||||
{
|
||||
fn enabled(&self, metadata: &tracing::Metadata<'_>, ctx: Context<'_, S>) -> bool {
|
||||
metadata.target() == self.target && self.inner.enabled(metadata, ctx)
|
||||
}
|
||||
|
||||
fn on_new_span(
|
||||
&self,
|
||||
attrs: &tracing::span::Attributes<'_>,
|
||||
id: &tracing::span::Id,
|
||||
ctx: Context<'_, S>,
|
||||
) {
|
||||
if attrs.metadata().target() == self.target {
|
||||
self.inner.on_new_span(attrs, id, ctx);
|
||||
}
|
||||
}
|
||||
|
||||
fn on_record(
|
||||
&self,
|
||||
span: &tracing::span::Id,
|
||||
values: &tracing::span::Record<'_>,
|
||||
ctx: Context<'_, S>,
|
||||
) {
|
||||
self.inner.on_record(span, values, ctx);
|
||||
}
|
||||
|
||||
fn on_follows_from(
|
||||
&self,
|
||||
span: &tracing::span::Id,
|
||||
follows: &tracing::span::Id,
|
||||
ctx: Context<'_, S>,
|
||||
) {
|
||||
self.inner.on_follows_from(span, follows, ctx);
|
||||
}
|
||||
|
||||
fn on_event(&self, event: &tracing::Event<'_>, ctx: Context<'_, S>) {
|
||||
if event.metadata().target() == self.target {
|
||||
self.inner.on_event(event, ctx);
|
||||
}
|
||||
}
|
||||
|
||||
fn on_enter(&self, id: &tracing::span::Id, ctx: Context<'_, S>) {
|
||||
self.inner.on_enter(id, ctx);
|
||||
}
|
||||
|
||||
fn on_exit(&self, id: &tracing::span::Id, ctx: Context<'_, S>) {
|
||||
self.inner.on_exit(id, ctx);
|
||||
}
|
||||
|
||||
fn on_close(&self, id: tracing::span::Id, ctx: Context<'_, S>) {
|
||||
self.inner.on_close(id, ctx);
|
||||
}
|
||||
}
|
||||
|
||||
/// A no-op layer that does nothing.
|
||||
/// Used when instrumentation is disabled to avoid any overhead.
|
||||
pub struct NoOpLayer<S>(PhantomData<fn(S)>);
|
||||
|
||||
impl<S> Default for NoOpLayer<S> {
|
||||
fn default() -> Self {
|
||||
Self(PhantomData)
|
||||
}
|
||||
}
|
||||
|
||||
impl<S> NoOpLayer<S> {
|
||||
pub fn new() -> Self {
|
||||
Self::default()
|
||||
}
|
||||
}
|
||||
|
||||
impl<S: Subscriber> Layer<S> for NoOpLayer<S> {
|
||||
// All methods use default implementations which do nothing
|
||||
}
|
||||
|
||||
fn resolve_output_path(mode: InstrumentationMode) -> Option<PathBuf> {
|
||||
if mode == InstrumentationMode::Disabled {
|
||||
return None;
|
||||
}
|
||||
|
||||
log_path_from_env().or_else(|| Some(default_output_path(mode)))
|
||||
}
|
||||
|
||||
fn resolve_log_path() -> Option<PathBuf> {
|
||||
if mode() != InstrumentationMode::Log {
|
||||
return None;
|
||||
}
|
||||
resolve_output_path(InstrumentationMode::Log)
|
||||
}
|
||||
|
||||
fn build_writer(path: Option<PathBuf>) -> BoxMakeWriter {
|
||||
let Some(path) = path else {
|
||||
return BoxMakeWriter::new(std::io::sink);
|
||||
};
|
||||
|
||||
if let Some(parent) = path.parent()
|
||||
&& let Err(err) = std::fs::create_dir_all(parent)
|
||||
{
|
||||
eprintln!(
|
||||
"Failed to create instrumentation log directory {:?}: {}",
|
||||
parent, err
|
||||
);
|
||||
return BoxMakeWriter::new(std::io::sink);
|
||||
}
|
||||
|
||||
let file = match std::fs::OpenOptions::new()
|
||||
.create(true)
|
||||
.append(true)
|
||||
.open(&path)
|
||||
{
|
||||
Ok(file) => file,
|
||||
Err(err) => {
|
||||
eprintln!(
|
||||
"Failed to open instrumentation log file {:?}: {}",
|
||||
path, err
|
||||
);
|
||||
return BoxMakeWriter::new(std::io::sink);
|
||||
}
|
||||
};
|
||||
|
||||
let (non_blocking, guard) = tracing_appender::non_blocking(file);
|
||||
let guard_slot = LOG_GUARD.get_or_init(|| Mutex::new(None));
|
||||
if let Ok(mut slot) = guard_slot.lock() {
|
||||
*slot = Some(guard);
|
||||
}
|
||||
BoxMakeWriter::new(non_blocking)
|
||||
}
|
||||
|
||||
fn build_log_layer<S>(mode: InstrumentationMode) -> Box<dyn Layer<S> + Send + Sync>
|
||||
where
|
||||
S: Subscriber + for<'span> LookupSpan<'span> + Send + Sync + 'static,
|
||||
{
|
||||
// When disabled, return a true no-op layer that does nothing.
|
||||
// This avoids any overhead and potential issues with complex layer types.
|
||||
if mode == InstrumentationMode::Disabled {
|
||||
return Box::new(NoOpLayer::new());
|
||||
}
|
||||
|
||||
let writer = build_writer(resolve_log_path());
|
||||
|
||||
// Use TargetFilterLayer instead of .with_filter() to avoid the FilterId
|
||||
// registration issue when the layer is boxed as Box<dyn Layer<S>>.
|
||||
let fmt_layer = tracing_subscriber::fmt::layer()
|
||||
.json()
|
||||
.with_current_span(false) // `spans` array already carries the full ancestor list
|
||||
.with_ansi(false)
|
||||
.with_timer(tracing_subscriber::fmt::time::UtcTime::rfc_3339())
|
||||
.with_thread_ids(true)
|
||||
.with_thread_names(true)
|
||||
.with_target(true)
|
||||
.with_writer(writer);
|
||||
|
||||
Box::new(TargetFilterLayer::new(fmt_layer, TARGET))
|
||||
}
|
||||
|
||||
fn build_chrome_layer<S>() -> Box<dyn Layer<S> + Send + Sync>
|
||||
where
|
||||
S: Subscriber + for<'span> LookupSpan<'span> + Send + Sync + 'static,
|
||||
{
|
||||
let Some(path) = resolve_output_path(InstrumentationMode::Chrome) else {
|
||||
return build_log_layer(InstrumentationMode::Disabled);
|
||||
};
|
||||
|
||||
if let Some(parent) = path.parent()
|
||||
&& let Err(err) = std::fs::create_dir_all(parent)
|
||||
{
|
||||
eprintln!(
|
||||
"Failed to create chrome trace directory {:?}: {}",
|
||||
parent, err
|
||||
);
|
||||
return build_log_layer(InstrumentationMode::Disabled);
|
||||
}
|
||||
|
||||
let file = match std::fs::OpenOptions::new()
|
||||
.create(true)
|
||||
.truncate(true)
|
||||
.write(true)
|
||||
.open(&path)
|
||||
{
|
||||
Ok(file) => file,
|
||||
Err(err) => {
|
||||
eprintln!("Failed to open chrome trace file {:?}: {}", path, err);
|
||||
return build_log_layer(InstrumentationMode::Disabled);
|
||||
}
|
||||
};
|
||||
|
||||
let (layer, guard) = ChromeLayerBuilder::<S>::new()
|
||||
.writer(file)
|
||||
.include_args(true)
|
||||
.trace_style(TraceStyle::Async)
|
||||
.build();
|
||||
|
||||
let guard_slot = CHROME_GUARD.get_or_init(|| Mutex::new(None));
|
||||
if let Ok(mut slot) = guard_slot.lock() {
|
||||
*slot = Some(guard);
|
||||
}
|
||||
|
||||
// Use TargetFilterLayer instead of .with_filter() to avoid the FilterId
|
||||
// registration issue when the layer is boxed as Box<dyn Layer<S>>.
|
||||
Box::new(TargetFilterLayer::new(layer, TARGET))
|
||||
}
|
||||
|
||||
pub fn layer<S>() -> Box<dyn Layer<S> + Send + Sync>
|
||||
where
|
||||
S: Subscriber + for<'span> LookupSpan<'span> + Send + Sync + 'static,
|
||||
{
|
||||
let mode = mode();
|
||||
match mode {
|
||||
InstrumentationMode::Chrome => build_chrome_layer(),
|
||||
InstrumentationMode::Log => build_log_layer(mode),
|
||||
// Server uses the OTEL layer in tracing.rs, not this instrumentation layer
|
||||
InstrumentationMode::Disabled | InstrumentationMode::Server => {
|
||||
build_log_layer(InstrumentationMode::Disabled)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Install a global panic hook that emits a structured tracing event before
|
||||
/// invoking the default hook. Call this once, early in `main`, after the
|
||||
/// tracing subscriber has been installed.
|
||||
pub fn install_panic_hook() {
|
||||
let default_hook = std::panic::take_hook();
|
||||
std::panic::set_hook(Box::new(move |info| {
|
||||
let message = if let Some(s) = info.payload().downcast_ref::<&str>() {
|
||||
s.to_string()
|
||||
} else if let Some(s) = info.payload().downcast_ref::<String>() {
|
||||
s.clone()
|
||||
} else {
|
||||
"unknown panic".to_string()
|
||||
};
|
||||
let location = info
|
||||
.location()
|
||||
.map(|l| format!("{}:{}:{}", l.file(), l.line(), l.column()));
|
||||
// `location` is the panic's source `file:line:col` (no user content);
|
||||
// path-scrubbed by the redact layer. Gives the panic counter a place
|
||||
// to point without exporting the message/stack.
|
||||
let err_span = tracing::info_span!(
|
||||
"internal_error",
|
||||
error_type = "panic",
|
||||
location = tracing::field::Empty,
|
||||
);
|
||||
if let Some(loc) = location.as_deref() {
|
||||
err_span.record("location", loc);
|
||||
}
|
||||
err_span.in_scope(|| {});
|
||||
// External OTEL stream: error class only — no message, no location
|
||||
// (RQ5). Synchronous queue hand-off; no-op unless the stream is
|
||||
// active. The internal pipelines keep the richer span/event above.
|
||||
crate::external::emit(&crate::events::InternalError {
|
||||
error_type: "panic".to_owned(),
|
||||
});
|
||||
tracing::error!(
|
||||
error_type = "panic",
|
||||
panic.message = %message,
|
||||
panic.location = ?location,
|
||||
"Process panicked"
|
||||
);
|
||||
default_hook(info);
|
||||
}));
|
||||
}
|
||||
|
||||
fn resolve_input_path(input: Option<PathBuf>) -> Result<PathBuf> {
|
||||
if let Some(path) = input {
|
||||
return Ok(path);
|
||||
}
|
||||
if let Some(path) = log_path_from_env() {
|
||||
return Ok(path);
|
||||
}
|
||||
Ok(default_log_path())
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, Default)]
|
||||
pub struct ChromeTraceOptions {
|
||||
pub input: Option<PathBuf>,
|
||||
pub output: Option<PathBuf>,
|
||||
}
|
||||
|
||||
pub fn generate_chrome_trace(options: ChromeTraceOptions) -> Result<PathBuf> {
|
||||
let input = resolve_input_path(options.input)?;
|
||||
let output = options
|
||||
.output
|
||||
.unwrap_or_else(|| input.with_extension("trace.json"));
|
||||
|
||||
let file = std::fs::File::open(&input)
|
||||
.map_err(|err| anyhow!("failed to open instrumentation log {:?}: {}", input, err))?;
|
||||
let reader = io::BufReader::new(file);
|
||||
|
||||
let mut events: Vec<Value> = Vec::new();
|
||||
let mut seen = 0usize;
|
||||
|
||||
for line in reader.lines() {
|
||||
let line = match line {
|
||||
Ok(line) => line,
|
||||
Err(_) => continue,
|
||||
};
|
||||
let value: Value = match serde_json::from_str(&line) {
|
||||
Ok(value) => value,
|
||||
Err(_) => continue,
|
||||
};
|
||||
|
||||
let target = value.get("target").and_then(Value::as_str);
|
||||
if target != Some(TARGET) {
|
||||
continue;
|
||||
}
|
||||
|
||||
let fields = match value.get("fields").and_then(Value::as_object) {
|
||||
Some(fields) => fields,
|
||||
None => continue,
|
||||
};
|
||||
|
||||
let event = fields.get("event").and_then(Value::as_str);
|
||||
if event != Some("timing") {
|
||||
continue;
|
||||
}
|
||||
|
||||
let name = match fields.get("name").and_then(Value::as_str) {
|
||||
Some(name) => name,
|
||||
None => continue,
|
||||
};
|
||||
|
||||
// Support both elapsed_us (new) and elapsed_ms (legacy) formats
|
||||
let dur_us = if let Some(us) = fields.get("elapsed_us").and_then(|v| v.as_u64()) {
|
||||
us
|
||||
} else if let Some(ms) = fields.get("elapsed_ms").and_then(|v| v.as_u64()) {
|
||||
ms.saturating_mul(1_000)
|
||||
} else {
|
||||
continue;
|
||||
};
|
||||
if dur_us == 0 {
|
||||
continue;
|
||||
}
|
||||
|
||||
let timestamp = value.get("timestamp").and_then(Value::as_str);
|
||||
let end_us = match timestamp.and_then(parse_timestamp_us) {
|
||||
Some(ts) => ts,
|
||||
None => continue,
|
||||
};
|
||||
let start_us = end_us.saturating_sub(dur_us as i64);
|
||||
|
||||
let mut args = serde_json::Map::new();
|
||||
args.insert("elapsed_us".to_string(), Value::Number(dur_us.into()));
|
||||
if let Some(extra) = fields.get("fields") {
|
||||
args.insert("fields".to_string(), extra.clone());
|
||||
}
|
||||
|
||||
let thread_name = value
|
||||
.get("thread_name")
|
||||
.or_else(|| value.get("threadName"))
|
||||
.and_then(Value::as_str);
|
||||
if let Some(name) = thread_name {
|
||||
args.insert("thread_name".to_string(), Value::String(name.to_string()));
|
||||
}
|
||||
|
||||
let thread_id = value
|
||||
.get("thread_id")
|
||||
.or_else(|| value.get("threadId"))
|
||||
.and_then(parse_thread_id)
|
||||
.unwrap_or(0);
|
||||
|
||||
let trace_event = serde_json::json!({
|
||||
"name": name,
|
||||
"cat": "instrumentation",
|
||||
"ph": "X",
|
||||
"ts": start_us,
|
||||
"dur": dur_us,
|
||||
"pid": 1,
|
||||
"tid": thread_id,
|
||||
"args": Value::Object(args),
|
||||
});
|
||||
|
||||
events.push(trace_event);
|
||||
seen += 1;
|
||||
}
|
||||
|
||||
if seen == 0 {
|
||||
return Err(anyhow!("no timing events found in {:?}", input));
|
||||
}
|
||||
|
||||
let trace = serde_json::json!({
|
||||
"displayTimeUnit": "ms",
|
||||
"traceEvents": events,
|
||||
});
|
||||
|
||||
let mut output_file = std::fs::File::create(&output)
|
||||
.map_err(|err| anyhow!("failed to create chrome trace {:?}: {}", output, err))?;
|
||||
serde_json::to_writer_pretty(&mut output_file, &trace)
|
||||
.map_err(|err| anyhow!("failed to write chrome trace: {}", err))?;
|
||||
|
||||
Ok(output)
|
||||
}
|
||||
|
||||
pub fn finalize() -> Result<()> {
|
||||
let mode = mode();
|
||||
if mode == InstrumentationMode::Disabled {
|
||||
return Ok(());
|
||||
}
|
||||
|
||||
drop_guard(LOG_GUARD.get());
|
||||
drop_guard(CHROME_GUARD.get());
|
||||
|
||||
Ok(())
|
||||
}
|
||||
|
||||
fn drop_guard<T>(guard: Option<&Mutex<Option<T>>>) {
|
||||
if let Some(lock) = guard
|
||||
&& let Ok(mut slot) = lock.lock()
|
||||
{
|
||||
let _ = slot.take();
|
||||
}
|
||||
}
|
||||
|
||||
pub struct InstrumentationFinalizer;
|
||||
|
||||
impl Drop for InstrumentationFinalizer {
|
||||
fn drop(&mut self) {
|
||||
let _ = finalize();
|
||||
}
|
||||
}
|
||||
|
||||
pub fn finalizer() -> InstrumentationFinalizer {
|
||||
InstrumentationFinalizer
|
||||
}
|
||||
|
||||
fn parse_timestamp_us(timestamp: &str) -> Option<i64> {
|
||||
let parsed: DateTime<FixedOffset> = DateTime::parse_from_rfc3339(timestamp).ok()?;
|
||||
Some(parsed.timestamp_micros())
|
||||
}
|
||||
|
||||
fn parse_thread_id(value: &Value) -> Option<i64> {
|
||||
match value {
|
||||
Value::Number(n) => n.as_i64(),
|
||||
Value::String(s) => s.parse::<i64>().ok(),
|
||||
_ => None,
|
||||
}
|
||||
}
|
||||
|
||||
pub struct InstrumentationTimer {
|
||||
name: &'static str,
|
||||
start: Instant,
|
||||
fields: Vec<(String, Value)>,
|
||||
mode: InstrumentationMode,
|
||||
_span_guard: Option<tracing::span::EnteredSpan>,
|
||||
}
|
||||
|
||||
impl InstrumentationTimer {
|
||||
pub fn new(name: &'static str) -> Self {
|
||||
Self {
|
||||
name,
|
||||
start: Instant::now(),
|
||||
fields: Vec::new(),
|
||||
mode: mode(),
|
||||
_span_guard: None,
|
||||
}
|
||||
}
|
||||
|
||||
pub fn new_with_span(
|
||||
name: &'static str,
|
||||
mode: InstrumentationMode,
|
||||
span_guard: Option<tracing::span::EnteredSpan>,
|
||||
) -> Self {
|
||||
Self {
|
||||
name,
|
||||
start: Instant::now(),
|
||||
fields: Vec::new(),
|
||||
mode,
|
||||
_span_guard: span_guard,
|
||||
}
|
||||
}
|
||||
|
||||
pub fn with_field(&mut self, key: impl Into<String>, value: impl Into<Value>) -> &mut Self {
|
||||
if self.mode != InstrumentationMode::Disabled && self.mode != InstrumentationMode::Chrome {
|
||||
self.fields.push((key.into(), value.into()));
|
||||
}
|
||||
self
|
||||
}
|
||||
}
|
||||
|
||||
impl Drop for InstrumentationTimer {
|
||||
fn drop(&mut self) {
|
||||
if self.mode == InstrumentationMode::Disabled {
|
||||
return;
|
||||
}
|
||||
if self.mode == InstrumentationMode::Chrome {
|
||||
let _ = self._span_guard.take();
|
||||
return;
|
||||
}
|
||||
let elapsed_us = self.start.elapsed().as_micros() as u64;
|
||||
if self.fields.is_empty() {
|
||||
tracing::info!(
|
||||
target: TARGET,
|
||||
event = "timing",
|
||||
name = self.name,
|
||||
elapsed_us = elapsed_us,
|
||||
);
|
||||
return;
|
||||
}
|
||||
|
||||
let mut map = serde_json::Map::new();
|
||||
for (key, value) in std::mem::take(&mut self.fields) {
|
||||
map.insert(key, value);
|
||||
}
|
||||
|
||||
let fields = Value::Object(map);
|
||||
tracing::info!(
|
||||
target: TARGET,
|
||||
event = "timing",
|
||||
name = self.name,
|
||||
elapsed_us = elapsed_us,
|
||||
fields = ?fields
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
pub fn timer(name: &'static str) -> InstrumentationTimer {
|
||||
InstrumentationTimer::new(name)
|
||||
}
|
||||
41
crates/codegen/xai-grok-telemetry/src/lib.rs
Normal file
41
crates/codegen/xai-grok-telemetry/src/lib.rs
Normal file
|
|
@ -0,0 +1,41 @@
|
|||
//! Telemetry engine for Grok Build sessions: product events + Mixpanel emission +
|
||||
//! Sentry error reporting + OpenTelemetry tracing + structured unified log.
|
||||
//!
|
||||
//! Extracted from `xai-file-utils` per review feedback so telemetry has
|
||||
//! its own ownership boundary (see CODEOWNERS) and so downstream consumers
|
||||
//! that only want event tracking + inference metrics no longer pull in
|
||||
//! Mixpanel/HTTP/identity dependencies.
|
||||
|
||||
mod appender;
|
||||
pub mod client;
|
||||
pub mod config;
|
||||
pub mod context;
|
||||
pub mod debug_log;
|
||||
pub mod enums;
|
||||
pub mod events;
|
||||
pub mod external;
|
||||
pub mod hooks_log;
|
||||
pub mod http;
|
||||
pub mod id;
|
||||
pub mod instrumentation;
|
||||
pub mod memory_log;
|
||||
pub mod memory_telemetry;
|
||||
pub mod otel_layer;
|
||||
pub(crate) mod otlp_http;
|
||||
pub mod prompt_timing;
|
||||
pub(crate) mod redact_common;
|
||||
pub mod sampling_log;
|
||||
pub mod sentry;
|
||||
pub mod session_ctx;
|
||||
pub mod session_metrics;
|
||||
pub mod unified_log;
|
||||
|
||||
pub use client::{
|
||||
Metadata, TelemetryClient, UserContext, init, init_if_needed, is_enabled,
|
||||
is_session_metrics_enabled,
|
||||
};
|
||||
pub use events::TelemetryEvent;
|
||||
pub use session_ctx::{
|
||||
EmitterOrigin, TelemetryCtx, emit_event, emit_event_with_origin, log_event, log_session_event,
|
||||
log_session_event_with_origin, with_session_ctx,
|
||||
};
|
||||
127
crates/codegen/xai-grok-telemetry/src/memory_log.rs
Normal file
127
crates/codegen/xai-grok-telemetry/src/memory_log.rs
Normal file
|
|
@ -0,0 +1,127 @@
|
|||
//! Memory system tracing target and optional file-based logging layer.
|
||||
//!
|
||||
//! Provides a dedicated tracing target (`xai_memory`) with an optional
|
||||
//! file logger that writes to `~/.grok/logs/memory.log`.
|
||||
//!
|
||||
//! ## When to use
|
||||
//!
|
||||
//! Use `tracing::info!(target: memory_log::TARGET, ...)` at memory system
|
||||
//! lifecycle points — config resolution, storage init, flush, search, etc.
|
||||
//! These events are always emitted (zero cost when the layer is absent).
|
||||
//!
|
||||
//! ## Enabling (debug builds)
|
||||
//!
|
||||
//! ```bash
|
||||
//! # build with memory logging enabled, then:
|
||||
//! GROK_MEMORY_LOG=0 grok # disable even when enabled
|
||||
//! tail -f ~/.grok/logs/memory.log # watch in another terminal
|
||||
//! ```
|
||||
|
||||
/// Tracing target for all memory system operations.
|
||||
pub const TARGET: &str = "xai_memory";
|
||||
|
||||
#[cfg(feature = "memory-log")]
|
||||
mod inner {
|
||||
use std::fmt;
|
||||
use std::path::PathBuf;
|
||||
use std::sync::Mutex;
|
||||
use std::time::Instant;
|
||||
|
||||
use tracing::Subscriber;
|
||||
use tracing_subscriber::fmt::format::Writer;
|
||||
use tracing_subscriber::fmt::time::FormatTime;
|
||||
use tracing_subscriber::fmt::writer::BoxMakeWriter;
|
||||
use tracing_subscriber::layer::Layer;
|
||||
use tracing_subscriber::registry::LookupSpan;
|
||||
|
||||
use super::TARGET;
|
||||
use xai_grok_config::grok_home;
|
||||
|
||||
const ENV_MEMORY_LOG: &str = "GROK_MEMORY_LOG";
|
||||
|
||||
static LOG_GUARD: std::sync::OnceLock<
|
||||
Mutex<Option<tracing_appender::non_blocking::WorkerGuard>>,
|
||||
> = std::sync::OnceLock::new();
|
||||
|
||||
#[derive(Clone)]
|
||||
struct UptimeTimer {
|
||||
epoch: Instant,
|
||||
}
|
||||
|
||||
impl UptimeTimer {
|
||||
fn new() -> Self {
|
||||
Self {
|
||||
epoch: Instant::now(),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl FormatTime for UptimeTimer {
|
||||
fn format_time(&self, w: &mut Writer<'_>) -> fmt::Result {
|
||||
let elapsed = self.epoch.elapsed();
|
||||
write!(w, "+{}.{:03}s", elapsed.as_secs(), elapsed.subsec_millis())
|
||||
}
|
||||
}
|
||||
|
||||
/// Build the memory log layer.
|
||||
///
|
||||
/// Writes to `~/.grok/logs/memory.log`. Filters to `xai_memory=trace`.
|
||||
/// Set `GROK_MEMORY_LOG=0` to disable, `GROK_MEMORY_LOG=/path` to redirect.
|
||||
pub fn layer<S>() -> Option<impl Layer<S>>
|
||||
where
|
||||
S: Subscriber + for<'span> LookupSpan<'span>,
|
||||
{
|
||||
let path = resolve_log_path()?;
|
||||
|
||||
if let Some(parent) = path.parent() {
|
||||
let _ = std::fs::create_dir_all(parent);
|
||||
}
|
||||
|
||||
let file = match std::fs::OpenOptions::new()
|
||||
.create(true)
|
||||
.append(true)
|
||||
.open(&path)
|
||||
{
|
||||
Ok(f) => f,
|
||||
Err(e) => {
|
||||
tracing::warn!("[memory-log] Failed to open {:?}: {}", path, e);
|
||||
return None;
|
||||
}
|
||||
};
|
||||
|
||||
let (non_blocking, guard) = tracing_appender::non_blocking(file);
|
||||
let guard_slot = LOG_GUARD.get_or_init(|| Mutex::new(None));
|
||||
if let Ok(mut slot) = guard_slot.lock() {
|
||||
*slot = Some(guard);
|
||||
}
|
||||
|
||||
let filter = tracing_subscriber::filter::EnvFilter::new(format!("{TARGET}=trace"));
|
||||
let fmt_layer = tracing_subscriber::fmt::layer()
|
||||
.with_target(true)
|
||||
.with_ansi(false)
|
||||
.with_thread_ids(true)
|
||||
.with_timer(UptimeTimer::new())
|
||||
.with_writer(BoxMakeWriter::new(non_blocking))
|
||||
.with_filter(filter);
|
||||
|
||||
tracing::info!("[memory-log] Memory logging enabled");
|
||||
Some(fmt_layer)
|
||||
}
|
||||
|
||||
fn resolve_log_path() -> Option<PathBuf> {
|
||||
let default_path = || grok_home().join("logs").join("memory.log");
|
||||
let raw = match std::env::var(ENV_MEMORY_LOG) {
|
||||
Ok(val) => val,
|
||||
Err(_) => return Some(default_path()),
|
||||
};
|
||||
let raw = raw.trim();
|
||||
match raw {
|
||||
"" | "0" | "false" | "off" | "no" => None,
|
||||
"1" | "true" | "on" | "yes" => Some(default_path()),
|
||||
path => Some(PathBuf::from(path)),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(feature = "memory-log")]
|
||||
pub use inner::layer;
|
||||
118
crates/codegen/xai-grok-telemetry/src/memory_telemetry.rs
Normal file
118
crates/codegen/xai-grok-telemetry/src/memory_telemetry.rs
Normal file
|
|
@ -0,0 +1,118 @@
|
|||
//! Memory subsystem telemetry. Routes through `log_event` (product tier,
|
||||
//! `Enabled` mode only). No PII or user content -- only counts, scores,
|
||||
//! durations, and config values.
|
||||
|
||||
use serde::Serialize;
|
||||
|
||||
#[derive(Serialize)]
|
||||
pub struct MemorySessionInit {
|
||||
pub session_id: String,
|
||||
pub memory_enabled: bool,
|
||||
pub watcher_config_enabled: bool,
|
||||
pub watcher_started: bool,
|
||||
pub temporal_decay_enabled: bool,
|
||||
pub mmr_enabled: bool,
|
||||
pub mmr_lambda: f64,
|
||||
pub half_life_days: f64,
|
||||
pub embedding_dimensions: usize,
|
||||
pub total_chunks: usize,
|
||||
pub total_files: usize,
|
||||
pub has_global_memory_md: bool,
|
||||
pub has_workspace_memory_md: bool,
|
||||
}
|
||||
|
||||
#[derive(Serialize)]
|
||||
pub struct MemorySearch {
|
||||
pub session_id: String,
|
||||
pub query_length: usize,
|
||||
pub keyword_count: usize,
|
||||
pub result_count: usize,
|
||||
pub top_score: f64,
|
||||
pub min_score_threshold: f64,
|
||||
pub search_mode: String,
|
||||
pub duration_ms: u64,
|
||||
pub vec_available: bool,
|
||||
pub source: String,
|
||||
}
|
||||
|
||||
#[derive(Serialize)]
|
||||
pub struct MemorySearchEmpty {
|
||||
pub session_id: String,
|
||||
pub query_length: usize,
|
||||
pub keyword_count: usize,
|
||||
pub min_score_threshold: f64,
|
||||
pub search_mode: String,
|
||||
pub duration_ms: u64,
|
||||
pub vec_available: bool,
|
||||
pub source: String,
|
||||
}
|
||||
|
||||
#[derive(Serialize)]
|
||||
pub struct MemoryFlushStart {
|
||||
pub session_id: String,
|
||||
pub trigger: String,
|
||||
pub conversation_len: usize,
|
||||
pub user_message_count: usize,
|
||||
}
|
||||
|
||||
#[derive(Serialize)]
|
||||
pub struct MemoryFlushComplete {
|
||||
pub session_id: String,
|
||||
pub trigger: String,
|
||||
pub outcome: String,
|
||||
pub duration_ms: u64,
|
||||
pub response_length: usize,
|
||||
pub accepted_length: usize,
|
||||
pub was_truncated: bool,
|
||||
}
|
||||
|
||||
#[derive(Serialize)]
|
||||
pub struct MemoryInjection {
|
||||
pub session_id: String,
|
||||
pub was_greeting_fallback: bool,
|
||||
pub result_count: usize,
|
||||
pub total_snippet_chars: usize,
|
||||
pub top_score: f64,
|
||||
pub configured_min_score: f64,
|
||||
pub injection_duration_ms: u64,
|
||||
}
|
||||
|
||||
#[derive(Serialize)]
|
||||
pub struct MemoryReindex {
|
||||
pub session_id: String,
|
||||
pub source: String,
|
||||
pub added: usize,
|
||||
pub updated: usize,
|
||||
pub removed: usize,
|
||||
pub embedded: usize,
|
||||
pub duration_ms: u64,
|
||||
pub trigger: String,
|
||||
}
|
||||
|
||||
#[derive(Serialize)]
|
||||
pub struct MemoryWatcherSync {
|
||||
pub session_id: String,
|
||||
pub dirty_file_count: usize,
|
||||
pub claimed: bool,
|
||||
pub reindexed_count: usize,
|
||||
pub embedded_count: usize,
|
||||
pub duration_ms: u64,
|
||||
}
|
||||
|
||||
#[derive(Serialize)]
|
||||
pub struct MemorySessionSummary {
|
||||
pub session_id: String,
|
||||
pub session_duration_secs: u64,
|
||||
pub flush_count: u64,
|
||||
pub flush_success_count: u64,
|
||||
pub flush_error_count: u64,
|
||||
pub tool_search_count: u64,
|
||||
pub injection_count: u64,
|
||||
pub recovery_search_count: u64,
|
||||
pub total_chunks_at_end: usize,
|
||||
pub chunks_added_this_session: usize,
|
||||
pub session_end_result: String,
|
||||
pub dream_count: u64,
|
||||
pub dream_success_count: u64,
|
||||
pub dream_error_count: u64,
|
||||
}
|
||||
772
crates/codegen/xai-grok-telemetry/src/otel_layer/mod.rs
Normal file
772
crates/codegen/xai-grok-telemetry/src/otel_layer/mod.rs
Normal file
|
|
@ -0,0 +1,772 @@
|
|||
//! Shared OpenTelemetry tracing layer for exporting spans to the cli-chat-proxy.
|
||||
//!
|
||||
//! `xai-grok-pager` uses this module to set up OTLP
|
||||
//! trace export so that session-level spans (with `session_id`, tool timings,
|
||||
//! inference latency, etc.) are available in the product observability backend.
|
||||
use crate::instrumentation;
|
||||
use opentelemetry::global;
|
||||
use opentelemetry::trace::TracerProvider as _;
|
||||
use opentelemetry_otlp::{WithExportConfig, WithHttpConfig};
|
||||
use opentelemetry_sdk::trace::SdkTracerProvider;
|
||||
use std::sync::{Arc, OnceLock};
|
||||
use tracing_opentelemetry::OpenTelemetryLayer;
|
||||
use tracing_subscriber::Layer as _;
|
||||
use tracing_subscriber::registry::LookupSpan;
|
||||
use xai_grok_auth::AuthCredentialProvider;
|
||||
mod redact;
|
||||
static TRACER_PROVIDER: OnceLock<SdkTracerProvider> = OnceLock::new();
|
||||
const ENV_OTEL_FILTER: &str = "GROK_OTEL_FILTER";
|
||||
const DEFAULT_OTEL_FILTER: &str = "info";
|
||||
/// Configuration for [`build_otel_layer`]. Encapsulates all the runtime values
|
||||
/// the layer needs that used to be reach-ins into shell-internal types
|
||||
/// (`AuthManager`, `EndpointsConfig`, `GrokComConfig`).
|
||||
///
|
||||
/// Built by the binaries (`xai-grok-pager`) from their own
|
||||
/// configuration. The credentials provider is constructed by shell's
|
||||
/// `xai_grok_shell::auth::credential_provider::build_otel_credential_provider`.
|
||||
pub struct OtelLayerConfig {
|
||||
/// Live credential source. Read on every batch export to obtain a fresh
|
||||
/// bearer token for the OTLP `Authorization` header.
|
||||
pub credentials: Arc<dyn AuthCredentialProvider>,
|
||||
/// Value for the `X-XAI-Token-Auth` header (typically `"xai-grok-cli"`).
|
||||
pub token_header_value: String,
|
||||
/// Optional extra access key for traces. Injection is honored only when
|
||||
/// the crate's optional non-production feature is enabled and only for
|
||||
/// matching first-party hosts. Field stays present so cross-crate
|
||||
/// constructors compile regardless of per-crate feature unification.
|
||||
pub alpha_test_key: Option<String>,
|
||||
pub exporter: OtelExporterConfig,
|
||||
}
|
||||
/// Static identity of the client emitting telemetry. Becomes resource
|
||||
/// attributes (`client.name`, `client.version`, `service.version`,
|
||||
/// `app.entrypoint`) on every span.
|
||||
#[derive(Debug, Clone, Copy)]
|
||||
pub struct OtelClientInfo {
|
||||
/// Binary name (`grok-pager`) -> `client.name`.
|
||||
pub client_name: &'static str,
|
||||
/// Front-end client version -> `client.version`.
|
||||
pub client_version: &'static str,
|
||||
/// Engine build (version + commit) -> `service.version`.
|
||||
pub service_version: &'static str,
|
||||
/// How the session was launched (`cli`/`headless`/`agent`) -> `app.entrypoint`.
|
||||
pub app_entrypoint: &'static str,
|
||||
}
|
||||
/// OTLP trace-export transport settings, resolved from the `OTEL_*` env vars /
|
||||
/// managed config.
|
||||
#[derive(Debug, Default, Clone)]
|
||||
pub struct OtelExporterConfig {
|
||||
/// Full OTLP traces endpoint URL (e.g. `https://cli-chat-proxy.grok.com/v1/traces`).
|
||||
pub traces_url: String,
|
||||
/// `OTEL_EXPORTER_OTLP_HEADERS` pairs.
|
||||
pub extra_headers: Vec<(String, String)>,
|
||||
/// `OTEL_TRACES_EXPORT_INTERVAL` batch flush interval. `None` = SDK default.
|
||||
pub export_interval: Option<std::time::Duration>,
|
||||
/// `OTEL_EXPORTER_OTLP_TIMEOUT` export timeout. `None` = 10s default.
|
||||
pub timeout: Option<std::time::Duration>,
|
||||
/// `false` when `OTEL_TRACES_EXPORTER=none`: spans created, never exported.
|
||||
pub enabled: bool,
|
||||
}
|
||||
/// Creates an OpenTelemetry layer that bridges tracing spans to OpenTelemetry.
|
||||
/// This enables trace context propagation and OTLP export to the cli-chat-proxy.
|
||||
///
|
||||
/// - `client_name`: binary name (e.g. `"grok-tui"`, `"grok-pager"`) -- stored as
|
||||
/// `client.name` resource attribute for dashboards to distinguish client types.
|
||||
/// - `client_version`: `CARGO_PKG_VERSION` -- sent in the `x-grok-client-version` header.
|
||||
/// - `service_version`: `VERSION_WITH_COMMIT` -- stored as `service.version` resource attribute.
|
||||
/// - `config`: runtime configuration; see [`OtelLayerConfig`].
|
||||
pub fn build_otel_layer<S>(
|
||||
client: OtelClientInfo,
|
||||
config: OtelLayerConfig,
|
||||
) -> impl tracing_subscriber::layer::Layer<S>
|
||||
where
|
||||
S: tracing::Subscriber + for<'span> LookupSpan<'span>,
|
||||
{
|
||||
let provider = TRACER_PROVIDER.get_or_init(|| build_tracer_provider(client, config));
|
||||
let tracer = provider.tracer("grok-cli");
|
||||
global::set_tracer_provider(provider.clone());
|
||||
global::set_text_map_propagator(opentelemetry_sdk::propagation::TraceContextPropagator::new());
|
||||
let otel_filter =
|
||||
std::env::var(ENV_OTEL_FILTER).unwrap_or_else(|_| DEFAULT_OTEL_FILTER.to_string());
|
||||
let otel_filter = tracing_subscriber::filter::EnvFilter::try_new(&otel_filter)
|
||||
.unwrap_or_else(|e| {
|
||||
eprintln!(
|
||||
"[otel] Invalid GROK_OTEL_FILTER '{}': {}. Using default '{}'.",
|
||||
otel_filter, e, DEFAULT_OTEL_FILTER
|
||||
);
|
||||
tracing_subscriber::filter::EnvFilter::try_new(DEFAULT_OTEL_FILTER)
|
||||
.expect("default otel filter must parse")
|
||||
})
|
||||
.add_directive(
|
||||
"sampling_log=off"
|
||||
.parse()
|
||||
.expect("static directive must parse"),
|
||||
);
|
||||
OpenTelemetryLayer::new(tracer)
|
||||
.with_context_activation(false)
|
||||
.with_filter(otel_filter)
|
||||
}
|
||||
fn build_tracer_provider(client: OtelClientInfo, config: OtelLayerConfig) -> SdkTracerProvider {
|
||||
match instrumentation::current_mode() {
|
||||
instrumentation::InstrumentationMode::Server => build_server_provider(client, config),
|
||||
_ => SdkTracerProvider::builder().build(),
|
||||
}
|
||||
}
|
||||
/// Wraps an OTLP `SpanExporter`, rebuilding it with a fresh auth token on each
|
||||
/// `export()` call if the in-memory auth token has changed. On export failure,
|
||||
/// attempts a token refresh and retries once.
|
||||
struct RefreshableSpanExporter {
|
||||
endpoint: Arc<str>,
|
||||
static_headers: Arc<std::collections::HashMap<String, String>>,
|
||||
credentials: Arc<dyn AuthCredentialProvider>,
|
||||
last_token: parking_lot::Mutex<String>,
|
||||
/// Pre-built HTTP client shared across all export calls. Created once at
|
||||
/// init (outside the batch processor thread) to avoid the "no reactor"
|
||||
/// panic that occurs when `hyper-util` tries DNS resolution on a non-Tokio
|
||||
/// thread.
|
||||
http_client: crate::otlp_http::BlockingOtlpClient,
|
||||
/// Resource set by the `BatchSpanProcessor` via `set_resource()`.
|
||||
/// Forwarded to each one-shot exporter so OTLP payloads include
|
||||
/// `service.name`, `service.version`, `user.id`, etc.
|
||||
resource: parking_lot::Mutex<opentelemetry_sdk::Resource>,
|
||||
/// Value for `X-XAI-Token-Auth`. Only sent when `credentials.needs_token_auth_header()`.
|
||||
token_header_value: Arc<str>,
|
||||
extra_headers: Arc<Vec<(String, String)>>,
|
||||
}
|
||||
impl std::fmt::Debug for RefreshableSpanExporter {
|
||||
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
|
||||
f.debug_struct("RefreshableSpanExporter")
|
||||
.field("endpoint", &self.endpoint)
|
||||
.field("static_headers", &self.static_headers)
|
||||
.field("credentials", &"configured")
|
||||
.field("last_token", &"***")
|
||||
.finish_non_exhaustive()
|
||||
}
|
||||
}
|
||||
/// Build the header map for an OTLP export request.
|
||||
fn build_export_headers(
|
||||
static_headers: &std::collections::HashMap<String, String>,
|
||||
token: &str,
|
||||
token_auth_header: Option<&str>,
|
||||
extra_headers: &[(String, String)],
|
||||
snapshot: &xai_grok_auth::CredentialSnapshot,
|
||||
) -> std::collections::HashMap<String, String> {
|
||||
let mut headers = static_headers.clone();
|
||||
for (name, value) in [
|
||||
("x-userid", &snapshot.user_id),
|
||||
("x-teamid", &snapshot.team_id),
|
||||
] {
|
||||
match value.as_deref().filter(|v| !v.is_empty()) {
|
||||
Some(v) => {
|
||||
headers.insert(name.to_string(), v.to_string());
|
||||
}
|
||||
None => {
|
||||
headers.remove(name);
|
||||
}
|
||||
}
|
||||
}
|
||||
headers.insert("Authorization".to_string(), format!("Bearer {token}"));
|
||||
if let Some(value) = token_auth_header {
|
||||
headers.insert("X-XAI-Token-Auth".to_string(), value.to_string());
|
||||
}
|
||||
for (k, v) in extra_headers {
|
||||
headers.insert(k.clone(), v.clone());
|
||||
}
|
||||
headers
|
||||
}
|
||||
fn build_otlp_exporter(
|
||||
endpoint: &str,
|
||||
static_headers: &std::collections::HashMap<String, String>,
|
||||
token: &str,
|
||||
token_auth_header: Option<&str>,
|
||||
extra_headers: &[(String, String)],
|
||||
http_client: crate::otlp_http::BlockingOtlpClient,
|
||||
snapshot: &xai_grok_auth::CredentialSnapshot,
|
||||
) -> Result<opentelemetry_otlp::SpanExporter, opentelemetry_otlp::ExporterBuildError> {
|
||||
let headers = build_export_headers(
|
||||
static_headers,
|
||||
token,
|
||||
token_auth_header,
|
||||
extra_headers,
|
||||
snapshot,
|
||||
);
|
||||
opentelemetry_otlp::SpanExporter::builder()
|
||||
.with_http()
|
||||
.with_http_client(http_client)
|
||||
.with_endpoint(endpoint)
|
||||
.with_headers(headers)
|
||||
.build()
|
||||
}
|
||||
/// Send a batch through an exporter with `set_resource` applied.
|
||||
async fn export_batch(
|
||||
exporter: &mut opentelemetry_otlp::SpanExporter,
|
||||
resource: &opentelemetry_sdk::Resource,
|
||||
batch: Vec<opentelemetry_sdk::trace::SpanData>,
|
||||
) -> opentelemetry_sdk::error::OTelSdkResult {
|
||||
use opentelemetry_sdk::trace::SpanExporter as _;
|
||||
exporter.set_resource(resource);
|
||||
exporter.export(batch).await
|
||||
}
|
||||
impl RefreshableSpanExporter {
|
||||
#[cfg(test)]
|
||||
fn current_token(&self) -> String {
|
||||
self.credentials.snapshot().token.unwrap_or_else(|| {
|
||||
tracing::debug!("auth: otel credential snapshot has no token, using cached last_token");
|
||||
self.last_token.lock().clone()
|
||||
})
|
||||
}
|
||||
}
|
||||
/// Stamp `deployment.id`/`api_key.id`/`organization.id`/`team.id`/`user.id`
|
||||
/// per-export (they're only known after auth is wired, post-init — stamping at
|
||||
/// init would leave them blank for a session that authenticates mid-run).
|
||||
fn resource_with_tenant_id(
|
||||
base: opentelemetry_sdk::Resource,
|
||||
snapshot: &xai_grok_auth::CredentialSnapshot,
|
||||
) -> opentelemetry_sdk::Resource {
|
||||
let tenant_attrs: Vec<opentelemetry::KeyValue> = [
|
||||
("deployment.id", &snapshot.deployment_id),
|
||||
("api_key.id", &snapshot.api_key_id),
|
||||
("organization.id", &snapshot.organization_id),
|
||||
("team.id", &snapshot.team_id),
|
||||
("user.id", &snapshot.user_id),
|
||||
]
|
||||
.into_iter()
|
||||
.filter_map(|(key, val)| {
|
||||
val.as_deref()
|
||||
.filter(|v| !v.is_empty())
|
||||
.map(|v| opentelemetry::KeyValue::new(key, v.to_string()))
|
||||
})
|
||||
.collect();
|
||||
if tenant_attrs.is_empty() {
|
||||
return base;
|
||||
}
|
||||
let mut attrs: Vec<opentelemetry::KeyValue> = base
|
||||
.iter()
|
||||
.map(|(k, v)| opentelemetry::KeyValue::new(k.clone(), v.clone()))
|
||||
.collect();
|
||||
attrs.extend(tenant_attrs);
|
||||
opentelemetry_sdk::Resource::builder_empty()
|
||||
.with_attributes(attrs)
|
||||
.build()
|
||||
}
|
||||
/// Inputs for one export attempt, built on the calling thread: the
|
||||
/// `BatchSpanProcessor` drives `export()` from a non-Tokio `std::thread`, so
|
||||
/// constructing the exporter/HTTP client inside the future would hit the
|
||||
/// "no reactor" panic.
|
||||
struct ExportInputs {
|
||||
one_shot: Result<opentelemetry_otlp::SpanExporter, opentelemetry_otlp::ExporterBuildError>,
|
||||
resource: opentelemetry_sdk::Resource,
|
||||
credentials: Arc<dyn AuthCredentialProvider>,
|
||||
endpoint: Arc<str>,
|
||||
static_headers: Arc<std::collections::HashMap<String, String>>,
|
||||
token_header_value: Arc<str>,
|
||||
http_client: crate::otlp_http::BlockingOtlpClient,
|
||||
extra_headers: Arc<Vec<(String, String)>>,
|
||||
}
|
||||
impl opentelemetry_sdk::trace::SpanExporter for RefreshableSpanExporter {
|
||||
fn export(
|
||||
&self,
|
||||
batch: Vec<opentelemetry_sdk::trace::SpanData>,
|
||||
) -> impl std::future::Future<Output = opentelemetry_sdk::error::OTelSdkResult> + Send {
|
||||
let prepared = (crate::client::is_session_metrics_enabled()
|
||||
&& self.credentials.has_usable_credential())
|
||||
.then(|| {
|
||||
let snapshot = self.credentials.snapshot();
|
||||
let token = snapshot.token.clone().unwrap_or_else(|| {
|
||||
tracing::debug!(
|
||||
"auth: otel credential snapshot has no token, using cached last_token"
|
||||
);
|
||||
self.last_token.lock().clone()
|
||||
});
|
||||
*self.last_token.lock() = token.clone();
|
||||
let token_auth = self
|
||||
.credentials
|
||||
.needs_token_auth_header()
|
||||
.then(|| Arc::clone(&self.token_header_value));
|
||||
ExportInputs {
|
||||
one_shot: build_otlp_exporter(
|
||||
&self.endpoint,
|
||||
&self.static_headers,
|
||||
&token,
|
||||
token_auth.as_deref(),
|
||||
&self.extra_headers,
|
||||
self.http_client.clone(),
|
||||
&snapshot,
|
||||
),
|
||||
resource: resource_with_tenant_id(self.resource.lock().clone(), &snapshot),
|
||||
credentials: Arc::clone(&self.credentials),
|
||||
endpoint: Arc::clone(&self.endpoint),
|
||||
static_headers: Arc::clone(&self.static_headers),
|
||||
token_header_value: Arc::clone(&self.token_header_value),
|
||||
http_client: self.http_client.clone(),
|
||||
extra_headers: Arc::clone(&self.extra_headers),
|
||||
}
|
||||
});
|
||||
async move {
|
||||
let Some(ExportInputs {
|
||||
one_shot,
|
||||
resource,
|
||||
credentials,
|
||||
endpoint,
|
||||
static_headers,
|
||||
token_header_value,
|
||||
http_client,
|
||||
extra_headers,
|
||||
}) = prepared
|
||||
else {
|
||||
return Ok(());
|
||||
};
|
||||
let mut exporter = match one_shot {
|
||||
Ok(e) => e,
|
||||
Err(e) => {
|
||||
return Err(opentelemetry_sdk::error::OTelSdkError::InternalFailure(
|
||||
format!("failed to build exporter: {e}"),
|
||||
));
|
||||
}
|
||||
};
|
||||
let mut batch = batch;
|
||||
redact::redact_batch(&mut batch);
|
||||
let batch_for_retry = tokio::runtime::Handle::try_current()
|
||||
.is_ok()
|
||||
.then(|| batch.clone());
|
||||
let result = export_batch(&mut exporter, &resource, batch).await;
|
||||
if result.is_ok() {
|
||||
return result;
|
||||
}
|
||||
let Some(batch_for_retry) = batch_for_retry else {
|
||||
return result;
|
||||
};
|
||||
tracing::debug!("otel export failed, attempting token refresh");
|
||||
if !credentials.refresh_after_unauthorized().await {
|
||||
return result;
|
||||
}
|
||||
let retry_snapshot = credentials.snapshot();
|
||||
let new_token = retry_snapshot.token.clone().unwrap_or_default();
|
||||
if new_token.is_empty() {
|
||||
tracing::warn!("token refresh reported success but snapshot returned no token");
|
||||
return result;
|
||||
}
|
||||
let retry_token_auth = credentials
|
||||
.needs_token_auth_header()
|
||||
.then(|| token_header_value.as_ref());
|
||||
match build_otlp_exporter(
|
||||
&endpoint,
|
||||
&static_headers,
|
||||
&new_token,
|
||||
retry_token_auth,
|
||||
&extra_headers,
|
||||
http_client,
|
||||
&retry_snapshot,
|
||||
) {
|
||||
Ok(mut retry_exporter) => {
|
||||
let retry_resource = resource_with_tenant_id(resource, &retry_snapshot);
|
||||
export_batch(&mut retry_exporter, &retry_resource, batch_for_retry)
|
||||
.await
|
||||
.or(result)
|
||||
}
|
||||
Err(e) => {
|
||||
tracing::debug!("failed to build retry exporter: {e}");
|
||||
result
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
fn set_resource(&mut self, resource: &opentelemetry_sdk::Resource) {
|
||||
*self.resource.lock() = resource.clone();
|
||||
}
|
||||
fn shutdown(&self) -> opentelemetry_sdk::error::OTelSdkResult {
|
||||
Ok(())
|
||||
}
|
||||
}
|
||||
fn build_server_provider(client: OtelClientInfo, config: OtelLayerConfig) -> SdkTracerProvider {
|
||||
let OtelClientInfo {
|
||||
client_name,
|
||||
client_version,
|
||||
service_version,
|
||||
app_entrypoint,
|
||||
} = client;
|
||||
let snapshot = config.credentials.snapshot();
|
||||
let initial_token = snapshot.token.unwrap_or_default();
|
||||
if initial_token.is_empty() {
|
||||
tracing::debug!(
|
||||
"No authentication credentials found at init. OTLP exporter will retry after auth."
|
||||
);
|
||||
}
|
||||
let mut resource_attrs = vec![
|
||||
opentelemetry::KeyValue::new("service.version", service_version.to_string()),
|
||||
opentelemetry::KeyValue::new("client.name", client_name.to_string()),
|
||||
opentelemetry::KeyValue::new("client.version", client_version.to_string()),
|
||||
opentelemetry::KeyValue::new("app.entrypoint", app_entrypoint.to_string()),
|
||||
];
|
||||
if let Some(terminal_type) = std::env::var("TERM_PROGRAM")
|
||||
.ok()
|
||||
.or_else(|| std::env::var("TERM").ok())
|
||||
.filter(|v| !v.is_empty())
|
||||
{
|
||||
resource_attrs.push(opentelemetry::KeyValue::new("terminal.type", terminal_type));
|
||||
}
|
||||
let mut provider = SdkTracerProvider::builder().with_resource(
|
||||
opentelemetry_sdk::Resource::builder_empty()
|
||||
.with_service_name("grok-cli")
|
||||
.with_attributes(resource_attrs)
|
||||
.build(),
|
||||
);
|
||||
if config.exporter.enabled {
|
||||
let traces_url = config.exporter.traces_url;
|
||||
let mut static_headers = std::collections::HashMap::new();
|
||||
static_headers.insert(
|
||||
"x-grok-client-version".to_string(),
|
||||
client_version.to_string(),
|
||||
);
|
||||
let timeout = config
|
||||
.exporter
|
||||
.timeout
|
||||
.unwrap_or(std::time::Duration::from_secs(10));
|
||||
let http_client = match crate::otlp_http::build_blocking_client(timeout) {
|
||||
Ok(client) => client,
|
||||
Err(err) => {
|
||||
tracing::warn!(
|
||||
error = % err,
|
||||
"otel: OTLP HTTP client build failed; span export disabled"
|
||||
);
|
||||
return provider.build();
|
||||
}
|
||||
};
|
||||
let refreshable_exporter = RefreshableSpanExporter {
|
||||
endpoint: Arc::from(traces_url),
|
||||
static_headers: Arc::new(static_headers),
|
||||
credentials: config.credentials,
|
||||
last_token: parking_lot::Mutex::new(initial_token),
|
||||
http_client,
|
||||
resource: parking_lot::Mutex::new(opentelemetry_sdk::Resource::builder_empty().build()),
|
||||
token_header_value: Arc::from(config.token_header_value.as_str()),
|
||||
extra_headers: Arc::new(config.exporter.extra_headers),
|
||||
};
|
||||
let mut batch_builder =
|
||||
opentelemetry_sdk::trace::BatchConfigBuilder::default().with_max_export_batch_size(64);
|
||||
if let Some(interval) = config.exporter.export_interval {
|
||||
batch_builder = batch_builder.with_scheduled_delay(interval);
|
||||
}
|
||||
let batch_processor =
|
||||
opentelemetry_sdk::trace::BatchSpanProcessor::builder(refreshable_exporter)
|
||||
.with_batch_config(batch_builder.build())
|
||||
.build();
|
||||
provider = provider.with_span_processor(batch_processor);
|
||||
}
|
||||
provider.build()
|
||||
}
|
||||
/// Flush and shut down the global tracer provider (and the external OTEL
|
||||
/// stream — both ride the same exit chokepoints).
|
||||
///
|
||||
/// Prefer [`OtelGuard`] for normal code paths. Use this directly only in
|
||||
/// signal handlers or `process::exit` paths where destructors won't run.
|
||||
/// Safe to call multiple times (second call logs a warning but does not panic;
|
||||
/// the external shutdown is idempotent).
|
||||
pub fn shutdown_otel() {
|
||||
crate::external::shutdown();
|
||||
if let Some(provider) = TRACER_PROVIDER.get()
|
||||
&& let Err(e) = provider.shutdown()
|
||||
{
|
||||
tracing::debug!("[otel] Failed to shutdown tracer provider: {}", e);
|
||||
}
|
||||
}
|
||||
/// RAII guard that calls [`shutdown_otel`] on drop.
|
||||
pub struct OtelGuard;
|
||||
impl Drop for OtelGuard {
|
||||
fn drop(&mut self) {
|
||||
shutdown_otel();
|
||||
}
|
||||
}
|
||||
/// Create an [`OtelGuard`] that flushes traces on drop.
|
||||
pub fn otel_guard() -> OtelGuard {
|
||||
OtelGuard
|
||||
}
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
use async_trait::async_trait;
|
||||
use std::sync::Mutex;
|
||||
use xai_grok_auth::{AuthCredentialProvider, CredentialSnapshot, HttpAuth};
|
||||
/// Test double for `AuthCredentialProvider`. When constructed with
|
||||
/// `with_refresh`, `refresh_after_unauthorized` rotates the token and
|
||||
/// returns `true`; otherwise it returns `false`.
|
||||
struct TestProvider {
|
||||
token: Mutex<Option<String>>,
|
||||
refreshed_token: Option<String>,
|
||||
refresh_count: std::sync::atomic::AtomicU32,
|
||||
}
|
||||
impl TestProvider {
|
||||
fn new(initial: Option<&str>) -> Self {
|
||||
Self {
|
||||
token: Mutex::new(initial.map(|s| s.to_owned())),
|
||||
refreshed_token: None,
|
||||
refresh_count: std::sync::atomic::AtomicU32::new(0),
|
||||
}
|
||||
}
|
||||
fn with_refresh(initial: &str, refreshed: &str) -> Self {
|
||||
Self {
|
||||
token: Mutex::new(Some(initial.to_owned())),
|
||||
refreshed_token: Some(refreshed.to_owned()),
|
||||
refresh_count: std::sync::atomic::AtomicU32::new(0),
|
||||
}
|
||||
}
|
||||
fn set(&self, value: Option<&str>) {
|
||||
*self.token.lock().unwrap() = value.map(|s| s.to_owned());
|
||||
}
|
||||
fn refresh_count(&self) -> u32 {
|
||||
self.refresh_count
|
||||
.load(std::sync::atomic::Ordering::Relaxed)
|
||||
}
|
||||
}
|
||||
impl HttpAuth for TestProvider {
|
||||
fn apply(
|
||||
&self,
|
||||
builder: reqwest::RequestBuilder,
|
||||
_base_url: &str,
|
||||
) -> reqwest::RequestBuilder {
|
||||
builder
|
||||
}
|
||||
}
|
||||
#[async_trait]
|
||||
impl AuthCredentialProvider for TestProvider {
|
||||
fn snapshot(&self) -> CredentialSnapshot {
|
||||
CredentialSnapshot {
|
||||
token: self.token.lock().unwrap().clone(),
|
||||
..Default::default()
|
||||
}
|
||||
}
|
||||
async fn refresh_after_unauthorized(&self) -> bool {
|
||||
if let Some(ref refreshed) = self.refreshed_token {
|
||||
*self.token.lock().unwrap() = Some(refreshed.clone());
|
||||
self.refresh_count
|
||||
.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
|
||||
true
|
||||
} else {
|
||||
false
|
||||
}
|
||||
}
|
||||
}
|
||||
/// Must be called from a non-async test (`#[test]`, not `#[tokio::test]`).
|
||||
/// The blocking client spawns an internal tokio runtime that panics if
|
||||
/// dropped inside an async executor.
|
||||
fn make_exporter(
|
||||
provider: Arc<dyn AuthCredentialProvider>,
|
||||
last_token: &str,
|
||||
) -> RefreshableSpanExporter {
|
||||
RefreshableSpanExporter {
|
||||
endpoint: Arc::from("http://localhost:4318/v1/traces"),
|
||||
static_headers: Arc::new(std::collections::HashMap::new()),
|
||||
credentials: provider,
|
||||
last_token: parking_lot::Mutex::new(last_token.to_string()),
|
||||
http_client: crate::otlp_http::build_blocking_client(std::time::Duration::from_secs(
|
||||
30,
|
||||
))
|
||||
.expect("test OTLP HTTP client must build"),
|
||||
resource: parking_lot::Mutex::new(opentelemetry_sdk::Resource::builder().build()),
|
||||
token_header_value: Arc::from("xai-grok-cli"),
|
||||
extra_headers: Arc::new(Vec::new()),
|
||||
}
|
||||
}
|
||||
#[test]
|
||||
fn build_export_headers_tracks_snapshot_and_respects_overrides() {
|
||||
let static_headers = std::collections::HashMap::new();
|
||||
for snapshot in [
|
||||
CredentialSnapshot::default(),
|
||||
CredentialSnapshot {
|
||||
user_id: Some(String::new()),
|
||||
team_id: Some(String::new()),
|
||||
..Default::default()
|
||||
},
|
||||
] {
|
||||
let headers = build_export_headers(&static_headers, "tok", None, &[], &snapshot);
|
||||
assert!(!headers.contains_key("x-userid"));
|
||||
assert!(!headers.contains_key("x-teamid"));
|
||||
}
|
||||
let snapshot = CredentialSnapshot {
|
||||
user_id: Some("u1".into()),
|
||||
team_id: Some("t9".into()),
|
||||
..Default::default()
|
||||
};
|
||||
let extra = vec![("Authorization".to_string(), "Bearer custom".to_string())];
|
||||
let headers = build_export_headers(&static_headers, "auto-token", None, &extra, &snapshot);
|
||||
assert_eq!(headers["x-userid"], "u1");
|
||||
assert_eq!(headers["x-teamid"], "t9");
|
||||
assert_eq!(headers["Authorization"], "Bearer custom");
|
||||
}
|
||||
#[test]
|
||||
fn refreshable_exporter_uses_updated_provider_token() {
|
||||
let provider = Arc::new(TestProvider::new(Some("token-a")));
|
||||
let exporter = make_exporter(provider.clone(), "cached-token");
|
||||
assert_eq!(exporter.current_token(), "token-a");
|
||||
provider.set(Some("token-b"));
|
||||
assert_eq!(exporter.current_token(), "token-b");
|
||||
}
|
||||
#[test]
|
||||
fn resource_injects_tenant_id_attrs() {
|
||||
use opentelemetry::Key;
|
||||
let base = opentelemetry_sdk::Resource::builder_empty()
|
||||
.with_attributes([opentelemetry::KeyValue::new("user.id", "")])
|
||||
.build();
|
||||
let plain = resource_with_tenant_id(base.clone(), &CredentialSnapshot::default());
|
||||
assert!(plain.get(&Key::from("deployment.id")).is_none());
|
||||
assert!(plain.get(&Key::from("api_key.id")).is_none());
|
||||
let snap = CredentialSnapshot {
|
||||
deployment_id: Some("dep-7b97".into()),
|
||||
..Default::default()
|
||||
};
|
||||
let r = resource_with_tenant_id(base.clone(), &snap);
|
||||
assert_eq!(
|
||||
r.get(&Key::from("deployment.id")).map(|v| v.to_string()),
|
||||
Some("dep-7b97".to_string())
|
||||
);
|
||||
assert!(
|
||||
r.get(&Key::from("user.id")).is_some(),
|
||||
"base attrs preserved"
|
||||
);
|
||||
let snap = CredentialSnapshot {
|
||||
api_key_id: Some("ak-0c2b".into()),
|
||||
..Default::default()
|
||||
};
|
||||
let r = resource_with_tenant_id(base.clone(), &snap);
|
||||
assert_eq!(
|
||||
r.get(&Key::from("api_key.id")).map(|v| v.to_string()),
|
||||
Some("ak-0c2b".to_string())
|
||||
);
|
||||
let snap = CredentialSnapshot {
|
||||
organization_id: Some("org-abc".into()),
|
||||
..Default::default()
|
||||
};
|
||||
let r = resource_with_tenant_id(base.clone(), &snap);
|
||||
assert_eq!(
|
||||
r.get(&Key::from("organization.id")).map(|v| v.to_string()),
|
||||
Some("org-abc".to_string())
|
||||
);
|
||||
let snap = CredentialSnapshot {
|
||||
user_id: Some("user-42".into()),
|
||||
..Default::default()
|
||||
};
|
||||
let r = resource_with_tenant_id(base.clone(), &snap);
|
||||
assert_eq!(
|
||||
r.get(&Key::from("user.id")).map(|v| v.to_string()),
|
||||
Some("user-42".to_string())
|
||||
);
|
||||
let snap = CredentialSnapshot {
|
||||
deployment_id: Some("dep-9".into()),
|
||||
user_id: Some(String::new()),
|
||||
organization_id: Some(String::new()),
|
||||
team_id: Some(String::new()),
|
||||
..Default::default()
|
||||
};
|
||||
let r =
|
||||
resource_with_tenant_id(opentelemetry_sdk::Resource::builder_empty().build(), &snap);
|
||||
assert!(r.get(&Key::from("user.id")).is_none());
|
||||
assert!(r.get(&Key::from("organization.id")).is_none());
|
||||
assert!(r.get(&Key::from("team.id")).is_none());
|
||||
assert_eq!(
|
||||
r.get(&Key::from("deployment.id")).map(|v| v.to_string()),
|
||||
Some("dep-9".to_string())
|
||||
);
|
||||
}
|
||||
#[test]
|
||||
fn refreshable_exporter_falls_back_to_cached_token_when_provider_empty() {
|
||||
let provider = Arc::new(TestProvider::new(Some("expired-token")));
|
||||
let exporter = make_exporter(provider.clone(), "cached-token");
|
||||
assert_eq!(exporter.current_token(), "expired-token");
|
||||
provider.set(None);
|
||||
assert_eq!(exporter.current_token(), "cached-token");
|
||||
}
|
||||
#[tokio::test]
|
||||
async fn refresh_after_unauthorized_rotates_token() {
|
||||
let provider = TestProvider::with_refresh("stale-token", "fresh-token");
|
||||
assert_eq!(provider.snapshot().token.as_deref(), Some("stale-token"));
|
||||
assert!(provider.refresh_after_unauthorized().await);
|
||||
assert_eq!(provider.refresh_count(), 1);
|
||||
assert_eq!(provider.snapshot().token.as_deref(), Some("fresh-token"));
|
||||
}
|
||||
#[tokio::test]
|
||||
async fn no_refresh_when_provider_cannot_refresh() {
|
||||
let provider = TestProvider::new(Some("only-token"));
|
||||
assert!(!provider.refresh_after_unauthorized().await);
|
||||
assert_eq!(provider.refresh_count(), 0);
|
||||
assert_eq!(provider.snapshot().token.as_deref(), Some("only-token"));
|
||||
}
|
||||
/// A provider whose `refresh_after_unauthorized` requires a Tokio
|
||||
/// runtime (calls `tokio::task::spawn_blocking`), mimicking the real
|
||||
/// `OtelAuthCredentialProvider` → `try_lock_auth_file_async` path.
|
||||
struct TokioDependentProvider {
|
||||
refresh_count: std::sync::atomic::AtomicU32,
|
||||
}
|
||||
impl HttpAuth for TokioDependentProvider {
|
||||
fn apply(
|
||||
&self,
|
||||
builder: reqwest::RequestBuilder,
|
||||
_base_url: &str,
|
||||
) -> reqwest::RequestBuilder {
|
||||
builder
|
||||
}
|
||||
}
|
||||
#[async_trait]
|
||||
impl AuthCredentialProvider for TokioDependentProvider {
|
||||
fn snapshot(&self) -> CredentialSnapshot {
|
||||
CredentialSnapshot {
|
||||
token: Some("test-token".into()),
|
||||
..Default::default()
|
||||
}
|
||||
}
|
||||
async fn refresh_after_unauthorized(&self) -> bool {
|
||||
let _ = tokio::task::spawn_blocking(|| {}).await;
|
||||
self.refresh_count
|
||||
.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
|
||||
true
|
||||
}
|
||||
}
|
||||
/// Regression test: export() must not call
|
||||
/// refresh_after_unauthorized() when driven by futures_executor on a
|
||||
/// plain std::thread (like the BatchSpanProcessor does).
|
||||
#[test]
|
||||
fn export_skips_refresh_without_tokio_runtime() {
|
||||
let provider = Arc::new(TokioDependentProvider {
|
||||
refresh_count: std::sync::atomic::AtomicU32::new(0),
|
||||
});
|
||||
let credentials: Arc<dyn AuthCredentialProvider> = provider.clone();
|
||||
let refresh_was_called = std::thread::spawn(move || {
|
||||
futures_executor::block_on(async {
|
||||
if tokio::runtime::Handle::try_current().is_err() {
|
||||
return false;
|
||||
}
|
||||
credentials.refresh_after_unauthorized().await
|
||||
})
|
||||
})
|
||||
.join()
|
||||
.expect("thread must not panic");
|
||||
assert!(!refresh_was_called, "refresh must be skipped without Tokio");
|
||||
assert_eq!(
|
||||
provider
|
||||
.refresh_count
|
||||
.load(std::sync::atomic::Ordering::Relaxed),
|
||||
0,
|
||||
"refresh_after_unauthorized must not be called"
|
||||
);
|
||||
}
|
||||
/// Verify refresh still works when a Tokio runtime IS present.
|
||||
#[tokio::test]
|
||||
async fn export_retries_refresh_with_tokio_runtime() {
|
||||
let provider = Arc::new(TokioDependentProvider {
|
||||
refresh_count: std::sync::atomic::AtomicU32::new(0),
|
||||
});
|
||||
let credentials: Arc<dyn AuthCredentialProvider> = provider.clone();
|
||||
let should_retry = async {
|
||||
if tokio::runtime::Handle::try_current().is_err() {
|
||||
return false;
|
||||
}
|
||||
credentials.refresh_after_unauthorized().await
|
||||
}
|
||||
.await;
|
||||
assert!(should_retry, "refresh should proceed with Tokio runtime");
|
||||
assert_eq!(
|
||||
provider
|
||||
.refresh_count
|
||||
.load(std::sync::atomic::Ordering::Relaxed),
|
||||
1
|
||||
);
|
||||
}
|
||||
}
|
||||
593
crates/codegen/xai-grok-telemetry/src/otel_layer/redact.rs
Normal file
593
crates/codegen/xai-grok-telemetry/src/otel_layer/redact.rs
Normal file
|
|
@ -0,0 +1,593 @@
|
|||
use std::borrow::Cow;
|
||||
use std::collections::HashSet;
|
||||
use std::sync::LazyLock;
|
||||
|
||||
use opentelemetry::trace::{Event, Status};
|
||||
use opentelemetry::{Array, KeyValue, StringValue, Value};
|
||||
use opentelemetry_sdk::trace::SpanData;
|
||||
|
||||
/// Adding a span attribute (default-deny via `enforce_allowlist`): record
|
||||
/// numerics as `i64` (`u64` serializes as a string and is dropped); derive
|
||||
/// label values from an enum `as_str()`; add string keys here and to the
|
||||
/// round-trip test pin, and only if they carry no user content.
|
||||
pub(super) static ALLOWED_STRING_KEYS: &[&str] = &[
|
||||
// tracing-opentelemetry / framework-injected
|
||||
"level",
|
||||
"target",
|
||||
"code.namespace",
|
||||
"code.filepath",
|
||||
"thread.name",
|
||||
// identifiers
|
||||
"session_id",
|
||||
"prompt_id",
|
||||
"req_id",
|
||||
"request_id",
|
||||
"child_session_id",
|
||||
"parent_session_id",
|
||||
"subagent_id",
|
||||
"agent_id",
|
||||
"task_id",
|
||||
"tool_call_id",
|
||||
"call_id",
|
||||
"event_id",
|
||||
"conv_id",
|
||||
"turn_id",
|
||||
// model / client
|
||||
"model_id",
|
||||
"model",
|
||||
"compact_model",
|
||||
"client_type",
|
||||
"client_version",
|
||||
"subagent_type",
|
||||
"persona",
|
||||
"role",
|
||||
// tool / skill / mcp / method NAMES (identifiers, not arguments)
|
||||
"skill_name",
|
||||
"server_name",
|
||||
"tool_name",
|
||||
"tool_names",
|
||||
"method",
|
||||
"operation",
|
||||
"endpoint",
|
||||
// paths / urls (additionally home-path- and url-scrubbed by redact_value)
|
||||
"path",
|
||||
"file_path",
|
||||
"repo_path",
|
||||
"gcs_path",
|
||||
"gcs_url",
|
||||
"url",
|
||||
"output_path",
|
||||
"dir",
|
||||
"dir_path",
|
||||
"notebook",
|
||||
"cwd",
|
||||
"original_cwd",
|
||||
"chosen_repo_root",
|
||||
"worktree",
|
||||
"source",
|
||||
"bucket_url",
|
||||
"object_path",
|
||||
"archive_name",
|
||||
"artifact",
|
||||
// enums / classifications
|
||||
"verdict",
|
||||
"pattern_class",
|
||||
"phase",
|
||||
"upload_reason",
|
||||
"suppress_reason",
|
||||
"error_kind",
|
||||
"error_category",
|
||||
"error_type",
|
||||
"outcome",
|
||||
"decision",
|
||||
"update_type",
|
||||
"kind",
|
||||
"step",
|
||||
"token_type",
|
||||
"stop_reason",
|
||||
"compaction_outcome",
|
||||
"compaction_stop_reason",
|
||||
"compaction_trigger",
|
||||
"compaction_prefire_outcome",
|
||||
"aspect_ratio",
|
||||
"resolution",
|
||||
"schedule",
|
||||
"interval",
|
||||
"mode",
|
||||
"detail",
|
||||
// span enums + plugin/auth/survey/mcp identifiers (categorical, no user content)
|
||||
"status",
|
||||
"action",
|
||||
"auth_method",
|
||||
// auth 401 attribution: fixed enum-ish consumer labels only
|
||||
// (e.g. "OaiCompatClient.chat_completions_stream"); never user content.
|
||||
// Key suffix fields stay denied — they are token fingerprints.
|
||||
"consumer",
|
||||
"to_mode",
|
||||
"trigger",
|
||||
"survey_type",
|
||||
"mention_type",
|
||||
"install_kind",
|
||||
"transport_type",
|
||||
"invocation_trigger",
|
||||
"skill_source",
|
||||
"plugin_name",
|
||||
"plugin_version",
|
||||
"plugin_scope",
|
||||
"hook_event",
|
||||
"hook_name",
|
||||
"hook_type",
|
||||
"hook_source",
|
||||
"server_scope",
|
||||
"mcp_server.name",
|
||||
"mcp_tool.name",
|
||||
"agent.name",
|
||||
"skill.name",
|
||||
"query_source",
|
||||
"effort",
|
||||
"start_type",
|
||||
"error",
|
||||
"location",
|
||||
"user_id",
|
||||
"parent_agent_id",
|
||||
"from_mode",
|
||||
"tool_use_id",
|
||||
"command_name",
|
||||
"command_source",
|
||||
"event_type",
|
||||
"appearance_id",
|
||||
// terminal telemetry
|
||||
"terminal.brand",
|
||||
"terminal.multiplexer",
|
||||
"terminal.tmux_version",
|
||||
"terminal.term_var",
|
||||
"skip_reason",
|
||||
"auto_cadence_reason",
|
||||
];
|
||||
|
||||
/// O(1) lookup view over [`ALLOWED_STRING_KEYS`].
|
||||
static ALLOWED_STRING_KEY_SET: LazyLock<HashSet<&'static str>> =
|
||||
LazyLock::new(|| ALLOWED_STRING_KEYS.iter().copied().collect());
|
||||
|
||||
/// Allowlisted keys holding full URLs: reduced to `scheme://host[:port]` so
|
||||
/// user-influenced path/query can't export. Storage-*path* keys (`gcs_path`,
|
||||
/// `object_path`, `output_path`) are excluded — those paths are wanted.
|
||||
static URL_VALUED_KEYS: &[&str] = &["url", "endpoint", "gcs_url", "bucket_url"];
|
||||
|
||||
/// Scrub every text-bearing surface of each span before export.
|
||||
pub(super) fn redact_batch(batch: &mut [SpanData]) {
|
||||
for span in batch.iter_mut() {
|
||||
// Exhaustive destructure (no `..`): a new `SpanData` field in a future
|
||||
// `opentelemetry_sdk` fails to compile here instead of exporting unscrubbed.
|
||||
let SpanData {
|
||||
name,
|
||||
attributes,
|
||||
events,
|
||||
links,
|
||||
status,
|
||||
span_context: _,
|
||||
parent_span_id: _,
|
||||
parent_span_is_remote: _,
|
||||
span_kind: _,
|
||||
start_time: _,
|
||||
end_time: _,
|
||||
dropped_attributes_count: _,
|
||||
instrumentation_scope: _,
|
||||
} = span;
|
||||
redact_in_place(name);
|
||||
scrub_attributes(attributes);
|
||||
for event in &mut events.events {
|
||||
neuter_event_name(event);
|
||||
// Re-scrub: synthesized callsite paths can be absolute (home dir).
|
||||
redact_in_place(&mut event.name);
|
||||
scrub_attributes(&mut event.attributes);
|
||||
}
|
||||
// Keep the error message (useful telemetry); scrub secrets/paths from it.
|
||||
if let Status::Error { description } = status {
|
||||
redact_in_place(description);
|
||||
}
|
||||
for link in &mut links.links {
|
||||
scrub_attributes(&mut link.attributes);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Numeric/bool scalars and their arrays are content-free; everything else —
|
||||
/// strings and any future `#[non_exhaustive]` variant — is content (fail-closed).
|
||||
fn is_content_value(value: &Value) -> bool {
|
||||
!matches!(
|
||||
value,
|
||||
Value::Bool(_)
|
||||
| Value::I64(_)
|
||||
| Value::F64(_)
|
||||
| Value::Array(Array::Bool(_) | Array::I64(_) | Array::F64(_))
|
||||
)
|
||||
}
|
||||
|
||||
/// Default-deny: drop content-valued attributes whose key isn't allowlisted.
|
||||
fn enforce_allowlist(attrs: &mut Vec<KeyValue>) {
|
||||
attrs.retain(|kv| {
|
||||
!is_content_value(&kv.value) || ALLOWED_STRING_KEY_SET.contains(kv.key.as_str())
|
||||
});
|
||||
}
|
||||
|
||||
fn scrub_attributes(attrs: &mut Vec<KeyValue>) {
|
||||
enforce_allowlist(attrs);
|
||||
for kv in attrs.iter_mut() {
|
||||
if URL_VALUED_KEYS.contains(&kv.key.as_str()) {
|
||||
reduce_url_to_origin(&mut kv.value);
|
||||
}
|
||||
redact_value(&mut kv.value);
|
||||
}
|
||||
}
|
||||
|
||||
/// An event's name is the formatted `tracing` message (`Event.name`) — free
|
||||
/// text the key allowlist can't gate, so replace it with the static callsite id
|
||||
/// (fail-closed). Rebuilt from the `code.filepath`/`code.lineno` attrs that
|
||||
/// `tracing-opentelemetry` attaches to every event (`with_location`, default-on).
|
||||
fn neuter_event_name(event: &mut Event) {
|
||||
let mut file: Option<String> = None;
|
||||
let mut line: Option<i64> = None;
|
||||
for kv in &event.attributes {
|
||||
match kv.key.as_str() {
|
||||
"code.filepath" => {
|
||||
if let Value::String(s) = &kv.value {
|
||||
file = Some(s.as_str().to_owned());
|
||||
}
|
||||
}
|
||||
"code.lineno" => {
|
||||
if let Value::I64(n) = &kv.value {
|
||||
line = Some(*n);
|
||||
}
|
||||
}
|
||||
_ => {}
|
||||
}
|
||||
}
|
||||
event.name = match (file, line) {
|
||||
(Some(f), Some(l)) => format!("{f}:{l}").into(),
|
||||
(Some(f), None) => f.into(),
|
||||
// No location attrs (e.g. a raw-API event): drop the message entirely.
|
||||
_ => Cow::Borrowed("event"),
|
||||
};
|
||||
}
|
||||
|
||||
/// Reduce a URL to `scheme://host[:port]` — its path/query can carry user
|
||||
/// content. Unparseable values pass through to the secret scrubber.
|
||||
fn reduce_url_to_origin(value: &mut Value) {
|
||||
if let Value::String(s) = value
|
||||
&& let Cow::Owned(origin) = crate::redact_common::url_origin(s.as_str())
|
||||
{
|
||||
*s = StringValue::from(origin);
|
||||
}
|
||||
}
|
||||
|
||||
/// Secret-shape then user-path scrub (shared with the external pipeline).
|
||||
/// Returns `Some` only when the input changed (owned, so callers can
|
||||
/// overwrite in place).
|
||||
fn redact_owned(input: &str) -> Option<String> {
|
||||
crate::redact_common::redact_owned(input)
|
||||
}
|
||||
|
||||
fn redact_in_place(s: &mut Cow<'static, str>) {
|
||||
if let Some(redacted) = redact_owned(s.as_ref()) {
|
||||
*s = Cow::Owned(redacted);
|
||||
}
|
||||
}
|
||||
|
||||
fn redact_value(value: &mut Value) {
|
||||
match value {
|
||||
Value::String(s) => {
|
||||
if let Some(redacted) = redact_owned(s.as_str()) {
|
||||
*s = StringValue::from(redacted);
|
||||
}
|
||||
}
|
||||
Value::Array(Array::String(items)) => {
|
||||
for s in items.iter_mut() {
|
||||
if let Some(redacted) = redact_owned(s.as_str()) {
|
||||
*s = StringValue::from(redacted);
|
||||
}
|
||||
}
|
||||
}
|
||||
// Non-string variants carry no free text; `Value` is `#[non_exhaustive]`.
|
||||
_ => {}
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
#[test]
|
||||
fn redact_value_scrubs_secret_string() {
|
||||
let mut v = Value::String(StringValue::from(
|
||||
"Authorization: Bearer eyJhbGciOiJIUzI1NiJ9.foo.bar.baz".to_string(),
|
||||
));
|
||||
redact_value(&mut v);
|
||||
let Value::String(s) = &v else {
|
||||
panic!("expected string value");
|
||||
};
|
||||
assert!(
|
||||
s.as_str().contains("[REDACTED_SECRET]"),
|
||||
"secret not scrubbed: {}",
|
||||
s.as_str()
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn event_name_neutered_to_callsite_drops_message_content() {
|
||||
let mut ev = Event::new(
|
||||
"received prompt: rm -rf /Users/alice/secret",
|
||||
std::time::SystemTime::now(),
|
||||
vec![
|
||||
KeyValue::new("code.filepath", "src/foo.rs"),
|
||||
KeyValue::new("code.lineno", 42_i64),
|
||||
],
|
||||
0,
|
||||
);
|
||||
neuter_event_name(&mut ev);
|
||||
assert_eq!(ev.name, "src/foo.rs:42");
|
||||
assert!(!ev.name.contains("prompt") && !ev.name.contains("rm -rf"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn event_name_without_location_drops_to_marker() {
|
||||
let mut ev = Event::new("SECRET {x:?}", std::time::SystemTime::now(), vec![], 0);
|
||||
neuter_event_name(&mut ev);
|
||||
assert_eq!(ev.name, "event");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn allowlist_drops_nonallowlisted_content_keeps_safe_and_numeric() {
|
||||
let mut attrs = vec![
|
||||
KeyValue::new("session_id", "sess-abc"), // allowlisted string
|
||||
KeyValue::new("path", "/tmp/x.rs"), // allowlisted string
|
||||
KeyValue::new("prompt", "CANARY_PROMPT secret user text"), // not allowlisted → drop
|
||||
KeyValue::new("command", "echo CANARY_SECRET"), // not allowlisted → drop
|
||||
KeyValue::new("turn_number", 7_i64), // numeric → keep
|
||||
KeyValue::new("is_background", true), // bool → keep
|
||||
];
|
||||
enforce_allowlist(&mut attrs);
|
||||
let keys: Vec<&str> = attrs.iter().map(|kv| kv.key.as_str()).collect();
|
||||
assert!(keys.contains(&"session_id"));
|
||||
assert!(keys.contains(&"path"));
|
||||
assert!(keys.contains(&"turn_number"));
|
||||
assert!(keys.contains(&"is_background"));
|
||||
assert!(
|
||||
!keys.contains(&"prompt"),
|
||||
"non-allowlisted content must be dropped"
|
||||
);
|
||||
assert!(
|
||||
!keys.contains(&"command"),
|
||||
"non-allowlisted content must be dropped"
|
||||
);
|
||||
// Canary: no dropped content survives anywhere in the attribute set.
|
||||
let blob = format!("{attrs:?}");
|
||||
assert!(
|
||||
!blob.contains("CANARY_PROMPT"),
|
||||
"prompt content leaked: {blob}"
|
||||
);
|
||||
assert!(
|
||||
!blob.contains("CANARY_SECRET"),
|
||||
"command content leaked: {blob}"
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn allowlist_contents_are_pinned() {
|
||||
// Keep this an independent copy — don't reference ALLOWED_STRING_KEYS, or
|
||||
// the assert becomes a tautology and stops gating allowlist changes.
|
||||
let expected: &[&str] = &[
|
||||
"level",
|
||||
"target",
|
||||
"code.namespace",
|
||||
"code.filepath",
|
||||
"thread.name",
|
||||
"session_id",
|
||||
"prompt_id",
|
||||
"req_id",
|
||||
"request_id",
|
||||
"child_session_id",
|
||||
"parent_session_id",
|
||||
"subagent_id",
|
||||
"agent_id",
|
||||
"task_id",
|
||||
"tool_call_id",
|
||||
"call_id",
|
||||
"event_id",
|
||||
"conv_id",
|
||||
"turn_id",
|
||||
"model_id",
|
||||
"model",
|
||||
"compact_model",
|
||||
"client_type",
|
||||
"client_version",
|
||||
"subagent_type",
|
||||
"persona",
|
||||
"role",
|
||||
"skill_name",
|
||||
"server_name",
|
||||
"tool_name",
|
||||
"tool_names",
|
||||
"method",
|
||||
"operation",
|
||||
"endpoint",
|
||||
"path",
|
||||
"file_path",
|
||||
"repo_path",
|
||||
"gcs_path",
|
||||
"gcs_url",
|
||||
"url",
|
||||
"output_path",
|
||||
"dir",
|
||||
"dir_path",
|
||||
"notebook",
|
||||
"cwd",
|
||||
"original_cwd",
|
||||
"chosen_repo_root",
|
||||
"worktree",
|
||||
"source",
|
||||
"bucket_url",
|
||||
"object_path",
|
||||
"archive_name",
|
||||
"artifact",
|
||||
"verdict",
|
||||
"pattern_class",
|
||||
"phase",
|
||||
"upload_reason",
|
||||
"suppress_reason",
|
||||
"error_kind",
|
||||
"error_category",
|
||||
"error_type",
|
||||
"outcome",
|
||||
"decision",
|
||||
"update_type",
|
||||
"kind",
|
||||
"step",
|
||||
"token_type",
|
||||
"stop_reason",
|
||||
"compaction_outcome",
|
||||
"compaction_stop_reason",
|
||||
"compaction_trigger",
|
||||
"compaction_prefire_outcome",
|
||||
"aspect_ratio",
|
||||
"resolution",
|
||||
"schedule",
|
||||
"interval",
|
||||
"mode",
|
||||
"detail",
|
||||
"status",
|
||||
"action",
|
||||
"auth_method",
|
||||
"consumer",
|
||||
"to_mode",
|
||||
"trigger",
|
||||
"survey_type",
|
||||
"mention_type",
|
||||
"install_kind",
|
||||
"transport_type",
|
||||
"invocation_trigger",
|
||||
"skill_source",
|
||||
"plugin_name",
|
||||
"plugin_version",
|
||||
"plugin_scope",
|
||||
"hook_event",
|
||||
"hook_name",
|
||||
"hook_type",
|
||||
"hook_source",
|
||||
"server_scope",
|
||||
"mcp_server.name",
|
||||
"mcp_tool.name",
|
||||
"agent.name",
|
||||
"skill.name",
|
||||
"query_source",
|
||||
"effort",
|
||||
"start_type",
|
||||
"error",
|
||||
"location",
|
||||
"user_id",
|
||||
"parent_agent_id",
|
||||
"from_mode",
|
||||
"tool_use_id",
|
||||
"command_name",
|
||||
"command_source",
|
||||
"event_type",
|
||||
"appearance_id",
|
||||
"terminal.brand",
|
||||
"terminal.multiplexer",
|
||||
"terminal.tmux_version",
|
||||
"terminal.term_var",
|
||||
"skip_reason",
|
||||
"auto_cadence_reason",
|
||||
];
|
||||
assert_eq!(
|
||||
ALLOWED_STRING_KEYS, expected,
|
||||
"ALLOWED_STRING_KEYS changed: adding a key exports a new field — confirm it carries no \
|
||||
user content, then update this pin."
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn error_status_message_retained_but_secret_scrubbed() {
|
||||
// Error messages are useful telemetry and must survive; only secret
|
||||
// shapes (and home/username paths) are scrubbed out of them.
|
||||
let mut status = Status::error("upstream auth failed: sk-CANARYabcdefghij1234567890");
|
||||
if let Status::Error { description } = &mut status {
|
||||
redact_in_place(description);
|
||||
}
|
||||
let Status::Error { description } = status else {
|
||||
panic!("status code must stay Error");
|
||||
};
|
||||
assert!(
|
||||
description.contains("upstream auth failed"),
|
||||
"useful message lost: {description}"
|
||||
);
|
||||
assert!(
|
||||
!description.contains("CANARY"),
|
||||
"secret survived: {description}"
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn url_value_reduced_to_origin_dropping_path_and_query() {
|
||||
let mut attrs = vec![KeyValue::new(
|
||||
"url",
|
||||
"https://example.com:8443/search?q=CANARY+secret+terms&u=bob#frag",
|
||||
)];
|
||||
scrub_attributes(&mut attrs);
|
||||
let blob = format!("{attrs:?}");
|
||||
assert!(
|
||||
blob.contains("https://example.com:8443"),
|
||||
"origin lost: {blob}"
|
||||
);
|
||||
assert!(!blob.contains("CANARY"), "query content survived: {blob}");
|
||||
assert!(!blob.contains("search"), "path survived: {blob}");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn url_valued_keys_reduced_to_origin_but_storage_paths_kept() {
|
||||
// Origin-reduction applies to every URL-valued key, not just `url`...
|
||||
let mut attrs = vec![
|
||||
KeyValue::new(
|
||||
"bucket_url",
|
||||
"https://store.example.com/b/CANARY/o?sig=CANARYSIG",
|
||||
),
|
||||
KeyValue::new(
|
||||
"endpoint",
|
||||
"https://api.example.com:8443/v1/chat?u=CANARYUSER",
|
||||
),
|
||||
// ...but storage *paths* are deliberately exported in full.
|
||||
KeyValue::new("gcs_path", "sessions/abc123/artifact-kept.tar"),
|
||||
];
|
||||
scrub_attributes(&mut attrs);
|
||||
let blob = format!("{attrs:?}");
|
||||
assert!(
|
||||
blob.contains("https://store.example.com"),
|
||||
"bucket_url origin lost: {blob}"
|
||||
);
|
||||
assert!(
|
||||
blob.contains("https://api.example.com:8443"),
|
||||
"endpoint origin lost: {blob}"
|
||||
);
|
||||
assert!(
|
||||
!blob.contains("CANARY"),
|
||||
"url path/query content survived: {blob}"
|
||||
);
|
||||
assert!(
|
||||
blob.contains("sessions/abc123/artifact-kept.tar"),
|
||||
"storage path was wrongly reduced: {blob}"
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn allowlisted_value_is_still_secret_scrubbed() {
|
||||
// Allowlisting a key permits the field; it does not exempt the value
|
||||
// from the shape scrub.
|
||||
let mut attrs = vec![KeyValue::new("source", "sk-CANARYabcdefghij1234567890")];
|
||||
scrub_attributes(&mut attrs);
|
||||
let blob = format!("{attrs:?}");
|
||||
assert!(
|
||||
!blob.contains("CANARY"),
|
||||
"secret in allowlisted value not scrubbed: {blob}"
|
||||
);
|
||||
}
|
||||
}
|
||||
77
crates/codegen/xai-grok-telemetry/src/otlp_http.rs
Normal file
77
crates/codegen/xai-grok-telemetry/src/otlp_http.rs
Normal file
|
|
@ -0,0 +1,77 @@
|
|||
//! Shared construction of the blocking `reqwest` client used by the OTLP
|
||||
//! HTTP exporters (spans in `otel_layer`, logs/metrics in `external`).
|
||||
//!
|
||||
//! This uses the workspace `reqwest` 0.12 (`rustls-tls`, embedded webpki
|
||||
//! roots) rather than reqwest 0.13. reqwest 0.13's blocking client runs its
|
||||
//! rustls/aws-lc-rs handshake on the fixed-stack, un-sizable
|
||||
//! `reqwest-internal-sync-runtime` thread; that handshake overflows the stack
|
||||
//! on the first OTLP export and crashes the CLI a few seconds after launch
|
||||
//! (observed on Windows arm64; `RUST_MIN_STACK` does not help because reqwest
|
||||
//! owns that thread). reqwest 0.12 shares the known-good TLS stack the rest of
|
||||
//! the CLI already uses, and its embedded roots keep the exporter working on
|
||||
//! hosts with no system CA store. `opentelemetry-http` only ships an
|
||||
//! `HttpClient` impl for its pinned reqwest 0.13, so the 0.12 client is wrapped
|
||||
//! below (orphan rule). Construction returns an error for callers to degrade on
|
||||
//! (disable the exporter, keep the session alive) instead of panicking.
|
||||
|
||||
use async_trait::async_trait;
|
||||
use bytes::Bytes;
|
||||
use opentelemetry_http::{HttpClient, HttpError};
|
||||
|
||||
/// `opentelemetry_http::HttpClient` over the workspace reqwest 0.12 blocking
|
||||
/// client. Mirrors `opentelemetry-http`'s built-in reqwest 0.13 blocking impl.
|
||||
#[derive(Debug, Clone)]
|
||||
pub(crate) struct BlockingOtlpClient(reqwest::blocking::Client);
|
||||
|
||||
#[async_trait]
|
||||
impl HttpClient for BlockingOtlpClient {
|
||||
async fn send_bytes(
|
||||
&self,
|
||||
request: http::Request<Bytes>,
|
||||
) -> Result<http::Response<Bytes>, HttpError> {
|
||||
let request = request.try_into()?;
|
||||
let mut response = self.0.execute(request)?.error_for_status()?;
|
||||
let headers = std::mem::take(response.headers_mut());
|
||||
let mut http_response = http::Response::builder()
|
||||
.status(response.status())
|
||||
.body(response.bytes()?)?;
|
||||
*http_response.headers_mut() = headers;
|
||||
Ok(http_response)
|
||||
}
|
||||
}
|
||||
|
||||
/// Build the blocking OTLP HTTP client on a dedicated thread.
|
||||
///
|
||||
/// The blocking client can't be built inside a Tokio runtime, and the batch
|
||||
/// processors drive exports from non-Tokio threads — building on a fresh
|
||||
/// thread avoids the "no reactor" panic for every caller.
|
||||
pub(crate) fn build_blocking_client(
|
||||
timeout: std::time::Duration,
|
||||
) -> Result<BlockingOtlpClient, String> {
|
||||
std::thread::Builder::new()
|
||||
.name("otlp-client-build".into())
|
||||
.spawn(move || {
|
||||
reqwest::blocking::Client::builder()
|
||||
.timeout(timeout)
|
||||
.build()
|
||||
.map(BlockingOtlpClient)
|
||||
.map_err(|e| format!("building blocking OTLP HTTP client: {e}"))
|
||||
})
|
||||
.map_err(|e| format!("spawning OTLP client builder thread: {e}"))?
|
||||
.join()
|
||||
.map_err(|_| "OTLP client builder thread panicked".to_string())?
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
/// The client must build without consulting the system CA store — reqwest
|
||||
/// 0.12 `rustls-tls` trusts embedded webpki roots, so this holds on hosts
|
||||
/// with no system CA store.
|
||||
#[test]
|
||||
fn blocking_otlp_client_builds_with_embedded_roots() {
|
||||
build_blocking_client(std::time::Duration::from_secs(5))
|
||||
.expect("client with embedded webpki roots must build on any host");
|
||||
}
|
||||
}
|
||||
57
crates/codegen/xai-grok-telemetry/src/prompt_timing.rs
Normal file
57
crates/codegen/xai-grok-telemetry/src/prompt_timing.rs
Normal file
|
|
@ -0,0 +1,57 @@
|
|||
//! Per-turn prompt latency measurement.
|
||||
//!
|
||||
//! Extracted from `xai-grok-shell::session::prompt_timing`.
|
||||
|
||||
use std::time::Instant;
|
||||
|
||||
use crate::events::PromptLatency;
|
||||
use crate::session_ctx::log_event;
|
||||
|
||||
pub use crate::enums::McpInitStrategy;
|
||||
|
||||
pub struct PromptTiming {
|
||||
turn_start: Instant,
|
||||
mcp_wait_ms: u64,
|
||||
tool_collection_ms: u64,
|
||||
}
|
||||
|
||||
impl PromptTiming {
|
||||
pub fn start() -> Self {
|
||||
Self {
|
||||
turn_start: Instant::now(),
|
||||
mcp_wait_ms: 0,
|
||||
tool_collection_ms: 0,
|
||||
}
|
||||
}
|
||||
|
||||
pub fn record_tool_prep(&mut self, mcp_wait_ms: u64, total_prep_ms: u64) {
|
||||
self.mcp_wait_ms = mcp_wait_ms;
|
||||
self.tool_collection_ms = total_prep_ms.saturating_sub(mcp_wait_ms);
|
||||
}
|
||||
|
||||
pub fn emit(
|
||||
self,
|
||||
model_call_ms: u64,
|
||||
turn_index: u32,
|
||||
mcp_server_count: u32,
|
||||
mcp_tools_registered: u32,
|
||||
mcp_strategy: McpInitStrategy,
|
||||
model_id: String,
|
||||
) {
|
||||
let total_ms = self.turn_start.elapsed().as_millis() as u64;
|
||||
let pre_model_ms = total_ms.saturating_sub(model_call_ms);
|
||||
|
||||
log_event(PromptLatency {
|
||||
turn_index,
|
||||
total_ms,
|
||||
mcp_wait_ms: self.mcp_wait_ms,
|
||||
tool_collection_ms: self.tool_collection_ms,
|
||||
model_call_ms,
|
||||
pre_model_ms,
|
||||
mcp_server_count,
|
||||
mcp_tools_registered,
|
||||
mcp_strategy,
|
||||
model_id,
|
||||
});
|
||||
}
|
||||
}
|
||||
72
crates/codegen/xai-grok-telemetry/src/redact_common.rs
Normal file
72
crates/codegen/xai-grok-telemetry/src/redact_common.rs
Normal file
|
|
@ -0,0 +1,72 @@
|
|||
//! Redaction helpers shared by the **internal** OTLP span pipeline
|
||||
//! ([`crate::otel_layer`]) and the **external** customer-collector pipeline
|
||||
//! ([`crate::external`]).
|
||||
//!
|
||||
//! Both pipelines are authoritative privacy chokepoints (see the crate
|
||||
//! `AGENTS.md`); these helpers are the string-level scrubbing primitives they
|
||||
//! share. Changes here affect every byte that leaves the process on either
|
||||
//! pipeline.
|
||||
|
||||
use std::borrow::Cow;
|
||||
|
||||
/// Secret-shape then user-path scrub. Returns `Some` only when the input
|
||||
/// changed (owned, so callers can overwrite in place).
|
||||
pub(crate) fn redact_owned(input: &str) -> Option<String> {
|
||||
let secrets = xai_grok_secrets::redact_secrets(input);
|
||||
match xai_grok_secrets::redact_user_paths(secrets.as_ref()) {
|
||||
Cow::Owned(paths) => Some(paths),
|
||||
Cow::Borrowed(_) => match secrets {
|
||||
Cow::Owned(s) => Some(s),
|
||||
Cow::Borrowed(_) => None,
|
||||
},
|
||||
}
|
||||
}
|
||||
|
||||
/// Scrub a string, returning the (possibly unchanged) owned value.
|
||||
pub(crate) fn redact_to_owned(input: &str) -> String {
|
||||
redact_owned(input).unwrap_or_else(|| input.to_owned())
|
||||
}
|
||||
|
||||
/// Reduce a URL to `scheme://host[:port]` — its path/query can carry user
|
||||
/// content. Unparseable values are returned unchanged (callers pass the result
|
||||
/// through the secret scrubber).
|
||||
pub(crate) fn url_origin(value: &str) -> Cow<'_, str> {
|
||||
if let Ok(url) = url::Url::parse(value)
|
||||
&& let Some(host) = url.host_str()
|
||||
{
|
||||
let origin = match url.port() {
|
||||
Some(port) => format!("{}://{}:{}", url.scheme(), host, port),
|
||||
None => format!("{}://{}", url.scheme(), host),
|
||||
};
|
||||
return Cow::Owned(origin);
|
||||
}
|
||||
Cow::Borrowed(value)
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
#[test]
|
||||
fn redact_owned_scrubs_secret_shapes() {
|
||||
let out = redact_owned("key sk-CANARYabcdefghij1234567890 end")
|
||||
.expect("secret must trigger a rewrite");
|
||||
assert!(!out.contains("CANARY"), "secret survived: {out}");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn redact_owned_returns_none_when_clean() {
|
||||
assert_eq!(redact_owned("no secrets here"), None);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn url_origin_drops_path_and_query() {
|
||||
let origin = url_origin("https://collector.corp.example:4318/v1/logs?token=CANARY");
|
||||
assert_eq!(origin, "https://collector.corp.example:4318");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn url_origin_passes_unparseable_through() {
|
||||
assert_eq!(url_origin("not a url"), "not a url");
|
||||
}
|
||||
}
|
||||
70
crates/codegen/xai-grok-telemetry/src/sampling_log.rs
Normal file
70
crates/codegen/xai-grok-telemetry/src/sampling_log.rs
Normal file
|
|
@ -0,0 +1,70 @@
|
|||
//! Tracing layer for `target: "sampling_log"` → `~/.grok/logs/sampling.jsonl`.
|
||||
//! Enable with `--log-sampling` or `GROK_LOG_SAMPLING=1`.
|
||||
|
||||
use std::sync::Mutex;
|
||||
|
||||
use tracing::Subscriber;
|
||||
use tracing_subscriber::fmt::writer::BoxMakeWriter;
|
||||
use tracing_subscriber::layer::Layer;
|
||||
use tracing_subscriber::registry::LookupSpan;
|
||||
|
||||
use xai_grok_config::grok_home;
|
||||
|
||||
use crate::instrumentation::{NoOpLayer, TargetFilterLayer};
|
||||
|
||||
const ENV_VAR: &str = "GROK_LOG_SAMPLING";
|
||||
const LOG_FILE: &str = "sampling.jsonl";
|
||||
const TARGET: &str = "sampling_log";
|
||||
|
||||
static GUARD: std::sync::OnceLock<Mutex<Option<tracing_appender::non_blocking::WorkerGuard>>> =
|
||||
std::sync::OnceLock::new();
|
||||
|
||||
pub fn layer<S>() -> Box<dyn Layer<S> + Send + Sync>
|
||||
where
|
||||
S: Subscriber + for<'span> LookupSpan<'span> + Send + Sync + 'static,
|
||||
{
|
||||
if !std::env::var(ENV_VAR).is_ok_and(|v| matches!(v.as_str(), "1" | "true" | "on")) {
|
||||
return Box::new(NoOpLayer::new());
|
||||
}
|
||||
|
||||
let path = grok_home().join(crate::unified_log::LOG_DIR).join(LOG_FILE);
|
||||
|
||||
if let Some(parent) = path.parent()
|
||||
&& let Err(e) = std::fs::create_dir_all(parent)
|
||||
{
|
||||
tracing::warn!("failed to create sampling log dir: {e}");
|
||||
return Box::new(NoOpLayer::new());
|
||||
}
|
||||
|
||||
if crate::unified_log::file_size(&path) >= crate::unified_log::MAX_SIZE {
|
||||
crate::unified_log::trim_file(&path);
|
||||
}
|
||||
|
||||
let file = match std::fs::OpenOptions::new()
|
||||
.create(true)
|
||||
.append(true)
|
||||
.open(&path)
|
||||
{
|
||||
Ok(f) => f,
|
||||
Err(e) => {
|
||||
tracing::warn!("failed to open sampling log: {e}");
|
||||
return Box::new(NoOpLayer::new());
|
||||
}
|
||||
};
|
||||
|
||||
let (non_blocking, guard) = tracing_appender::non_blocking(file);
|
||||
let guard_slot = GUARD.get_or_init(|| Mutex::new(None));
|
||||
if let Ok(mut slot) = guard_slot.lock() {
|
||||
*slot = Some(guard);
|
||||
}
|
||||
|
||||
let fmt_layer = tracing_subscriber::fmt::layer()
|
||||
.json()
|
||||
.with_current_span(false) // `spans` array already carries the full ancestor list
|
||||
.with_ansi(false)
|
||||
.with_timer(tracing_subscriber::fmt::time::UtcTime::rfc_3339())
|
||||
.with_target(false)
|
||||
.with_writer(BoxMakeWriter::new(non_blocking));
|
||||
|
||||
Box::new(TargetFilterLayer::new(fmt_layer, TARGET))
|
||||
}
|
||||
427
crates/codegen/xai-grok-telemetry/src/sentry.rs
Normal file
427
crates/codegen/xai-grok-telemetry/src/sentry.rs
Normal file
|
|
@ -0,0 +1,427 @@
|
|||
use std::borrow::Cow;
|
||||
use std::sync::Arc;
|
||||
use std::sync::OnceLock;
|
||||
use std::time::Duration;
|
||||
|
||||
use sentry::ClientInitGuard;
|
||||
use sentry::ClientOptions;
|
||||
use sentry::protocol::{Event, Value};
|
||||
|
||||
const TRACES_SAMPLE_RATE: f32 = 0.01;
|
||||
const FLUSH_TIMEOUT: Duration = Duration::from_secs(2);
|
||||
|
||||
// ─── Host integration ─────────────────────────────────────────────────────
|
||||
|
||||
/// Per-host config; everything that varies between binaries lives here.
|
||||
pub struct Config {
|
||||
/// Sentry tag `client`, e.g. `"grok-pager"`.
|
||||
pub client: &'static str,
|
||||
pub client_version: &'static str,
|
||||
pub release: &'static str,
|
||||
/// When `true`, [`init`] returns a no-op guard regardless of `SENTRY_DSN`.
|
||||
pub disabled: bool,
|
||||
}
|
||||
|
||||
static CONFIG: OnceLock<Config> = OnceLock::new();
|
||||
|
||||
// ─── Public API ────────────────────────────────────────────────────────────
|
||||
|
||||
/// Init Sentry + apply the process-wide scope tags. Call once at process
|
||||
/// start; the returned guard must outlive the process. No-op guard when
|
||||
/// `config.disabled`.
|
||||
pub fn init(config: Config) -> ClientInitGuard {
|
||||
let config = CONFIG.get_or_init(|| config);
|
||||
|
||||
if config.disabled {
|
||||
return sentry::init(ClientOptions::default());
|
||||
}
|
||||
|
||||
let dsn = std::env::var("SENTRY_DSN")
|
||||
.ok()
|
||||
.or_else(|| option_env!("SENTRY_DSN").map(|s| s.to_string()))
|
||||
.unwrap_or_default();
|
||||
|
||||
let scrubber = Scrubber::from_env();
|
||||
|
||||
let guard = sentry::init((
|
||||
dsn.as_str(),
|
||||
ClientOptions {
|
||||
release: Some(config.release.into()),
|
||||
send_default_pii: false,
|
||||
server_name: Some("".into()),
|
||||
attach_stacktrace: true,
|
||||
traces_sample_rate: TRACES_SAMPLE_RATE,
|
||||
environment: Some(environment().into()),
|
||||
before_send: Some(Arc::new(move |event| before_send(event, &scrubber))),
|
||||
..Default::default()
|
||||
},
|
||||
));
|
||||
|
||||
sentry::configure_scope(|scope| {
|
||||
scope.set_tag("client", config.client);
|
||||
scope.set_tag("client_version", config.client_version);
|
||||
scope.set_tag("os", std::env::consts::OS);
|
||||
scope.set_tag("arch", std::env::consts::ARCH);
|
||||
});
|
||||
|
||||
guard
|
||||
}
|
||||
|
||||
/// Flush in-flight events. Call before `std::process::exit` in signal handlers.
|
||||
pub fn flush_on_shutdown() {
|
||||
if let Some(client) = sentry::Hub::current().client() {
|
||||
client.flush(Some(FLUSH_TIMEOUT));
|
||||
}
|
||||
}
|
||||
|
||||
// ─── Internals ─────────────────────────────────────────────────────────────
|
||||
|
||||
#[derive(Clone, Debug)]
|
||||
struct Scrubber {
|
||||
home_dir: Option<String>,
|
||||
usernames: Vec<String>,
|
||||
}
|
||||
|
||||
impl Scrubber {
|
||||
fn from_env() -> Self {
|
||||
Self {
|
||||
home_dir: dirs::home_dir().map(|p| p.to_string_lossy().to_string()),
|
||||
usernames: collect_usernames_from_env(),
|
||||
}
|
||||
}
|
||||
|
||||
fn scrub(&self, s: &str) -> String {
|
||||
let out = xai_grok_secrets::redact_secrets(s);
|
||||
let out = match self.home_dir.as_deref() {
|
||||
Some(home) => Cow::Owned(replace_home_prefix(out.as_ref(), home)),
|
||||
None => out,
|
||||
};
|
||||
redact_username_segments(out.as_ref(), &self.usernames)
|
||||
}
|
||||
|
||||
fn scrub_value(&self, val: &mut Value) {
|
||||
xai_grok_secrets::walk_json_strings(val, &mut |s| *s = self.scrub(s));
|
||||
}
|
||||
}
|
||||
|
||||
const REDACTED_USER: &str = "<user>";
|
||||
|
||||
/// `$USERNAME` then `$USER`, deduped, 3-char floor to avoid over-matching.
|
||||
fn collect_usernames_from_env() -> Vec<String> {
|
||||
let mut usernames: Vec<String> = Vec::new();
|
||||
for var in ["USERNAME", "USER"] {
|
||||
if let Ok(name) = std::env::var(var) {
|
||||
let trimmed = name.trim();
|
||||
if trimmed.len() >= 3 && !usernames.iter().any(|u| u.eq_ignore_ascii_case(trimmed)) {
|
||||
usernames.push(trimmed.to_owned());
|
||||
}
|
||||
}
|
||||
}
|
||||
usernames
|
||||
}
|
||||
|
||||
/// Replace whole `/`- or `\`-delimited segments matching any entry in
|
||||
/// `usernames` with `<user>`. Substrings inside a segment are untouched.
|
||||
/// Case-insensitive on Windows (NTFS), case-sensitive elsewhere (POSIX).
|
||||
fn redact_username_segments(value: &str, usernames: &[String]) -> String {
|
||||
if usernames.is_empty() {
|
||||
return value.to_owned();
|
||||
}
|
||||
let mut out = String::with_capacity(value.len());
|
||||
let mut buf = String::new();
|
||||
for ch in value.chars() {
|
||||
if ch == '/' || ch == '\\' {
|
||||
push_segment(&mut out, &buf, usernames);
|
||||
buf.clear();
|
||||
out.push(ch);
|
||||
} else {
|
||||
buf.push(ch);
|
||||
}
|
||||
}
|
||||
push_segment(&mut out, &buf, usernames);
|
||||
out
|
||||
}
|
||||
|
||||
fn push_segment(out: &mut String, segment: &str, usernames: &[String]) {
|
||||
let matches = if cfg!(windows) {
|
||||
usernames.iter().any(|u| u.eq_ignore_ascii_case(segment))
|
||||
} else {
|
||||
usernames.iter().any(|u| u == segment)
|
||||
};
|
||||
if matches {
|
||||
out.push_str(REDACTED_USER);
|
||||
} else {
|
||||
out.push_str(segment);
|
||||
}
|
||||
}
|
||||
|
||||
/// Whole-segment `home` -> `~` so `/Users/bob` doesn't fold over `/Users/bobby/...`.
|
||||
fn replace_home_prefix(input: &str, home: &str) -> String {
|
||||
if home.is_empty() || !input.contains(home) {
|
||||
return input.to_owned();
|
||||
}
|
||||
let mut out = String::with_capacity(input.len());
|
||||
let mut rest = input;
|
||||
while let Some(idx) = rest.find(home) {
|
||||
let (before, tail) = rest.split_at(idx);
|
||||
out.push_str(before);
|
||||
let after = &tail[home.len()..];
|
||||
let prev_ok = before.chars().last().is_none_or(is_segment_boundary_char);
|
||||
let next_ok = after
|
||||
.chars()
|
||||
.next()
|
||||
.is_none_or(|c| c == '/' || c == '\\' || is_segment_boundary_char(c));
|
||||
if prev_ok && next_ok {
|
||||
out.push('~');
|
||||
} else {
|
||||
out.push_str(home);
|
||||
}
|
||||
rest = after;
|
||||
}
|
||||
out.push_str(rest);
|
||||
out
|
||||
}
|
||||
|
||||
fn is_segment_boundary_char(c: char) -> bool {
|
||||
c.is_whitespace()
|
||||
|| matches!(
|
||||
c,
|
||||
'"' | '\'' | '(' | '[' | '{' | ',' | ':' | ';' | '=' | '<'
|
||||
)
|
||||
}
|
||||
|
||||
fn before_send(mut event: Event<'static>, scrubber: &Scrubber) -> Option<Event<'static>> {
|
||||
if is_broken_pipe_panic(&event) {
|
||||
return None;
|
||||
}
|
||||
|
||||
if let Some(ref msg) = event.message {
|
||||
event.message = Some(scrubber.scrub(msg));
|
||||
}
|
||||
|
||||
for ex in &mut event.exception.values {
|
||||
if let Some(ref val) = ex.value {
|
||||
ex.value = Some(scrubber.scrub(val));
|
||||
}
|
||||
if let Some(ref mut stacktrace) = ex.stacktrace {
|
||||
for frame in &mut stacktrace.frames {
|
||||
if let Some(ref f) = frame.filename {
|
||||
frame.filename = Some(scrubber.scrub(f));
|
||||
}
|
||||
if let Some(ref f) = frame.abs_path {
|
||||
frame.abs_path = Some(scrubber.scrub(f));
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
for bc in &mut event.breadcrumbs.values {
|
||||
if let Some(ref msg) = bc.message {
|
||||
bc.message = Some(scrubber.scrub(msg));
|
||||
}
|
||||
for val in bc.data.values_mut() {
|
||||
scrubber.scrub_value(val);
|
||||
}
|
||||
}
|
||||
|
||||
event.extra.remove("cwd");
|
||||
for val in event.extra.values_mut() {
|
||||
scrubber.scrub_value(val);
|
||||
}
|
||||
|
||||
for tag in event.tags.values_mut() {
|
||||
*tag = scrubber.scrub(tag);
|
||||
}
|
||||
|
||||
event.server_name = None;
|
||||
|
||||
Some(event)
|
||||
}
|
||||
|
||||
/// Drop panics caused by broken pipe or disk-full, both user-environment noise.
|
||||
fn is_broken_pipe_panic(event: &Event<'_>) -> bool {
|
||||
event.exception.values.iter().any(|ex| {
|
||||
ex.value.as_deref().is_some_and(|v| {
|
||||
v.contains("Broken pipe")
|
||||
|| v.contains("os error 32")
|
||||
|| v.contains("No space left on device")
|
||||
|| v.contains("os error 28")
|
||||
})
|
||||
})
|
||||
}
|
||||
|
||||
fn environment() -> &'static str {
|
||||
if cfg!(debug_assertions) {
|
||||
"development"
|
||||
} else {
|
||||
"production"
|
||||
}
|
||||
}
|
||||
|
||||
// ─── Tests ─────────────────────────────────────────────────────────────────
|
||||
|
||||
#[cfg(test)]
|
||||
#[allow(clippy::field_reassign_with_default)] // setup reads better field-by-field
|
||||
mod tests {
|
||||
use super::*;
|
||||
use sentry::protocol::{Breadcrumb, Exception, Frame, Stacktrace};
|
||||
|
||||
fn make_scrubber() -> Scrubber {
|
||||
Scrubber {
|
||||
home_dir: Some("/Users/alice".to_string()),
|
||||
usernames: vec!["alice".to_owned()],
|
||||
}
|
||||
}
|
||||
|
||||
/// `aliceapp.log` case guards against a refactor to naive `str::replace`.
|
||||
#[test]
|
||||
fn scrub_applies_all_layers() {
|
||||
let s = make_scrubber();
|
||||
assert_eq!(s.scrub("/Users/alice/code/foo.rs"), "~/code/foo.rs");
|
||||
assert_eq!(s.scrub("/srv/alice/data"), "/srv/<user>/data");
|
||||
assert_eq!(s.scrub("aliceapp.log"), "aliceapp.log");
|
||||
assert!(s.scrub("token=longvalue123").contains("[REDACTED_SECRET]"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn before_send_drops_broken_pipe_keeps_others() {
|
||||
let s = make_scrubber();
|
||||
|
||||
let mut event = Event::default();
|
||||
event.exception.values.push(Exception {
|
||||
value: Some("Broken pipe (os error 32)".into()),
|
||||
..Default::default()
|
||||
});
|
||||
assert!(before_send(event, &s).is_none());
|
||||
|
||||
let mut event = Event::default();
|
||||
event.exception.values.push(Exception {
|
||||
value: Some("unrelated panic".into()),
|
||||
..Default::default()
|
||||
});
|
||||
assert!(before_send(event, &s).is_some());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn before_send_drops_no_space_left_panic() {
|
||||
let s = make_scrubber();
|
||||
let mut event = Event::default();
|
||||
event.exception.values.push(Exception {
|
||||
value: Some("failed printing to stderr: No space left on device (os error 28)".into()),
|
||||
..Default::default()
|
||||
});
|
||||
assert!(before_send(event, &s).is_none());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn before_send_scrubs_message_exception_stacktrace_breadcrumbs() {
|
||||
let s = make_scrubber();
|
||||
let mut event = Event::default();
|
||||
event.message = Some("error in /Users/alice/foo".into());
|
||||
event.exception.values.push(Exception {
|
||||
value: Some("/Users/alice/x failed".into()),
|
||||
stacktrace: Some(Stacktrace {
|
||||
frames: vec![Frame {
|
||||
filename: Some("/Users/alice/src/lib.rs".into()),
|
||||
abs_path: Some("/Users/alice/src/lib.rs".into()),
|
||||
..Default::default()
|
||||
}],
|
||||
..Default::default()
|
||||
}),
|
||||
..Default::default()
|
||||
});
|
||||
event.breadcrumbs.values.push(Breadcrumb {
|
||||
message: Some("opened /srv/alice/log".into()),
|
||||
..Default::default()
|
||||
});
|
||||
|
||||
let out = before_send(event, &s).unwrap();
|
||||
assert_eq!(out.message.as_deref(), Some("error in ~/foo"));
|
||||
let ex = &out.exception.values[0];
|
||||
assert_eq!(ex.value.as_deref(), Some("~/x failed"));
|
||||
let frame = &ex.stacktrace.as_ref().unwrap().frames[0];
|
||||
assert_eq!(frame.filename.as_deref(), Some("~/src/lib.rs"));
|
||||
assert_eq!(frame.abs_path.as_deref(), Some("~/src/lib.rs"));
|
||||
assert_eq!(
|
||||
out.breadcrumbs.values[0].message.as_deref(),
|
||||
Some("opened /srv/<user>/log"),
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn before_send_handles_extras_and_server_name() {
|
||||
let s = make_scrubber();
|
||||
let mut event = Event::default();
|
||||
event.extra.insert("cwd".into(), "/Users/alice/proj".into());
|
||||
event
|
||||
.extra
|
||||
.insert("other".into(), "/Users/alice/foo".into());
|
||||
event.server_name = Some("hostname.example.com".into());
|
||||
|
||||
let out = before_send(event, &s).unwrap();
|
||||
assert!(!out.extra.contains_key("cwd"));
|
||||
assert_eq!(
|
||||
out.extra.get("other").and_then(|v| v.as_str()),
|
||||
Some("~/foo")
|
||||
);
|
||||
assert!(out.server_name.is_none());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn before_send_scrubs_breadcrumb_data() {
|
||||
let s = make_scrubber();
|
||||
let mut event = Event::default();
|
||||
let mut bc = Breadcrumb::default();
|
||||
bc.data.insert("path".into(), "/Users/alice/foo".into());
|
||||
event.breadcrumbs.values.push(bc);
|
||||
|
||||
let out = before_send(event, &s).unwrap();
|
||||
assert_eq!(
|
||||
out.breadcrumbs.values[0]
|
||||
.data
|
||||
.get("path")
|
||||
.and_then(|v| v.as_str()),
|
||||
Some("~/foo"),
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn before_send_scrubs_tags() {
|
||||
let s = make_scrubber();
|
||||
let mut event = Event::default();
|
||||
event.tags.insert("workspace".into(), "/srv/alice/x".into());
|
||||
|
||||
let out = before_send(event, &s).unwrap();
|
||||
assert_eq!(
|
||||
out.tags.get("workspace").map(String::as_str),
|
||||
Some("/srv/<user>/x")
|
||||
);
|
||||
}
|
||||
|
||||
/// `/Users/bob` must not partial-match the prefix of `/Users/bobby/...`.
|
||||
#[test]
|
||||
fn home_dir_replacement_is_segment_aware() {
|
||||
let s = Scrubber {
|
||||
home_dir: Some("/Users/bob".to_string()),
|
||||
usernames: vec![],
|
||||
};
|
||||
assert_eq!(s.scrub("/Users/bobby/code"), "/Users/bobby/code");
|
||||
assert_eq!(s.scrub("opened /Users/bob/x"), "opened ~/x");
|
||||
}
|
||||
|
||||
#[cfg(not(windows))]
|
||||
#[test]
|
||||
fn unix_match_is_case_sensitive() {
|
||||
let out = redact_username_segments("/Users/Alice/proj", &["alice".to_owned()]);
|
||||
assert_eq!(out, "/Users/Alice/proj");
|
||||
let out = redact_username_segments("/Users/alice/proj", &["alice".to_owned()]);
|
||||
assert_eq!(out, "/Users/<user>/proj");
|
||||
}
|
||||
|
||||
#[cfg(windows)]
|
||||
#[test]
|
||||
fn windows_match_is_case_insensitive() {
|
||||
let out = redact_username_segments(r"C:\Users\Alice\proj", &["alice".to_owned()]);
|
||||
assert_eq!(out, r"C:\Users\<user>\proj");
|
||||
}
|
||||
}
|
||||
326
crates/codegen/xai-grok-telemetry/src/session_ctx.rs
Normal file
326
crates/codegen/xai-grok-telemetry/src/session_ctx.rs
Normal file
|
|
@ -0,0 +1,326 @@
|
|||
//! Ambient session context for telemetry — product events + Mixpanel via
|
||||
//! [`log_event`]. `session_id` and `turn_number` are injected from the
|
||||
//! task-local [`TelemetryCtx`] active for the duration of a session.
|
||||
//!
|
||||
//! Extracted from `xai-grok-shell::agent::telemetry`.
|
||||
|
||||
use std::sync::Arc;
|
||||
|
||||
use serde::Serialize;
|
||||
use serde_json::json;
|
||||
|
||||
use crate::client::{self, Metadata, UserContext};
|
||||
use crate::events::TelemetryEvent;
|
||||
|
||||
/// Ambient session context for telemetry. Snapshotted synchronously by
|
||||
/// `log_event` at call time to avoid racing with turn increments.
|
||||
#[derive(Clone)]
|
||||
pub struct TelemetryCtx {
|
||||
pub session_id: String,
|
||||
pub prompt_index: Arc<tokio::sync::Mutex<usize>>,
|
||||
/// Per-prompt correlation UUID for the external OTEL stream (`prompt.id`,
|
||||
/// events only — never metrics). Set at turn start where `prompt_index`
|
||||
/// increments; `None` outside a prompt.
|
||||
pub prompt_id: Arc<parking_lot::Mutex<Option<String>>>,
|
||||
}
|
||||
|
||||
impl TelemetryCtx {
|
||||
pub fn new(session_id: String, prompt_index: Arc<tokio::sync::Mutex<usize>>) -> Self {
|
||||
Self {
|
||||
session_id,
|
||||
prompt_index,
|
||||
prompt_id: Arc::new(parking_lot::Mutex::new(None)),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Snapshot of the ambient ctx for the external OTEL stream.
|
||||
pub(crate) struct ExternalCtxSnapshot {
|
||||
pub session_id: String,
|
||||
pub turn_number: Option<u32>,
|
||||
pub prompt_id: Option<String>,
|
||||
}
|
||||
|
||||
/// Rotate the per-prompt correlation UUID at turn start (where
|
||||
/// `prompt_index` increments). No-op outside a session ctx scope. The id is
|
||||
/// attached as `prompt.id` to external OTEL events only.
|
||||
pub fn begin_prompt_id() {
|
||||
let _ = TELEMETRY_CTX.try_with(|c| {
|
||||
*c.prompt_id.lock() = Some(uuid::Uuid::new_v4().to_string());
|
||||
});
|
||||
}
|
||||
|
||||
/// Snapshot the task-local ctx (if any) for external emission. Non-blocking:
|
||||
/// a contended `prompt_index` lock yields `turn_number = None` rather than
|
||||
/// stalling the emitting task.
|
||||
pub(crate) fn external_ctx_snapshot() -> Option<ExternalCtxSnapshot> {
|
||||
TELEMETRY_CTX
|
||||
.try_with(|c| ExternalCtxSnapshot {
|
||||
session_id: c.session_id.clone(),
|
||||
turn_number: c.prompt_index.try_lock().map(|g| *g as u32).ok(),
|
||||
prompt_id: c.prompt_id.lock().clone(),
|
||||
})
|
||||
.ok()
|
||||
}
|
||||
|
||||
tokio::task_local! {
|
||||
static TELEMETRY_CTX: Arc<TelemetryCtx>;
|
||||
}
|
||||
|
||||
/// The `session_id` field name the debug-log firehose router keys on:
|
||||
/// `debug_log::SessionIdVisitor` stashes a `SessionId` extension on any span
|
||||
/// carrying this field — the span *name* is not load-bearing for routing. Shared
|
||||
/// so the `info_span!` here and the router in `debug_log` can't silently drift; a
|
||||
/// rename trips `session_span_exposes_router_field` below.
|
||||
pub(crate) const SESSION_ID_FIELD: &str = "session_id";
|
||||
|
||||
/// Build the per-session tracing span the firehose router routes by. The field
|
||||
/// name MUST be the literal `session_id` (tracing field names can't come from a
|
||||
/// const); the test below pins it against [`SESSION_ID_FIELD`].
|
||||
fn session_span(session_id: &str) -> tracing::Span {
|
||||
tracing::info_span!("session", session_id = %session_id)
|
||||
}
|
||||
|
||||
/// Run `fut` with telemetry context active. Also sets a `tracing` span.
|
||||
pub async fn with_session_ctx<F: std::future::Future>(ctx: TelemetryCtx, fut: F) -> F::Output {
|
||||
use tracing::Instrument;
|
||||
let span = session_span(&ctx.session_id);
|
||||
TELEMETRY_CTX
|
||||
.scope(Arc::new(ctx), fut.instrument(span))
|
||||
.await
|
||||
}
|
||||
|
||||
/// Product surface that emitted a telemetry event. Selects the analytics
|
||||
/// event-name prefix so shell and workspace events are distinguishable on the
|
||||
/// wire while sharing this emitter (and the `event_value` derivation in
|
||||
/// [`crate::client`]).
|
||||
#[derive(Debug, Clone, Copy, PartialEq, Eq, strum::EnumCount)]
|
||||
pub enum EmitterOrigin {
|
||||
/// `xai-grok-shell` (and the pager/TUI that emit through it).
|
||||
Shell,
|
||||
/// `xai-grok-workspace` (remote sampler / workspace server).
|
||||
Workspace,
|
||||
}
|
||||
|
||||
impl EmitterOrigin {
|
||||
/// Every emitter origin. [`crate::client::event_value`] iterates this to
|
||||
/// strip whichever prefix an event name carries. Iteration *order* is
|
||||
/// irrelevant: the prefixes are mutually exclusive (no
|
||||
/// [`EmitterOrigin::event_prefix`] is a prefix of another — pinned by
|
||||
/// `client`'s `emitter_prefixes_are_mutually_exclusive` test), so at most
|
||||
/// one entry ever matches a given name. Completeness is compiler-enforced
|
||||
/// by the `EmitterOrigin::ALL` length assertion below, so a newly added
|
||||
/// variant that is omitted here fails to compile.
|
||||
pub const ALL: [EmitterOrigin; 2] = [EmitterOrigin::Shell, EmitterOrigin::Workspace];
|
||||
|
||||
/// Analytics event-name prefix for this origin. [`crate::client::event_value`]
|
||||
/// strips the same prefix to derive the wire `event_value`, so the two must
|
||||
/// stay in lockstep.
|
||||
pub fn event_prefix(self) -> &'static str {
|
||||
match self {
|
||||
EmitterOrigin::Shell => "grok-shell-",
|
||||
EmitterOrigin::Workspace => "grok-workspace-",
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Compile-time completeness guard for [`EmitterOrigin::ALL`]: adding a variant
|
||||
/// without listing it in `ALL` makes `ALL.len()` diverge from the
|
||||
/// `strum::EnumCount`-derived variant count and fails this assertion, so
|
||||
/// `client::event_value` can never silently stop stripping an origin's prefix.
|
||||
const _: () = assert!(EmitterOrigin::ALL.len() == <EmitterOrigin as strum::EnumCount>::COUNT);
|
||||
|
||||
/// Product analytics event (type-safe). Only fires in `Enabled` mode.
|
||||
/// Unconditionally fans out to the external OTEL stream first ("one call
|
||||
/// site, two sinks, independent gates"): the external gate is
|
||||
/// `external::is_active()`, independent of `TelemetryMode`.
|
||||
pub fn log_event<T: TelemetryEvent>(data: T) {
|
||||
crate::external::emit(&data);
|
||||
if !client::is_enabled() {
|
||||
return;
|
||||
}
|
||||
emit_event(T::NAME, data);
|
||||
}
|
||||
|
||||
/// Emit one event to the external stream always (no-op unless the stream is
|
||||
/// active) and to the product events/Mixpanel funnel only when `internal_enabled`.
|
||||
///
|
||||
/// Used by call sites whose internal sink is gated by a *stricter* predicate
|
||||
/// than [`log_event`]'s own `TelemetryMode::Enabled` check (the shell's
|
||||
/// `telemetry_enabled` = `Enabled && !ZDR`, or `!is_data_collection_disabled()`).
|
||||
/// Because [`log_event`] already fans out to the external sink before its
|
||||
/// internal gate, the two branches are **mutually exclusive**: routing through
|
||||
/// `log_event` when internal is enabled reaches both sinks, and calling
|
||||
/// [`crate::external::emit`] directly otherwise keeps `session.count` /
|
||||
/// `turn.count` exactly-once on every path while never sending an internal
|
||||
/// record under ZDR.
|
||||
pub fn log_event_dual<T: TelemetryEvent>(internal_enabled: bool, data: T) {
|
||||
if internal_enabled {
|
||||
log_event(data);
|
||||
} else {
|
||||
crate::external::emit(&data);
|
||||
}
|
||||
}
|
||||
|
||||
/// Session lifecycle event (type-safe). Fires in both `Enabled` and
|
||||
/// `SessionMetrics` modes. Emits with the [`EmitterOrigin::Shell`] prefix;
|
||||
/// workspace-side callers use [`log_session_event_with_origin`].
|
||||
/// Unconditionally fans out to the external OTEL stream first (independent
|
||||
/// gate; see [`log_event`]).
|
||||
pub fn log_session_event<T: TelemetryEvent>(data: T) {
|
||||
crate::external::emit(&data);
|
||||
if !client::is_session_metrics_enabled() {
|
||||
return;
|
||||
}
|
||||
emit_event_with_origin(EmitterOrigin::Shell, T::NAME, data);
|
||||
}
|
||||
|
||||
/// Session lifecycle event tagged with the emitting [`EmitterOrigin`]. Fires in
|
||||
/// both `Enabled` and `SessionMetrics` modes; the origin selects the analytics
|
||||
/// event-name prefix (`grok-shell-*` vs `grok-workspace-*`).
|
||||
///
|
||||
/// Deliberately **no external fan-out** here: workspace-side callers
|
||||
/// (`EmitterOrigin::Workspace` — remote sampler / workspace server, a
|
||||
/// different process and monitoring audience) invoke this directly, and the
|
||||
/// external stream is Shell-origin only. An `external = …` macro arm on a
|
||||
/// workspace-only event therefore has no effect (pinned by test in
|
||||
/// `external::tests`). If the external stream ever needs workspace events,
|
||||
/// the hook moves here behind an explicit `origin == Shell` filter.
|
||||
pub fn log_session_event_with_origin<T: TelemetryEvent>(origin: EmitterOrigin, data: T) {
|
||||
if !client::is_session_metrics_enabled() {
|
||||
return;
|
||||
}
|
||||
emit_event_with_origin(origin, T::NAME, data);
|
||||
}
|
||||
|
||||
/// Emit an event with the default [`EmitterOrigin::Shell`] prefix.
|
||||
pub fn emit_event<T: Serialize + Send + 'static>(event_suffix: impl Into<String>, data: T) {
|
||||
emit_event_with_origin(EmitterOrigin::Shell, event_suffix, data);
|
||||
}
|
||||
|
||||
/// Emit an event whose analytics name is `{origin prefix}{event_suffix}`.
|
||||
pub fn emit_event_with_origin<T: Serialize + Send + 'static>(
|
||||
origin: EmitterOrigin,
|
||||
event_suffix: impl Into<String>,
|
||||
data: T,
|
||||
) {
|
||||
let event_name = format!("{}{}", origin.event_prefix(), event_suffix.into());
|
||||
let ctx_snapshot = TELEMETRY_CTX
|
||||
.try_with(|c| {
|
||||
(
|
||||
c.session_id.clone(),
|
||||
c.prompt_index.try_lock().map(|g| *g as u32).ok(),
|
||||
)
|
||||
})
|
||||
.ok();
|
||||
|
||||
tokio::spawn(async move {
|
||||
let user_ctx = UserContext::collect();
|
||||
let request_id = format!("{}-{}", event_name, uuid::Uuid::new_v4());
|
||||
|
||||
let mut metadata = match serde_json::to_value(data) {
|
||||
Ok(serde_json::Value::Object(map)) => map,
|
||||
Ok(other) => {
|
||||
let mut m = Metadata::new();
|
||||
m.insert("value".into(), other);
|
||||
m
|
||||
}
|
||||
Err(_) => Metadata::new(),
|
||||
};
|
||||
|
||||
if let Some((session_id, turn_number)) = ctx_snapshot {
|
||||
metadata.insert("session_id".into(), json!(session_id));
|
||||
if let Some(turn) = turn_number {
|
||||
metadata.insert("turn_number".into(), json!(turn));
|
||||
}
|
||||
}
|
||||
|
||||
client::track(&event_name, &request_id, &user_ctx, metadata).await;
|
||||
});
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
/// The debug-log firehose router (`debug_log`) finds the session span by its
|
||||
/// `session_id` field (not by name). That field name is a literal in
|
||||
/// `session_span` (tracing field names can't be a const), so pin it against the
|
||||
/// shared const here — a rename of either breaks this test instead of silently
|
||||
/// degrading routing to the per-pid fallback.
|
||||
#[test]
|
||||
fn session_span_exposes_router_field() {
|
||||
// A bare registry enables every callsite, so the span has live metadata.
|
||||
let subscriber = tracing_subscriber::registry();
|
||||
tracing::subscriber::with_default(subscriber, || {
|
||||
let span = session_span("test-id");
|
||||
let meta = span
|
||||
.metadata()
|
||||
.expect("session span must have metadata under an enabling subscriber");
|
||||
assert!(
|
||||
meta.fields().field(SESSION_ID_FIELD).is_some(),
|
||||
"session span must expose `{SESSION_ID_FIELD}` for debug-log routing",
|
||||
);
|
||||
});
|
||||
}
|
||||
|
||||
/// Event-name prefixes are wire contract — analytics queries match on them, so
|
||||
/// they must not drift.
|
||||
#[test]
|
||||
fn event_prefix_is_stable_per_origin() {
|
||||
assert_eq!(EmitterOrigin::Shell.event_prefix(), "grok-shell-");
|
||||
assert_eq!(EmitterOrigin::Workspace.event_prefix(), "grok-workspace-");
|
||||
}
|
||||
|
||||
/// The `Shell` reroute must reproduce the historical
|
||||
/// `format!("grok-shell-{suffix}")` event name byte-for-byte, since every
|
||||
/// existing `log_session_event` / `log_event` / `emit_event` call funnels
|
||||
/// through `EmitterOrigin::Shell`.
|
||||
#[test]
|
||||
fn shell_origin_event_name_matches_legacy_format() {
|
||||
let suffix = "trace_upload_attempted";
|
||||
let rerouted = format!("{}{}", EmitterOrigin::Shell.event_prefix(), suffix);
|
||||
let legacy = format!("grok-shell-{suffix}");
|
||||
assert_eq!(rerouted, legacy);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn workspace_origin_event_name_uses_workspace_prefix() {
|
||||
let name = format!("{}turn", EmitterOrigin::Workspace.event_prefix());
|
||||
assert_eq!(name, "grok-workspace-turn");
|
||||
}
|
||||
|
||||
/// `ALL` must enumerate every variant so the stripper in `client` can
|
||||
/// recover the `event_value` for any origin the emitter produces. Length
|
||||
/// completeness is also compiler-enforced by the `const _` assertion in
|
||||
/// this module (via `strum::EnumCount`); this test additionally pins that
|
||||
/// the known variants are present and that every origin yields a distinct,
|
||||
/// non-empty prefix (which `EnumCount` alone does not guarantee).
|
||||
#[test]
|
||||
fn all_covers_every_origin_with_distinct_nonempty_prefixes() {
|
||||
assert!(EmitterOrigin::ALL.contains(&EmitterOrigin::Shell));
|
||||
assert!(EmitterOrigin::ALL.contains(&EmitterOrigin::Workspace));
|
||||
assert_eq!(
|
||||
EmitterOrigin::ALL.len(),
|
||||
<EmitterOrigin as strum::EnumCount>::COUNT,
|
||||
"ALL must list every EmitterOrigin variant",
|
||||
);
|
||||
|
||||
let mut prefixes: Vec<&str> = EmitterOrigin::ALL
|
||||
.iter()
|
||||
.map(|o| o.event_prefix())
|
||||
.collect();
|
||||
assert!(
|
||||
prefixes.iter().all(|p| !p.is_empty()),
|
||||
"every origin must have a non-empty prefix",
|
||||
);
|
||||
let total = prefixes.len();
|
||||
prefixes.sort_unstable();
|
||||
prefixes.dedup();
|
||||
assert_eq!(
|
||||
prefixes.len(),
|
||||
total,
|
||||
"every origin must yield a distinct prefix",
|
||||
);
|
||||
}
|
||||
}
|
||||
213
crates/codegen/xai-grok-telemetry/src/session_metrics.rs
Normal file
213
crates/codegen/xai-grok-telemetry/src/session_metrics.rs
Normal file
|
|
@ -0,0 +1,213 @@
|
|||
//! Session lifecycle event structs.
|
||||
//!
|
||||
//! Fires in both `Enabled` and `SessionMetrics` telemetry modes via
|
||||
//! `log_session_event`.
|
||||
|
||||
use serde::Serialize;
|
||||
|
||||
#[derive(Serialize)]
|
||||
pub struct SessionStarted {
|
||||
pub session_id: String,
|
||||
}
|
||||
|
||||
#[derive(Serialize)]
|
||||
pub struct Turn {
|
||||
pub session_id: String,
|
||||
pub turn_number: u64,
|
||||
}
|
||||
|
||||
#[derive(Serialize)]
|
||||
pub struct TurnCompletedLifecycle {
|
||||
pub session_id: String,
|
||||
pub turn_number: u64,
|
||||
}
|
||||
|
||||
/// Doom-loop recovery acted this turn: poisoned attempts were resampled
|
||||
/// and/or a response was accepted with confident signals after the budget
|
||||
/// was spent. Trigger labels only — never generation content.
|
||||
#[derive(Serialize)]
|
||||
pub struct DoomLoopRecovery {
|
||||
pub session_id: String,
|
||||
pub turn_number: u64,
|
||||
/// Resamples this turn (doomed attempts discarded).
|
||||
pub attempts: u32,
|
||||
/// Whether the final response kept confident signals (budget spent).
|
||||
pub accepted_after_budget: bool,
|
||||
/// Tightest raw trigger label observed this turn.
|
||||
#[serde(skip_serializing_if = "Option::is_none")]
|
||||
pub top_trigger: Option<String>,
|
||||
/// Model that produced the doomed attempts.
|
||||
pub model: String,
|
||||
}
|
||||
|
||||
#[derive(Serialize)]
|
||||
pub struct TraceUploadAttempted {
|
||||
pub session_id: String,
|
||||
pub turn_number: u64,
|
||||
pub upload_method: String,
|
||||
}
|
||||
|
||||
#[derive(Serialize)]
|
||||
pub struct TraceUploadSucceeded {
|
||||
pub session_id: String,
|
||||
pub turn_number: u64,
|
||||
pub upload_method: String,
|
||||
pub fully_uploaded: bool,
|
||||
}
|
||||
|
||||
#[derive(Serialize)]
|
||||
pub struct TraceUploadSkipped {
|
||||
pub session_id: String,
|
||||
pub turn_number: u64,
|
||||
pub reason: String,
|
||||
}
|
||||
|
||||
#[derive(Serialize)]
|
||||
pub struct TraceUploadFailed {
|
||||
pub session_id: String,
|
||||
pub turn_number: u64,
|
||||
pub upload_method: String,
|
||||
pub error_category: String,
|
||||
#[serde(skip_serializing_if = "Option::is_none")]
|
||||
pub status_code: Option<u16>,
|
||||
}
|
||||
|
||||
/// Why trace uploads are enabled or disabled for a given prompt.
|
||||
/// Recorded on the `agent.prompt` span as `upload_reason` for analytics queries.
|
||||
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
|
||||
pub enum TraceUploadReason {
|
||||
/// ZDR (zero data retention) team — all uploads disabled.
|
||||
ZdrTeam,
|
||||
/// `[telemetry] trace_upload = false` in config.
|
||||
FeatureOff,
|
||||
/// No grok.com auth or deployment key.
|
||||
NoCredentials,
|
||||
/// Direct-to-bucket S3 upload.
|
||||
DirectS3,
|
||||
/// Proxy mode via grok.com auth.
|
||||
Proxy,
|
||||
/// Direct GCS with service account key.
|
||||
DirectGcs,
|
||||
/// Session handle not found (edge case).
|
||||
SessionNotFound,
|
||||
}
|
||||
|
||||
impl TraceUploadReason {
|
||||
pub fn as_str(self) -> &'static str {
|
||||
match self {
|
||||
Self::ZdrTeam => "zdr_team",
|
||||
Self::FeatureOff => "feature_off",
|
||||
Self::NoCredentials => "no_credentials",
|
||||
Self::DirectS3 => "direct_s3",
|
||||
Self::Proxy => "proxy",
|
||||
Self::DirectGcs => "direct_gcs",
|
||||
Self::SessionNotFound => "session_not_found",
|
||||
}
|
||||
}
|
||||
|
||||
pub fn from_upload_method(method: &Option<xai_file_utils::UploadMethod>) -> Self {
|
||||
match method {
|
||||
Some(xai_file_utils::UploadMethod::Proxy { .. }) => Self::Proxy,
|
||||
Some(xai_file_utils::UploadMethod::S3 { .. }) => Self::DirectS3,
|
||||
Some(xai_file_utils::UploadMethod::Direct { .. }) => Self::DirectGcs,
|
||||
None => Self::NoCredentials,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use xai_file_utils::UploadMethod;
|
||||
|
||||
use super::TraceUploadReason;
|
||||
|
||||
/// The `grok-shell-doom_loop_recovery` Mixpanel event's name and
|
||||
/// property keys are dashboard contracts — pin them.
|
||||
#[test]
|
||||
fn doom_loop_recovery_event_shape_is_stable() {
|
||||
use crate::events::TelemetryEvent;
|
||||
assert_eq!(super::DoomLoopRecovery::NAME, "doom_loop_recovery");
|
||||
let with_trigger = serde_json::to_value(super::DoomLoopRecovery {
|
||||
session_id: "s1".to_string(),
|
||||
turn_number: 7,
|
||||
attempts: 2,
|
||||
accepted_after_budget: true,
|
||||
top_trigger: Some("tail_repetition:4@thinking".to_string()),
|
||||
model: "grok-4.5".to_string(),
|
||||
})
|
||||
.unwrap();
|
||||
assert_eq!(
|
||||
with_trigger,
|
||||
serde_json::json!({
|
||||
"session_id": "s1",
|
||||
"turn_number": 7,
|
||||
"attempts": 2,
|
||||
"accepted_after_budget": true,
|
||||
"top_trigger": "tail_repetition:4@thinking",
|
||||
"model": "grok-4.5",
|
||||
})
|
||||
);
|
||||
let no_trigger = serde_json::to_value(super::DoomLoopRecovery {
|
||||
session_id: "s1".to_string(),
|
||||
turn_number: 7,
|
||||
attempts: 1,
|
||||
accepted_after_budget: false,
|
||||
top_trigger: None,
|
||||
model: "grok-4.5".to_string(),
|
||||
})
|
||||
.unwrap();
|
||||
assert!(no_trigger.get("top_trigger").is_none(), "None is omitted");
|
||||
}
|
||||
|
||||
/// `as_str` values are recorded on the `agent.prompt` span as
|
||||
/// `upload_reason` and queried in analytics — they are a wire contract and
|
||||
/// must not drift.
|
||||
#[test]
|
||||
fn as_str_values_are_stable() {
|
||||
assert_eq!(TraceUploadReason::ZdrTeam.as_str(), "zdr_team");
|
||||
assert_eq!(TraceUploadReason::FeatureOff.as_str(), "feature_off");
|
||||
assert_eq!(TraceUploadReason::NoCredentials.as_str(), "no_credentials");
|
||||
assert_eq!(TraceUploadReason::DirectS3.as_str(), "direct_s3");
|
||||
assert_eq!(TraceUploadReason::Proxy.as_str(), "proxy");
|
||||
assert_eq!(TraceUploadReason::DirectGcs.as_str(), "direct_gcs");
|
||||
assert_eq!(
|
||||
TraceUploadReason::SessionNotFound.as_str(),
|
||||
"session_not_found"
|
||||
);
|
||||
}
|
||||
|
||||
/// Each `UploadMethod` maps to its corresponding reason; `None` (no
|
||||
/// credentials resolved) maps to `NoCredentials`.
|
||||
#[test]
|
||||
fn from_upload_method_maps_each_variant() {
|
||||
assert_eq!(
|
||||
TraceUploadReason::from_upload_method(&None),
|
||||
TraceUploadReason::NoCredentials
|
||||
);
|
||||
assert_eq!(
|
||||
TraceUploadReason::from_upload_method(&Some(UploadMethod::Direct {
|
||||
service_account_key: None,
|
||||
})),
|
||||
TraceUploadReason::DirectGcs
|
||||
);
|
||||
assert_eq!(
|
||||
TraceUploadReason::from_upload_method(&Some(UploadMethod::Proxy {
|
||||
proxy_base_url: String::new(),
|
||||
user_token: String::new(),
|
||||
deployment_key: None,
|
||||
alpha_test_key: None,
|
||||
})),
|
||||
TraceUploadReason::Proxy
|
||||
);
|
||||
assert_eq!(
|
||||
TraceUploadReason::from_upload_method(&Some(UploadMethod::S3 {
|
||||
bucket: String::new(),
|
||||
region: String::new(),
|
||||
credentials_file: None,
|
||||
credentials_content: None,
|
||||
endpoint_url: None,
|
||||
})),
|
||||
TraceUploadReason::DirectS3
|
||||
);
|
||||
}
|
||||
}
|
||||
505
crates/codegen/xai-grok-telemetry/src/unified_log.rs
Normal file
505
crates/codegen/xai-grok-telemetry/src/unified_log.rs
Normal file
|
|
@ -0,0 +1,505 @@
|
|||
//! Centralized unified log for cross-component session observability.
|
||||
//!
|
||||
//! Shell writes directly via [`emit()`]. Pager and desktop forward entries
|
||||
//! over ACP (`x.ai/log` notifications); shell receives them in
|
||||
//! [`ingest_client_entries()`] and writes on their behalf.
|
||||
|
||||
use std::fs::{self, File, OpenOptions};
|
||||
use std::io::Write;
|
||||
use std::path::PathBuf;
|
||||
use std::sync::{LazyLock, Mutex, OnceLock};
|
||||
|
||||
use chrono::Utc;
|
||||
use serde::{Deserialize, Serialize};
|
||||
|
||||
use xai_grok_config::grok_home;
|
||||
|
||||
/// Binary version stamped into every log entry. Set once at startup via
|
||||
/// [`set_version()`]; entries emitted before that get `None`.
|
||||
static VERSION: OnceLock<String> = OnceLock::new();
|
||||
|
||||
/// Register the binary version (e.g. shell's `CARGO_PKG_VERSION`).
|
||||
/// Call once at startup; subsequent calls are no-ops.
|
||||
pub fn set_version(ver: &str) {
|
||||
let _ = VERSION.set(ver.to_owned());
|
||||
}
|
||||
|
||||
pub const LOG_DIR: &str = "logs";
|
||||
const LOG_FILE: &str = "unified.jsonl";
|
||||
pub const MAX_SIZE: u64 = 5 * 1024 * 1024; // 5 MB
|
||||
|
||||
/// ACP method name for unified log notifications.
|
||||
pub const LOG_METHOD: &str = "x.ai/log";
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// Log entry types
|
||||
// ---------------------------------------------------------------------------
|
||||
|
||||
/// Log level for a unified log entry.
|
||||
#[derive(Debug, Clone, Copy, PartialEq, Eq, strum::Display, Serialize, Deserialize)]
|
||||
#[strum(serialize_all = "lowercase")]
|
||||
#[serde(rename_all = "lowercase")]
|
||||
pub enum LogLevel {
|
||||
Error,
|
||||
Warn,
|
||||
Info,
|
||||
Debug,
|
||||
}
|
||||
|
||||
/// Component that produced a log entry.
|
||||
#[derive(Debug, Clone, Copy, PartialEq, Eq, strum::Display, Serialize, Deserialize)]
|
||||
pub enum LogSource {
|
||||
#[strum(serialize = "shell")]
|
||||
#[serde(rename = "shell")]
|
||||
Shell,
|
||||
#[strum(serialize = "grok-pager")]
|
||||
#[serde(rename = "grok-pager")]
|
||||
GrokPager,
|
||||
#[strum(serialize = "grok-desktop")]
|
||||
#[serde(rename = "grok-desktop")]
|
||||
GrokDesktop,
|
||||
}
|
||||
|
||||
/// A single unified log entry, written as one JSONL line.
|
||||
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||
pub struct LogEntry {
|
||||
/// RFC 3339 timestamp (millisecond precision, UTC).
|
||||
pub ts: String,
|
||||
/// Component that produced the entry.
|
||||
pub src: LogSource,
|
||||
/// OS process id of the producer. Critical for cross-process trace
|
||||
/// reconstruction because shell/pager/desktop all append to the same
|
||||
/// `unified.jsonl`, so multiple shell processes' lines interleave
|
||||
/// indistinguishably without it.
|
||||
///
|
||||
/// `Option<u32>` is for wire compatibility only -- shell, pager, and
|
||||
/// desktop all stamp `Some(std::process::id())` at emit time. A
|
||||
/// `None` here means the entry came from an older client/server that
|
||||
/// predates this field; current code never emits one.
|
||||
#[serde(skip_serializing_if = "Option::is_none")]
|
||||
pub pid: Option<u32>,
|
||||
/// Binary version (e.g. `"0.1.211"`). Stamped by [`set_version()`]
|
||||
/// at startup so stale zombie processes are identifiable in logs.
|
||||
/// `None` for entries from older binaries that predate this field.
|
||||
#[serde(skip_serializing_if = "Option::is_none")]
|
||||
pub ver: Option<String>,
|
||||
/// Log level.
|
||||
pub lvl: LogLevel,
|
||||
/// Session ID, if one exists.
|
||||
#[serde(skip_serializing_if = "Option::is_none")]
|
||||
pub sid: Option<String>,
|
||||
/// Human-readable message.
|
||||
pub msg: String,
|
||||
/// Structured context fields.
|
||||
#[serde(skip_serializing_if = "Option::is_none")]
|
||||
pub ctx: Option<serde_json::Value>,
|
||||
}
|
||||
|
||||
/// Wire format for the `x.ai/log` ACP notification params.
|
||||
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||
pub struct LogNotificationParams {
|
||||
/// Source component identifier.
|
||||
pub src: LogSource,
|
||||
pub entries: Vec<ClientLogEntry>,
|
||||
}
|
||||
|
||||
/// Entry as sent by a client (no `src` field -- shell stamps it).
|
||||
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||
pub struct ClientLogEntry {
|
||||
pub ts: String,
|
||||
/// Client process id. Stamped by the client when the entry is
|
||||
/// created; preserved through ACP forwarding so the on-disk log
|
||||
/// reflects the originating process.
|
||||
///
|
||||
/// Optional only for wire compatibility with clients that predate
|
||||
/// this field; in-tree clients always populate it.
|
||||
#[serde(skip_serializing_if = "Option::is_none")]
|
||||
pub pid: Option<u32>,
|
||||
/// Binary version. Optional for wire compatibility with older clients.
|
||||
#[serde(skip_serializing_if = "Option::is_none")]
|
||||
pub ver: Option<String>,
|
||||
pub lvl: LogLevel,
|
||||
#[serde(skip_serializing_if = "Option::is_none")]
|
||||
pub sid: Option<String>,
|
||||
pub msg: String,
|
||||
#[serde(skip_serializing_if = "Option::is_none")]
|
||||
pub ctx: Option<serde_json::Value>,
|
||||
}
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// Writer
|
||||
// ---------------------------------------------------------------------------
|
||||
|
||||
struct LogWriter {
|
||||
file: File,
|
||||
path: PathBuf,
|
||||
written: u64,
|
||||
}
|
||||
|
||||
static WRITER: LazyLock<Mutex<Option<LogWriter>>> = LazyLock::new(|| Mutex::new(open_writer()));
|
||||
|
||||
fn log_path() -> PathBuf {
|
||||
grok_home().join(LOG_DIR).join(LOG_FILE)
|
||||
}
|
||||
|
||||
pub fn file_size(path: &std::path::Path) -> u64 {
|
||||
fs::metadata(path).map(|m| m.len()).unwrap_or(0)
|
||||
}
|
||||
|
||||
fn open_writer() -> Option<LogWriter> {
|
||||
let path = log_path();
|
||||
if let Some(parent) = path.parent()
|
||||
&& let Err(e) = fs::create_dir_all(parent)
|
||||
{
|
||||
tracing::warn!("[unified_log] failed to create log dir: {e}");
|
||||
return None;
|
||||
}
|
||||
|
||||
if file_size(&path) >= MAX_SIZE {
|
||||
trim_file(&path);
|
||||
}
|
||||
|
||||
match OpenOptions::new().create(true).append(true).open(&path) {
|
||||
Ok(file) => Some(LogWriter {
|
||||
written: file_size(&path),
|
||||
file,
|
||||
path,
|
||||
}),
|
||||
Err(e) => {
|
||||
tracing::warn!("[unified_log] failed to open log file: {e}");
|
||||
None
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
fn write_lines(lines: &[u8]) {
|
||||
let Ok(mut guard) = WRITER.lock() else { return };
|
||||
let writer = match guard.as_mut() {
|
||||
Some(w) => w,
|
||||
None => return,
|
||||
};
|
||||
|
||||
let len = lines.len() as u64;
|
||||
if let Err(e) = writer.file.write_all(lines) {
|
||||
tracing::warn!("[unified_log] write failed: {e}");
|
||||
return;
|
||||
}
|
||||
writer.written += len;
|
||||
|
||||
// Trim under the lock to avoid a race where concurrent writers see stale
|
||||
// state between drop + re-acquire. Trim is fast (~2.5 MB read+write) and
|
||||
// this is a low-volume diagnostic log.
|
||||
if writer.written >= MAX_SIZE {
|
||||
let _ = writer.file.flush();
|
||||
trim_file(&writer.path);
|
||||
if let Ok(new_file) = OpenOptions::new()
|
||||
.create(true)
|
||||
.append(true)
|
||||
.open(&writer.path)
|
||||
{
|
||||
writer.file = new_file;
|
||||
writer.written = file_size(&writer.path);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
fn write_entry(entry: &LogEntry) {
|
||||
let Ok(mut line) = serde_json::to_vec(entry) else {
|
||||
return;
|
||||
};
|
||||
line.push(b'\n');
|
||||
write_lines(&line);
|
||||
}
|
||||
|
||||
/// Drop the oldest lines from the file, keeping roughly the last half.
|
||||
///
|
||||
/// Uses write-to-temp + rename so a crash mid-trim cannot lose the entire log.
|
||||
pub fn trim_file(path: &std::path::Path) {
|
||||
let Ok(data) = fs::read(path) else { return };
|
||||
let half = data.len() / 2;
|
||||
// Find the first newline after the halfway point so we don't split a line.
|
||||
let start = match data[half..].iter().position(|&b| b == b'\n') {
|
||||
Some(pos) => half + pos + 1,
|
||||
None => return,
|
||||
};
|
||||
let tmp = path.with_extension("jsonl.tmp");
|
||||
if fs::write(&tmp, &data[start..]).is_ok() {
|
||||
let _ = fs::rename(&tmp, path);
|
||||
}
|
||||
}
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// Public API
|
||||
// ---------------------------------------------------------------------------
|
||||
|
||||
/// Return a new timestamp string in the unified log format.
|
||||
fn now_ts() -> String {
|
||||
Utc::now().to_rfc3339_opts(chrono::SecondsFormat::Millis, true)
|
||||
}
|
||||
|
||||
/// Emit a log entry from shell itself.
|
||||
pub fn emit(lvl: LogLevel, msg: &str, sid: Option<&str>, ctx: Option<serde_json::Value>) {
|
||||
let entry = LogEntry {
|
||||
ts: now_ts(),
|
||||
src: LogSource::Shell,
|
||||
pid: Some(std::process::id()),
|
||||
ver: VERSION.get().cloned(),
|
||||
lvl,
|
||||
sid: sid.map(Into::into),
|
||||
msg: msg.into(),
|
||||
ctx,
|
||||
};
|
||||
write_entry(&entry);
|
||||
}
|
||||
|
||||
/// Ingest a batch of log entries from a client (pager or desktop).
|
||||
///
|
||||
/// Called by the `x.ai/log` notification handler. Entries from
|
||||
/// [`LogSource::Shell`] are rejected to prevent spoofing.
|
||||
pub fn ingest_client_entries(src: LogSource, entries: &[ClientLogEntry]) {
|
||||
if matches!(src, LogSource::Shell) || entries.is_empty() {
|
||||
return;
|
||||
}
|
||||
// Serialize all entries up front, then write in a single lock acquisition.
|
||||
let mut buf = Vec::new();
|
||||
for client_entry in entries {
|
||||
let entry = LogEntry {
|
||||
ts: client_entry.ts.clone(),
|
||||
src,
|
||||
pid: client_entry.pid,
|
||||
ver: client_entry.ver.clone(),
|
||||
lvl: client_entry.lvl,
|
||||
sid: client_entry.sid.clone(),
|
||||
msg: client_entry.msg.clone(),
|
||||
ctx: client_entry.ctx.clone(),
|
||||
};
|
||||
if let Ok(mut line) = serde_json::to_vec(&entry) {
|
||||
line.push(b'\n');
|
||||
buf.extend_from_slice(&line);
|
||||
}
|
||||
}
|
||||
if !buf.is_empty() {
|
||||
write_lines(&buf);
|
||||
}
|
||||
}
|
||||
|
||||
/// Convenience: emit an info-level entry from shell.
|
||||
pub fn info(msg: &str, sid: Option<&str>, ctx: Option<serde_json::Value>) {
|
||||
emit(LogLevel::Info, msg, sid, ctx);
|
||||
}
|
||||
|
||||
/// Convenience: emit a warn-level entry from shell.
|
||||
pub fn warn(msg: &str, sid: Option<&str>, ctx: Option<serde_json::Value>) {
|
||||
emit(LogLevel::Warn, msg, sid, ctx);
|
||||
}
|
||||
|
||||
/// Convenience: emit an error-level entry from shell.
|
||||
pub fn error(msg: &str, sid: Option<&str>, ctx: Option<serde_json::Value>) {
|
||||
emit(LogLevel::Error, msg, sid, ctx);
|
||||
}
|
||||
|
||||
/// Convenience: emit a debug-level entry from shell.
|
||||
pub fn debug(msg: &str, sid: Option<&str>, ctx: Option<serde_json::Value>) {
|
||||
emit(LogLevel::Debug, msg, sid, ctx);
|
||||
}
|
||||
|
||||
/// Read the current unified log file and return its contents.
|
||||
///
|
||||
/// Returns `None` if the log file doesn't exist or can't be read.
|
||||
/// Used by diagnostic uploads to capture the log state at a point in time.
|
||||
pub fn snapshot_log() -> Option<Vec<u8>> {
|
||||
let path = log_path();
|
||||
// Flush pending writes before reading.
|
||||
if let Ok(mut guard) = WRITER.lock()
|
||||
&& let Some(ref mut w) = *guard
|
||||
{
|
||||
let _ = w.file.flush();
|
||||
}
|
||||
// Lock released intentionally — snapshot is approximate.
|
||||
match fs::read(&path) {
|
||||
Ok(data) if !data.is_empty() => Some(data),
|
||||
_ => None,
|
||||
}
|
||||
}
|
||||
|
||||
/// Read the unified log and return only entries belonging to the given session.
|
||||
///
|
||||
/// Parses each JSONL line, keeps entries where `"sid"` matches `session_id`,
|
||||
/// and returns the filtered lines as JSONL bytes. Returns `None` if the log
|
||||
/// is empty or contains no entries for this session.
|
||||
pub fn snapshot_session_log(session_id: &str) -> Option<Vec<u8>> {
|
||||
let path = log_path();
|
||||
if let Ok(mut guard) = WRITER.lock()
|
||||
&& let Some(ref mut w) = *guard
|
||||
{
|
||||
let _ = w.file.flush();
|
||||
}
|
||||
let data = match fs::read(&path) {
|
||||
Ok(d) if !d.is_empty() => d,
|
||||
_ => return None,
|
||||
};
|
||||
let mut out = Vec::new();
|
||||
for line in data.split(|&b| b == b'\n') {
|
||||
if line.is_empty() {
|
||||
continue;
|
||||
}
|
||||
if let Ok(entry) = serde_json::from_slice::<serde_json::Value>(line)
|
||||
&& entry.get("sid").and_then(|v| v.as_str()) == Some(session_id)
|
||||
{
|
||||
out.extend_from_slice(line);
|
||||
out.push(b'\n');
|
||||
}
|
||||
}
|
||||
if out.is_empty() { None } else { Some(out) }
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
#[test]
|
||||
fn log_entry_serializes_minimal() {
|
||||
let entry = LogEntry {
|
||||
ts: "2025-07-14T10:30:00.123Z".into(),
|
||||
src: LogSource::Shell,
|
||||
pid: None,
|
||||
ver: None,
|
||||
lvl: LogLevel::Info,
|
||||
sid: None,
|
||||
msg: "test".into(),
|
||||
ctx: None,
|
||||
};
|
||||
let json = serde_json::to_string(&entry).unwrap();
|
||||
assert!(!json.contains("sid"));
|
||||
assert!(!json.contains("ctx"));
|
||||
assert!(!json.contains("pid"));
|
||||
assert!(!json.contains("ver"));
|
||||
assert!(json.contains("\"src\":\"shell\""));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn log_entry_serializes_full() {
|
||||
let entry = LogEntry {
|
||||
ts: "2025-07-14T10:30:00.123Z".into(),
|
||||
src: LogSource::GrokPager,
|
||||
pid: Some(4242),
|
||||
ver: Some("0.1.211".into()),
|
||||
lvl: LogLevel::Warn,
|
||||
sid: Some("abc123".into()),
|
||||
msg: "connection lost".into(),
|
||||
ctx: Some(serde_json::json!({"retry": 3})),
|
||||
};
|
||||
let json = serde_json::to_string(&entry).unwrap();
|
||||
assert!(json.contains("\"sid\":\"abc123\""));
|
||||
assert!(json.contains("\"retry\":3"));
|
||||
assert!(json.contains("\"pid\":4242"));
|
||||
assert!(json.contains("\"ver\":\"0.1.211\""));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn client_entry_round_trip() {
|
||||
let wire = r#"{"ts":"2025-07-14T10:30:00.123Z","lvl":"info","msg":"hello"}"#;
|
||||
let entry: ClientLogEntry = serde_json::from_str(wire).unwrap();
|
||||
assert_eq!(entry.msg, "hello");
|
||||
assert!(entry.sid.is_none());
|
||||
assert!(entry.ctx.is_none());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn trim_file_keeps_recent_half() {
|
||||
let dir = tempfile::tempdir().unwrap();
|
||||
let path = dir.path().join("test.jsonl");
|
||||
let mut content = String::new();
|
||||
for i in 0..10 {
|
||||
content.push_str(&format!("line {i}\n"));
|
||||
}
|
||||
fs::write(&path, &content).unwrap();
|
||||
trim_file(&path);
|
||||
let result = fs::read_to_string(&path).unwrap();
|
||||
// Should keep roughly the second half, starting at a line boundary.
|
||||
assert!(!result.contains("line 0"));
|
||||
assert!(result.contains("line 9"));
|
||||
assert!(result.len() < content.len());
|
||||
// Every line should be complete (no partial lines).
|
||||
for line in result.lines() {
|
||||
assert!(line.starts_with("line "));
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn trim_file_no_newline_in_second_half_is_noop() {
|
||||
let dir = tempfile::tempdir().unwrap();
|
||||
let path = dir.path().join("test.jsonl");
|
||||
let content = "single-line-no-newline";
|
||||
fs::write(&path, content).unwrap();
|
||||
trim_file(&path);
|
||||
assert_eq!(fs::read_to_string(&path).unwrap(), content);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn trim_file_missing_file_is_noop() {
|
||||
let dir = tempfile::tempdir().unwrap();
|
||||
let path = dir.path().join("nonexistent.jsonl");
|
||||
trim_file(&path);
|
||||
assert!(!path.exists());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn ingest_rejects_shell_src() {
|
||||
ingest_client_entries(
|
||||
LogSource::Shell,
|
||||
&[ClientLogEntry {
|
||||
ts: "2025-01-01T00:00:00.000Z".into(),
|
||||
pid: None,
|
||||
ver: None,
|
||||
lvl: LogLevel::Info,
|
||||
sid: None,
|
||||
msg: "sneaky".into(),
|
||||
ctx: None,
|
||||
}],
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn unknown_src_rejected_at_deserialization() {
|
||||
for bad in &[
|
||||
r#"{"src":"evil","entries":[]}"#,
|
||||
r#"{"src":"","entries":[]}"#,
|
||||
r#"{"src":"GROK-PAGER","entries":[]}"#,
|
||||
] {
|
||||
assert!(serde_json::from_str::<LogNotificationParams>(bad).is_err());
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn notification_params_round_trip() {
|
||||
let params = LogNotificationParams {
|
||||
src: LogSource::GrokPager,
|
||||
entries: vec![
|
||||
ClientLogEntry {
|
||||
ts: "2025-07-14T10:30:00.123Z".into(),
|
||||
pid: Some(1234),
|
||||
ver: None,
|
||||
lvl: LogLevel::Info,
|
||||
sid: Some("s1".into()),
|
||||
msg: "first".into(),
|
||||
ctx: None,
|
||||
},
|
||||
ClientLogEntry {
|
||||
ts: "2025-07-14T10:30:00.456Z".into(),
|
||||
pid: Some(1234),
|
||||
ver: Some("0.1.211".into()),
|
||||
lvl: LogLevel::Error,
|
||||
sid: None,
|
||||
msg: "second".into(),
|
||||
ctx: Some(serde_json::json!({"code": 42})),
|
||||
},
|
||||
],
|
||||
};
|
||||
let json = serde_json::to_string(¶ms).unwrap();
|
||||
let parsed: LogNotificationParams = serde_json::from_str(&json).unwrap();
|
||||
assert_eq!(parsed.entries.len(), 2);
|
||||
assert_eq!(parsed.entries[0].msg, "first");
|
||||
assert_eq!(parsed.entries[1].msg, "second");
|
||||
}
|
||||
}
|
||||
Loading…
Reference in a new issue