From 2cc39d0fffed005e319b0143e5f5a17f5d8cc562 Mon Sep 17 00:00:00 2001 From: Ashish Kolhe <35160958+oasisk@users.noreply.github.com> Date: Thu, 30 Jul 2026 13:22:57 +0530 Subject: [PATCH] fix: persist canonical gen_ai agent env/version on UDS streams (#13544) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit ## Problem AI Observability pages fail with `Search field not found: Schema error: No field named gen_ai_agent_version` when an agent version/env filter is applied (cascade picker), on any LLM trace stream that has a user-defined schema. Root cause chain (verified live on introspect against `sre_agent_traces_production_eu`): 1. The agent registry resolves env/version from source attributes (e.g. `service_service_version`) and writes canonical `gen_ai_agent_env` / `gen_ai_agent_version` onto the span at ingest. 2. On UDS streams, `refactor_map` rebuilds the record keeping only `defined_schema_fields` — which never included the two new columns — so they are silently dropped. 3. `restore_canonical_agent_fields` exists to protect canonical agent fields from exactly this, but only restored name/id (env/version write-back landed later and this path was never extended). Result: the registry advertises env/version variants, the picker offers them, but the column never reaches the stream schema — every filtered query errors. ## Fix - **`GEN_AI_SCHEMA_FIELDS`** now includes `gen_ai_agent_env` and `gen_ai_agent_version`. This provisions the columns (Arrow schema + UDS field list) through the existing paths: - existing LLM streams: `ensure_gen_ai_fields_in_schema` runs on every OTLP ingest → self-heals on the next span, no migration needed - new streams: `set_stream_is_llm` provisions on first LLM detection - merging into the Arrow schema alone already stops the hard query error (old spans read as NULL) - **`restore_canonical_agent_fields`** now carries env/version via a `CanonicalAgentFields` struct and re-inserts them after the UDS refactor — covers the first-ever LLM batch, where the stream is only marked LLM at end of request so the UDS list doesn't have the fields yet. ## Testing - New unit test `test_restore_canonical_agent_fields_restores_env_and_version` - Extended `test_append_gen_ai_fields_to_defined_schema_fields_adds_migrated_fields` for the two new fields (kept the `_o2_ingest_ts` assertion from #13506) - `cargo check` / `clippy` / `fmt` clean on `openobserve-core` and `db` - End-to-end write-back verified against a live enterprise build: probe span with `service.version` + `deployment.environment.name` resource attrs came back queryable as `gen_ai_agent_version` / `gen_ai_agent_env` --- src/core/src/traces/mod.rs | 65 ++++++++++++++++++++++++++++++++++---- src/db/src/schema.rs | 7 ++++ 2 files changed, 66 insertions(+), 6 deletions(-) diff --git a/src/core/src/traces/mod.rs b/src/core/src/traces/mod.rs index e447c164cd..9609598f8e 100644 --- a/src/core/src/traces/mod.rs +++ b/src/core/src/traces/mod.rs @@ -141,6 +141,16 @@ type AgentObservationBuffer = std::collections::BTreeMap< #[cfg(not(feature = "enterprise"))] type AgentObservationBuffer = Option; +/// Canonical agent-identity fields the registry wrote onto the span, captured +/// before UDS refactoring so they can be re-inserted if the stream's UDS field +/// list does not (yet) include them (see [`restore_canonical_agent_fields`]). +struct CanonicalAgentFields { + agent_name: Option, + agent_id: Option, + env: Option, + version: Option, +} + fn normalize_llm_field_types(record_val: &mut Map) { for &field in GEN_AI_INT64_FIELDS.iter() { if let Some(value) = record_val.get_mut(field) @@ -178,7 +188,7 @@ fn collect_gen_ai_agent_observation( record_val: &mut Map, mapping_config: &GenAiAgentMappingConfig, observations: &mut AgentObservationBuffer, -) -> Option<(Option, Option)> { +) -> Option { #[cfg(feature = "enterprise")] { let observation = @@ -190,7 +200,12 @@ fn collect_gen_ai_agent_observation( record_val, mapping_config, )?; - let canonical_fields = (observation.agent_name.clone(), observation.agent_id.clone()); + let canonical_fields = CanonicalAgentFields { + agent_name: observation.agent_name.clone(), + agent_id: observation.agent_id.clone(), + env: observation.env.clone(), + version: observation.version.clone(), + }; let agent_key = observation.agent_key.clone(); let identity_source = observation.identity_source.clone(); let buffer_size_before = observations.len(); @@ -223,18 +238,24 @@ fn collect_gen_ai_agent_observation( fn restore_canonical_agent_fields( record_val: &mut Map, - canonical_fields: Option<(Option, Option)>, + canonical_fields: Option, ) { - let Some((agent_name, agent_id)) = canonical_fields else { + let Some(fields) = canonical_fields else { return; }; - if let Some(agent_name) = agent_name { + if let Some(agent_name) = fields.agent_name { record_val.insert("gen_ai_agent_name".to_string(), json::json!(agent_name)); } - if let Some(agent_id) = agent_id { + if let Some(agent_id) = fields.agent_id { record_val.insert("gen_ai_agent_id".to_string(), json::json!(agent_id)); } + if let Some(env) = fields.env { + record_val.insert("gen_ai_agent_env".to_string(), json::json!(env)); + } + if let Some(version) = fields.version { + record_val.insert("gen_ai_agent_version".to_string(), json::json!(version)); + } } #[cfg(feature = "enterprise")] @@ -1694,6 +1715,38 @@ mod tests { ); } + #[test] + fn test_restore_canonical_agent_fields_restores_env_and_version() { + // Simulates a UDS stream whose field list dropped the canonical agent + // columns during refactor_map: every canonical field must be restored, + // env/version included, so version-scoped queries can filter on them. + let mut record = json!({"_timestamp": 1_i64}).as_object().unwrap().clone(); + + super::restore_canonical_agent_fields( + &mut record, + Some(super::CanonicalAgentFields { + agent_name: Some("o2_ai_agent".to_string()), + agent_id: None, + env: Some("production".to_string()), + version: Some("0.1.0".to_string()), + }), + ); + + assert_eq!( + record.get("gen_ai_agent_name").and_then(|v| v.as_str()), + Some("o2_ai_agent") + ); + assert!(!record.contains_key("gen_ai_agent_id")); + assert_eq!( + record.get("gen_ai_agent_env").and_then(|v| v.as_str()), + Some("production") + ); + assert_eq!( + record.get("gen_ai_agent_version").and_then(|v| v.as_str()), + Some("0.1.0") + ); + } + #[test] fn test_normalize_llm_field_types() { let mut record = json!({ diff --git a/src/db/src/schema.rs b/src/db/src/schema.rs index 94503d2026..6f87c41826 100644 --- a/src/db/src/schema.rs +++ b/src/db/src/schema.rs @@ -215,6 +215,11 @@ static GEN_AI_SCHEMA_FIELDS: std::sync::LazyLock> = std::sync::LazyLo Field::new("gen_ai_provider_name", DataType::Utf8, true), Field::new("gen_ai_agent_name", DataType::Utf8, true), Field::new("gen_ai_agent_id", DataType::Utf8, true), + // Canonical env/version columns the agent registry writes back onto + // spans; must be provisioned (and UDS-listed) so version-scoped agent + // queries never hit "No field named" on LLM streams. + Field::new("gen_ai_agent_env", DataType::Utf8, true), + Field::new("gen_ai_agent_version", DataType::Utf8, true), Field::new("gen_ai_input_messages", DataType::Utf8, true), Field::new("gen_ai_output_messages", DataType::Utf8, true), Field::new("gen_ai_system_instructions", DataType::Utf8, true), @@ -830,6 +835,8 @@ mod tests { assert!(fields.contains(&"gen_ai_usage_cache_read_input_tokens".to_string())); assert!(fields.contains(&"gen_ai_usage_cache_creation_input_tokens".to_string())); assert!(fields.contains(&"gen_ai_usage_cost_net_cache_impact".to_string())); + assert!(fields.contains(&"gen_ai_agent_env".to_string())); + assert!(fields.contains(&"gen_ai_agent_version".to_string())); assert!(fields.contains(&O2_INGEST_TS_COL_NAME.to_string())); }