codex/codex-rs/otel/tests/suite/otel_export_routing_policy.rs

use codex_api::AgentIdentityTelemetry;
use codex_otel::AuthEnvTelemetryMetadata;
use codex_otel::OtelProvider;
use codex_otel::SessionTelemetry;
use codex_otel::TelemetryAuthMode;
use opentelemetry::KeyValue;
use opentelemetry::logs::AnyValue;
use opentelemetry::trace::TracerProvider as _;
use opentelemetry_sdk::logs::InMemoryLogExporter;
use opentelemetry_sdk::logs::SdkLogRecord;
use opentelemetry_sdk::logs::SdkLoggerProvider;
use opentelemetry_sdk::trace::InMemorySpanExporter;
use opentelemetry_sdk::trace::SdkTracerProvider;
use pretty_assertions::assert_eq;
use std::borrow::Cow;
use std::collections::BTreeMap;
use std::path::PathBuf;
use tracing_subscriber::Layer;
use tracing_subscriber::filter::filter_fn;
use tracing_subscriber::layer::SubscriberExt;

use codex_protocol::AgentPath;
use codex_protocol::ThreadId;
use codex_protocol::ToolName;
use codex_protocol::config_types::ReasoningSummary;
use codex_protocol::config_types::ToolResultLogConfig;
use codex_protocol::models::ImageReference;
use codex_protocol::protocol::AskForApproval;
use codex_protocol::protocol::GuardianAssessmentEvent;
use codex_protocol::protocol::GuardianAssessmentOutcome;
use codex_protocol::protocol::GuardianAssessmentStatus;
use codex_protocol::protocol::SandboxPolicy;
use codex_protocol::protocol::SessionSource;
use codex_protocol::protocol::SubAgentSource;
use codex_protocol::user_input::UserInput;

fn log_attributes(record: &SdkLogRecord) -> BTreeMap<String, String> {
    record
        .attributes_iter()
        .map(|(key, value)| (key.as_str().to_string(), any_value_to_string(value)))
        .collect()
}

fn span_event_attributes(event: &opentelemetry::trace::Event) -> BTreeMap<String, String> {
    event
        .attributes
        .iter()
        .map(|KeyValue { key, value, .. }| (key.as_str().to_string(), value.to_string()))
        .collect()
}

fn any_value_to_string(value: &AnyValue) -> String {
    match value {
        AnyValue::Int(value) => value.to_string(),
        AnyValue::Double(value) => value.to_string(),
        AnyValue::String(value) => value.as_str().to_string(),
        AnyValue::Boolean(value) => value.to_string(),
        AnyValue::Bytes(value) => String::from_utf8_lossy(value).into_owned(),
        AnyValue::ListAny(value) => format!("{value:?}"),
        AnyValue::Map(value) => format!("{value:?}"),
        _ => format!("{value:?}"),
    }
}

fn find_log_by_event_name<'a>(
    logs: &'a [opentelemetry_sdk::logs::in_memory_exporter::LogDataWithResource],
    event_name: &str,
) -> &'a opentelemetry_sdk::logs::in_memory_exporter::LogDataWithResource {
    logs.iter()
        .find(|log| {
            log_attributes(&log.record)
                .get("event.name")
                .is_some_and(|value| value == event_name)
        })
        .expect("log event should exist")
}

fn find_span_event_by_name_attr<'a>(
    events: &'a [opentelemetry::trace::Event],
    event_name: &str,
) -> &'a opentelemetry::trace::Event {
    events
        .iter()
        .find(|event| {
            span_event_attributes(event)
                .get("event.name")
                .is_some_and(|value| value == event_name)
        })
        .expect("span event should exist")
}

fn auth_env_metadata() -> AuthEnvTelemetryMetadata {
    AuthEnvTelemetryMetadata {
        openai_api_key_env_present: true,
        codex_api_key_env_present: false,
        codex_api_key_env_enabled: true,
        provider_env_key_name: Some("configured".to_string()),
        provider_env_key_present: Some(true),
        refresh_token_url_override_present: true,
    }
}

#[test]
fn otel_export_routing_policy_routes_prompt_and_guardian_payloads() {
    let log_exporter = InMemoryLogExporter::default();
    let logger_provider = SdkLoggerProvider::builder()
        .with_simple_exporter(log_exporter.clone())
        .build();
    let span_exporter = InMemorySpanExporter::default();
    let tracer_provider = SdkTracerProvider::builder()
        .with_simple_exporter(span_exporter.clone())
        .build();
    let tracer = tracer_provider.tracer("sink-split-test");

    let subscriber = tracing_subscriber::registry()
        .with(
            opentelemetry_appender_tracing::layer::OpenTelemetryTracingBridge::new(
                &logger_provider,
            )
            .with_filter(filter_fn(OtelProvider::log_export_filter)),
        )
        .with(
            tracing_opentelemetry::layer()
                .with_tracer(tracer)
                .with_filter(filter_fn(OtelProvider::trace_export_filter)),
        );

    tracing::subscriber::with_default(subscriber, || {
        tracing::callsite::rebuild_interest_cache();
        let manager = SessionTelemetry::new(
            ThreadId::new(),
            "gpt-5.1",
            "gpt-5.1",
            Some("account-id".to_string()),
            Some("[email protected]".to_string()),
            Some(TelemetryAuthMode::ApiKey),
            "codex_exec".to_string(),
            /*log_user_prompts*/ true,
            "tty".to_string(),
            SessionSource::Cli,
        );
        let root_span = tracing::info_span!("root");
        let _root_guard = root_span.enter();
        manager.user_prompt(&[
            UserInput::Text {
                text: "super secret prompt".to_string(),
                text_elements: Vec::new(),
            },
            UserInput::Image {
                image: ImageReference::Inline {
                    image_url: "https://example.com/image.png".to_string(),
                },
                detail: None,
            },
            UserInput::LocalImage {
                path: PathBuf::from("/tmp/secret.png"),
                detail: None,
            },
        ]);
        let mut assessment: GuardianAssessmentEvent = serde_json::from_value(serde_json::json!({
            "id": "review",
            "turn_id": "turn",
            "target_item_id": "tool-call",
            "status": "in_progress",
            "risk_level": "medium",
            "user_authorization": "high",
            "rationale": format!("{}🦀", "x".repeat(/*n*/ 65_535)),
            "action": {"type": "mcp_tool_call", "server": "private-server", "tool_name": "tool"},
        }))
        .expect("assessment");
        manager.guardian_assessment(&assessment, /*outcome*/ None);
        for (status, outcome) in [
            (
                GuardianAssessmentStatus::Approved,
                Some(GuardianAssessmentOutcome::Allow),
            ),
            (
                GuardianAssessmentStatus::Denied,
                Some(GuardianAssessmentOutcome::Deny),
            ),
            (GuardianAssessmentStatus::Denied, None),
            (GuardianAssessmentStatus::TimedOut, None),
            (GuardianAssessmentStatus::Aborted, None),
        ] {
            assessment.status = status;
            if outcome.is_none() {
                assessment.rationale = None;
            }
            manager.guardian_assessment(&assessment, outcome);
        }
    });

    logger_provider.force_flush().expect("flush logs");
    tracer_provider.force_flush().expect("flush traces");

    let logs = log_exporter.get_emitted_logs().expect("log export");
    assert!(
        logs.iter()
            .all(|log| { log.record.target().map(Cow::as_ref) == Some("codex_otel.log_only") })
    );

    let prompt_log = find_log_by_event_name(&logs, "codex.user_prompt");
    let prompt_log_attrs = log_attributes(&prompt_log.record);
    assert_eq!(
        prompt_log_attrs.get("prompt").map(String::as_str),
        Some("super secret prompt")
    );
    assert_eq!(
        prompt_log_attrs.get("user.email").map(String::as_str),
        Some("[email protected]")
    );

    let assessments: Vec<_> = logs
        .iter()
        .map(|log| log_attributes(&log.record))
        .filter(|attrs| {
            attrs.get("event.name").map(String::as_str) == Some("codex.guardian_assessment")
        })
        .collect();
    assert_eq!(
        assessments
            .iter()
            .map(|attrs| (
                attrs["status"].as_str(),
                attrs.get("outcome").map(String::as_str),
                attrs.get("rationale").map(String::len),
                attrs.get("rationale_length").map(String::as_str),
                attrs.get("rationale_truncated").map(String::as_str),
            ))
            .collect::<Vec<_>>(),
        vec![
            (
                "approved",
                Some("allow"),
                Some(65_535),
                Some("65539"),
                Some("true")
            ),
            (
                "denied",
                Some("deny"),
                Some(65_535),
                Some("65539"),
                Some("true")
            ),
            ("denied", None, None, None, None),
            ("timed_out", None, None, None, None),
            ("aborted", None, None, None, None),
        ],
    );
    assert!(!format!("{assessments:?}").contains("private-server"));

    let spans = span_exporter.get_finished_spans().expect("span export");
    assert_eq!(spans.len(), 1);
    let span_events = &spans[0].events.events;
    assert_eq!(span_events.len(), 1);

    let prompt_trace_event = find_span_event_by_name_attr(span_events, "codex.user_prompt");
    let prompt_trace_attrs = span_event_attributes(prompt_trace_event);
    assert_eq!(
        prompt_trace_attrs.get("prompt_length").map(String::as_str),
        Some("19")
    );
    assert_eq!(
        prompt_trace_attrs
            .get("text_input_count")
            .map(String::as_str),
        Some("1")
    );
    assert_eq!(
        prompt_trace_attrs
            .get("image_input_count")
            .map(String::as_str),
        Some("1")
    );
    assert_eq!(
        prompt_trace_attrs
            .get("local_image_input_count")
            .map(String::as_str),
        Some("1")
    );
    assert!(!prompt_trace_attrs.contains_key("prompt"));
    assert!(!prompt_trace_attrs.contains_key("user.email"));
    assert!(!prompt_trace_attrs.contains_key("user.account_id"));
}

#[test]
fn otel_export_routing_policy_routes_tool_result_log_and_trace_events() {
    let output = "secret output\nsecond line\n".repeat(100);
    let root_id = ThreadId::new();
    let child_id = ThreadId::new();
    let legacy_id = ThreadId::new();
    let unnamed_id = ThreadId::new();
    let make_manager = |thread_id, source| {
        SessionTelemetry::new(
            thread_id,
            "gpt-5.1",
            "gpt-5.1",
            Some("account-id".to_string()),
            Some("[email protected]".to_string()),
            Some(TelemetryAuthMode::ApiKey),
            "codex_exec".to_string(),
            /*log_user_prompts*/ true,
            "tty".to_string(),
            source,
        )
    };
    let log_exporter = InMemoryLogExporter::default();
    let logger_provider = SdkLoggerProvider::builder()
        .with_simple_exporter(log_exporter.clone())
        .build();
    let span_exporter = InMemorySpanExporter::default();
    let tracer_provider = SdkTracerProvider::builder()
        .with_simple_exporter(span_exporter.clone())
        .build();
    let tracer = tracer_provider.tracer("sink-split-test");

    let subscriber = tracing_subscriber::registry()
        .with(
            opentelemetry_appender_tracing::layer::OpenTelemetryTracingBridge::new(
                &logger_provider,
            )
            .with_filter(filter_fn(OtelProvider::log_export_filter)),
        )
        .with(
            tracing_opentelemetry::layer()
                .with_tracer(tracer)
                .with_filter(filter_fn(OtelProvider::trace_export_filter)),
        );

    tracing::subscriber::with_default(subscriber, || {
        tracing::callsite::rebuild_interest_cache();
        let manager = make_manager(root_id, SessionSource::Cli).with_product_sku(Some("codex"));
        let root_span = tracing::info_span!("root");
        let _root_guard = root_span.enter();
        let manager = manager.with_tool_result_log_config(ToolResultLogConfig {
            max_bytes: output.len(),
        });
        manager.tool_result_with_tags(
            &ToolName::namespaced("mcp__example", "shell"),
            "call-1",
            "secret arguments",
            std::time::Duration::from_millis(42),
            /*success*/ true,
            &output,
            &[],
            &[
                ("mcp_server", "internal-mcp"),
                ("mcp_server_origin", "stdio"),
            ],
        );
        let child = make_manager(
            child_id,
            SessionSource::SubAgent(SubAgentSource::ThreadSpawn {
                parent_thread_id: root_id,
                depth: 1,
                agent_path: Some(AgentPath::try_from("/root/reviewer").expect("agent path")),
                agent_nickname: Some("legacy nickname".to_string()),
                agent_role: None,
            }),
        )
        .with_product_sku(Some("customer-specific-value"));
        child.tool_result_with_tags(
            &ToolName::plain("shell"),
            "call-2",
            "{}",
            std::time::Duration::ZERO,
            /*success*/ true,
            &output,
            &[],
            &[],
        );
        let legacy = make_manager(
            legacy_id,
            SessionSource::SubAgent(SubAgentSource::ThreadSpawn {
                parent_thread_id: root_id,
                depth: 1,
                agent_path: None,
                agent_nickname: Some("Mendel".to_string()),
                agent_role: None,
            }),
        );
        let unnamed = make_manager(unnamed_id, SessionSource::SubAgent(SubAgentSource::Review));
        for (manager, name, error) in [
            (&manager, "root_failure", "failure"),
            (&legacy, "legacy_failure", output.as_str()),
            (&unnamed, "unnamed_failure", "failure"),
        ] {
            manager.tool_result_with_tags(
                &ToolName::plain(name),
                name,
                "{}",
                std::time::Duration::ZERO,
                /*success*/ false,
                error,
                &[],
                &[],
            );
        }
    });

    logger_provider.force_flush().expect("flush logs");
    tracer_provider.force_flush().expect("flush traces");

    let logs = log_exporter.get_emitted_logs().expect("log export");
    assert!(
        logs.iter()
            .all(|log| { log.record.target().map(Cow::as_ref) == Some("codex_otel.log_only") })
    );

    let tool_log = find_log_by_event_name(&logs, "codex.tool_result");
    let tool_log_attrs = log_attributes(&tool_log.record);
    assert_eq!(
        tool_log_attrs.get("tool_name").map(String::as_str),
        Some("shell")
    );
    assert_eq!(
        tool_log_attrs.get("tool_namespace").map(String::as_str),
        Some("mcp__example")
    );
    assert_eq!(
        tool_log_attrs.get("arguments").map(String::as_str),
        Some("secret arguments")
    );
    assert_eq!(
        tool_log_attrs.get("output").map(String::as_str),
        Some(output.as_str())
    );
    let tool_logs: Vec<_> = logs
        .iter()
        .map(|log| log_attributes(&log.record))
        .filter(|attrs| attrs.get("event.name").map(String::as_str) == Some("codex.tool_result"))
        .collect();
    assert_eq!(
        tool_logs
            .iter()
            .map(|attrs| attrs.get("product_sku").map(String::as_str))
            .collect::<Vec<_>>(),
        vec![Some("codex"), Some("other"), Some("codex"), None, None],
    );
    assert_eq!(
        tool_logs
            .iter()
            .map(|attrs| (
                attrs["conversation.id"].clone(),
                attrs["agent_name"].clone(),
                attrs["output_truncated"].clone(),
            ))
            .collect::<Vec<_>>(),
        vec![
            (
                root_id.to_string(),
                "/root".to_string(),
                "false".to_string()
            ),
            (
                child_id.to_string(),
                "/root/reviewer".to_string(),
                "true".to_string()
            ),
            (
                root_id.to_string(),
                "/root".to_string(),
                "false".to_string()
            ),
            (
                legacy_id.to_string(),
                "Mendel".to_string(),
                "true".to_string()
            ),
            (
                unnamed_id.to_string(),
                unnamed_id.to_string(),
                "false".to_string()
            ),
        ]
    );
    assert!(tool_logs[1]["output"].ends_with("[... telemetry preview truncated ...]"));
    let sequences: Vec<u64> = tool_logs
        .iter()
        .map(|attrs| attrs["tool_result_seq"].parse().expect("numeric sequence"))
        .collect();
    assert!(sequences[0] > 0);
    assert!(sequences.windows(2).all(|pair| pair[0] < pair[1]));
    assert_eq!(
        tool_log_attrs.get("mcp_server").map(String::as_str),
        Some("internal-mcp")
    );
    assert_eq!(
        tool_log_attrs.get("mcp_server_origin").map(String::as_str),
        Some("stdio")
    );

    let spans = span_exporter.get_finished_spans().expect("span export");
    assert_eq!(spans.len(), 1);
    let span_events = &spans[0].events.events;
    assert_eq!(span_events.len(), tool_logs.len());
    for (event, log) in span_events.iter().zip(&tool_logs) {
        let attrs = span_event_attributes(event);
        assert_eq!(attrs.get("tool_name"), log.get("tool_name"));
        assert_eq!(attrs.get("tool_namespace"), log.get("tool_namespace"));
        assert_eq!(attrs.get("tool_result_seq"), log.get("tool_result_seq"));
        assert_eq!(attrs.get("product_sku"), log.get("product_sku"));
        assert_eq!(attrs.get("output_truncated"), log.get("output_truncated"));
        assert!(!attrs.contains_key("agent_name"));
    }

    let tool_trace_event = find_span_event_by_name_attr(span_events, "codex.tool_result");
    let tool_trace_attrs = span_event_attributes(tool_trace_event);
    assert_eq!(
        tool_trace_attrs.get("arguments_length").map(String::as_str),
        Some("16")
    );
    assert_eq!(
        tool_trace_attrs.get("output_length").map(String::as_str),
        Some(output.len().to_string().as_str())
    );
    assert_eq!(
        tool_trace_attrs
            .get("output_line_count")
            .map(String::as_str),
        Some(output.lines().count().to_string().as_str())
    );
    assert!(!tool_trace_attrs.contains_key("arguments"));
    assert!(!tool_trace_attrs.contains_key("output"));
    assert!(!tool_trace_attrs.contains_key("mcp_server"));
    assert!(!tool_trace_attrs.contains_key("mcp_server_origin"));
}

#[test]
fn otel_export_routing_policy_routes_auth_recovery_log_and_trace_events() {
    let log_exporter = InMemoryLogExporter::default();
    let logger_provider = SdkLoggerProvider::builder()
        .with_simple_exporter(log_exporter.clone())
        .build();
    let span_exporter = InMemorySpanExporter::default();
    let tracer_provider = SdkTracerProvider::builder()
        .with_simple_exporter(span_exporter.clone())
        .build();
    let tracer = tracer_provider.tracer("sink-split-test");

    let subscriber = tracing_subscriber::registry()
        .with(
            opentelemetry_appender_tracing::layer::OpenTelemetryTracingBridge::new(
                &logger_provider,
            )
            .with_filter(filter_fn(OtelProvider::log_export_filter)),
        )
        .with(
            tracing_opentelemetry::layer()
                .with_tracer(tracer)
                .with_filter(filter_fn(OtelProvider::trace_export_filter)),
        );

    tracing::subscriber::with_default(subscriber, || {
        tracing::callsite::rebuild_interest_cache();
        let manager = SessionTelemetry::new(
            ThreadId::new(),
            "gpt-5.1",
            "gpt-5.1",
            Some("account-id".to_string()),
            Some("[email protected]".to_string()),
            Some(TelemetryAuthMode::Chatgpt),
            "codex_exec".to_string(),
            /*log_user_prompts*/ true,
            "tty".to_string(),
            SessionSource::Cli,
        );
        let root_span = tracing::info_span!("root");
        let _root_guard = root_span.enter();
        manager.record_auth_recovery(
            "managed",
            "reload",
            "recovery_succeeded",
            Some("req-401"),
            Some("ray-401"),
            Some("missing_authorization_header"),
            Some("token_expired"),
            /*recovery_reason*/ None,
            Some(true),
        );
    });

    logger_provider.force_flush().expect("flush logs");
    tracer_provider.force_flush().expect("flush traces");

    let logs = log_exporter.get_emitted_logs().expect("log export");
    let recovery_log = find_log_by_event_name(&logs, "codex.auth_recovery");
    let recovery_log_attrs = log_attributes(&recovery_log.record);
    assert_eq!(
        recovery_log_attrs.get("auth.mode").map(String::as_str),
        Some("managed")
    );
    assert_eq!(
        recovery_log_attrs.get("auth.step").map(String::as_str),
        Some("reload")
    );
    assert_eq!(
        recovery_log_attrs.get("auth.outcome").map(String::as_str),
        Some("recovery_succeeded")
    );
    assert_eq!(
        recovery_log_attrs
            .get("auth.request_id")
            .map(String::as_str),
        Some("req-401")
    );
    assert_eq!(
        recovery_log_attrs.get("auth.cf_ray").map(String::as_str),
        Some("ray-401")
    );
    assert_eq!(
        recovery_log_attrs.get("auth.error").map(String::as_str),
        Some("missing_authorization_header")
    );
    assert_eq!(
        recovery_log_attrs
            .get("auth.error_code")
            .map(String::as_str),
        Some("token_expired")
    );
    assert_eq!(
        recovery_log_attrs
            .get("auth.state_changed")
            .map(String::as_str),
        Some("true")
    );

    let spans = span_exporter.get_finished_spans().expect("span export");
    assert_eq!(spans.len(), 1);
    let span_events = &spans[0].events.events;
    assert_eq!(span_events.len(), 1);

    let recovery_trace_event = find_span_event_by_name_attr(span_events, "codex.auth_recovery");
    let recovery_trace_attrs = span_event_attributes(recovery_trace_event);
    assert_eq!(
        recovery_trace_attrs.get("auth.mode").map(String::as_str),
        Some("managed")
    );
    assert_eq!(
        recovery_trace_attrs.get("auth.step").map(String::as_str),
        Some("reload")
    );
    assert_eq!(
        recovery_trace_attrs.get("auth.outcome").map(String::as_str),
        Some("recovery_succeeded")
    );
    assert_eq!(
        recovery_trace_attrs
            .get("auth.request_id")
            .map(String::as_str),
        Some("req-401")
    );
    assert_eq!(
        recovery_trace_attrs.get("auth.cf_ray").map(String::as_str),
        Some("ray-401")
    );
    assert_eq!(
        recovery_trace_attrs.get("auth.error").map(String::as_str),
        Some("missing_authorization_header")
    );
    assert_eq!(
        recovery_trace_attrs
            .get("auth.error_code")
            .map(String::as_str),
        Some("token_expired")
    );
    assert_eq!(
        recovery_trace_attrs
            .get("auth.state_changed")
            .map(String::as_str),
        Some("true")
    );
}

#[test]
fn otel_export_routing_policy_routes_api_request_auth_observability() {
    let log_exporter = InMemoryLogExporter::default();
    let logger_provider = SdkLoggerProvider::builder()
        .with_simple_exporter(log_exporter.clone())
        .build();
    let span_exporter = InMemorySpanExporter::default();
    let tracer_provider = SdkTracerProvider::builder()
        .with_simple_exporter(span_exporter.clone())
        .build();
    let tracer = tracer_provider.tracer("sink-split-test");

    let subscriber = tracing_subscriber::registry()
        .with(
            opentelemetry_appender_tracing::layer::OpenTelemetryTracingBridge::new(
                &logger_provider,
            )
            .with_filter(filter_fn(OtelProvider::log_export_filter)),
        )
        .with(
            tracing_opentelemetry::layer()
                .with_tracer(tracer)
                .with_filter(filter_fn(OtelProvider::trace_export_filter)),
        );

    tracing::subscriber::with_default(subscriber, || {
        tracing::callsite::rebuild_interest_cache();
        let manager = SessionTelemetry::new(
            ThreadId::new(),
            "gpt-5.1",
            "gpt-5.1",
            Some("account-id".to_string()),
            Some("[email protected]".to_string()),
            Some(TelemetryAuthMode::Chatgpt),
            "codex_exec".to_string(),
            /*log_user_prompts*/ true,
            "tty".to_string(),
            SessionSource::Cli,
        )
        .with_auth_env(auth_env_metadata());
        let root_span = tracing::info_span!("root");
        let _root_guard = root_span.enter();
        manager.conversation_starts(
            "openai",
            /*reasoning_effort*/ None,
            ReasoningSummary::Auto,
            /*context_window*/ None,
            /*auto_compact_token_limit*/ None,
            AskForApproval::Never,
            SandboxPolicy::DangerFullAccess,
            Vec::new(),
        );
        let agent_identity_telemetry = AgentIdentityTelemetry {
            agent_id: "agent-runtime-otel".to_string(),
            task_id: "task-run-otel".to_string(),
        };
        manager.record_api_request(
            /*attempt*/ 1,
            Some(401),
            Some("http 401"),
            std::time::Duration::from_millis(42),
            /*auth_header_attached*/ true,
            Some("authorization"),
            /*retry_after_unauthorized*/ true,
            Some("managed"),
            Some("refresh_token"),
            "/responses",
            Some("req-401"),
            Some("ray-401"),
            Some("missing_authorization_header"),
            Some("token_expired"),
            Some(&agent_identity_telemetry),
        );
    });

    logger_provider.force_flush().expect("flush logs");
    tracer_provider.force_flush().expect("flush traces");

    let logs = log_exporter.get_emitted_logs().expect("log export");
    let conversation_log = find_log_by_event_name(&logs, "codex.conversation_starts");
    let conversation_log_attrs = log_attributes(&conversation_log.record);
    assert_eq!(
        conversation_log_attrs
            .get("auth.env_openai_api_key_present")
            .map(String::as_str),
        Some("true")
    );
    assert_eq!(
        conversation_log_attrs
            .get("auth.env_provider_key_name")
            .map(String::as_str),
        Some("configured")
    );
    let request_log = find_log_by_event_name(&logs, "codex.api_request");
    let request_log_attrs = log_attributes(&request_log.record);
    assert_eq!(
        request_log_attrs
            .get("auth.header_attached")
            .map(String::as_str),
        Some("true")
    );
    assert_eq!(
        request_log_attrs
            .get("auth.header_name")
            .map(String::as_str),
        Some("authorization")
    );
    assert_eq!(
        request_log_attrs
            .get("auth.retry_after_unauthorized")
            .map(String::as_str),
        Some("true")
    );
    assert_eq!(
        request_log_attrs
            .get("auth.recovery_mode")
            .map(String::as_str),
        Some("managed")
    );
    assert_eq!(
        request_log_attrs
            .get("auth.recovery_phase")
            .map(String::as_str),
        Some("refresh_token")
    );
    assert_eq!(
        request_log_attrs.get("endpoint").map(String::as_str),
        Some("/responses")
    );
    assert_eq!(
        request_log_attrs.get("auth.error").map(String::as_str),
        Some("missing_authorization_header")
    );
    assert_eq!(
        request_log_attrs
            .get("auth.env_codex_api_key_enabled")
            .map(String::as_str),
        Some("true")
    );
    assert_eq!(
        request_log_attrs
            .get("auth.env_refresh_token_url_override_present")
            .map(String::as_str),
        Some("true")
    );
    assert_eq!(
        request_log_attrs.get("auth.agent_id").map(String::as_str),
        Some("agent-runtime-otel")
    );
    assert_eq!(
        request_log_attrs.get("auth.task_id").map(String::as_str),
        Some("task-run-otel")
    );

    let spans = span_exporter.get_finished_spans().expect("span export");
    let conversation_trace_event =
        find_span_event_by_name_attr(&spans[0].events.events, "codex.conversation_starts");
    let conversation_trace_attrs = span_event_attributes(conversation_trace_event);
    assert_eq!(
        conversation_trace_attrs
            .get("auth.env_provider_key_present")
            .map(String::as_str),
        Some("true")
    );
    let request_trace_event =
        find_span_event_by_name_attr(&spans[0].events.events, "codex.api_request");
    let request_trace_attrs = span_event_attributes(request_trace_event);
    assert_eq!(
        request_trace_attrs
            .get("auth.header_attached")
            .map(String::as_str),
        Some("true")
    );
    assert_eq!(
        request_trace_attrs
            .get("auth.header_name")
            .map(String::as_str),
        Some("authorization")
    );
    assert_eq!(
        request_trace_attrs
            .get("auth.retry_after_unauthorized")
            .map(String::as_str),
        Some("true")
    );
    assert_eq!(
        request_trace_attrs.get("endpoint").map(String::as_str),
        Some("/responses")
    );
    assert_eq!(
        request_trace_attrs
            .get("auth.env_openai_api_key_present")
            .map(String::as_str),
        Some("true")
    );
    assert_eq!(
        request_trace_attrs.get("auth.agent_id").map(String::as_str),
        Some("agent-runtime-otel")
    );
    assert_eq!(
        request_trace_attrs.get("auth.task_id").map(String::as_str),
        Some("task-run-otel")
    );
}

#[test]
fn otel_export_routing_policy_routes_websocket_connect_auth_observability() {
    let log_exporter = InMemoryLogExporter::default();
    let logger_provider = SdkLoggerProvider::builder()
        .with_simple_exporter(log_exporter.clone())
        .build();
    let span_exporter = InMemorySpanExporter::default();
    let tracer_provider = SdkTracerProvider::builder()
        .with_simple_exporter(span_exporter.clone())
        .build();
    let tracer = tracer_provider.tracer("sink-split-test");

    let subscriber = tracing_subscriber::registry()
        .with(
            opentelemetry_appender_tracing::layer::OpenTelemetryTracingBridge::new(
                &logger_provider,
            )
            .with_filter(filter_fn(OtelProvider::log_export_filter)),
        )
        .with(
            tracing_opentelemetry::layer()
                .with_tracer(tracer)
                .with_filter(filter_fn(OtelProvider::trace_export_filter)),
        );

    tracing::subscriber::with_default(subscriber, || {
        tracing::callsite::rebuild_interest_cache();
        let manager = SessionTelemetry::new(
            ThreadId::new(),
            "gpt-5.1",
            "gpt-5.1",
            Some("account-id".to_string()),
            Some("[email protected]".to_string()),
            Some(TelemetryAuthMode::Chatgpt),
            "codex_exec".to_string(),
            /*log_user_prompts*/ true,
            "tty".to_string(),
            SessionSource::Cli,
        )
        .with_auth_env(auth_env_metadata());
        let root_span = tracing::info_span!("root");
        let _root_guard = root_span.enter();
        let agent_identity_telemetry = AgentIdentityTelemetry {
            agent_id: "agent-runtime-ws".to_string(),
            task_id: "task-run-ws".to_string(),
        };
        manager.record_websocket_connect(
            std::time::Duration::from_millis(17),
            Some(401),
            Some("http 401"),
            /*auth_header_attached*/ true,
            Some("authorization"),
            /*retry_after_unauthorized*/ true,
            Some("managed"),
            Some("reload"),
            "/responses",
            /*connection_reused*/ false,
            Some("req-ws-401"),
            Some("ray-ws-401"),
            Some("missing_authorization_header"),
            Some("token_expired"),
            Some(&agent_identity_telemetry),
        );
    });

    logger_provider.force_flush().expect("flush logs");
    tracer_provider.force_flush().expect("flush traces");

    let logs = log_exporter.get_emitted_logs().expect("log export");
    let connect_log = find_log_by_event_name(&logs, "codex.websocket_connect");
    let connect_log_attrs = log_attributes(&connect_log.record);
    assert_eq!(
        connect_log_attrs
            .get("auth.header_attached")
            .map(String::as_str),
        Some("true")
    );
    assert_eq!(
        connect_log_attrs
            .get("auth.header_name")
            .map(String::as_str),
        Some("authorization")
    );
    assert_eq!(
        connect_log_attrs.get("auth.error").map(String::as_str),
        Some("missing_authorization_header")
    );
    assert_eq!(
        connect_log_attrs.get("endpoint").map(String::as_str),
        Some("/responses")
    );
    assert_eq!(
        connect_log_attrs
            .get("auth.connection_reused")
            .map(String::as_str),
        Some("false")
    );
    assert_eq!(
        connect_log_attrs
            .get("auth.env_provider_key_name")
            .map(String::as_str),
        Some("configured")
    );
    assert_eq!(
        connect_log_attrs.get("auth.agent_id").map(String::as_str),
        Some("agent-runtime-ws")
    );
    assert_eq!(
        connect_log_attrs.get("auth.task_id").map(String::as_str),
        Some("task-run-ws")
    );

    let spans = span_exporter.get_finished_spans().expect("span export");
    let connect_trace_event =
        find_span_event_by_name_attr(&spans[0].events.events, "codex.websocket_connect");
    let connect_trace_attrs = span_event_attributes(connect_trace_event);
    assert_eq!(
        connect_trace_attrs
            .get("auth.recovery_phase")
            .map(String::as_str),
        Some("reload")
    );
    assert_eq!(
        connect_trace_attrs
            .get("auth.env_refresh_token_url_override_present")
            .map(String::as_str),
        Some("true")
    );
    assert_eq!(
        connect_trace_attrs.get("auth.agent_id").map(String::as_str),
        Some("agent-runtime-ws")
    );
    assert_eq!(
        connect_trace_attrs.get("auth.task_id").map(String::as_str),
        Some("task-run-ws")
    );
}

#[test]
fn otel_export_routing_policy_routes_websocket_request_transport_observability() {
    let log_exporter = InMemoryLogExporter::default();
    let logger_provider = SdkLoggerProvider::builder()
        .with_simple_exporter(log_exporter.clone())
        .build();
    let span_exporter = InMemorySpanExporter::default();
    let tracer_provider = SdkTracerProvider::builder()
        .with_simple_exporter(span_exporter.clone())
        .build();
    let tracer = tracer_provider.tracer("sink-split-test");

    let subscriber = tracing_subscriber::registry()
        .with(
            opentelemetry_appender_tracing::layer::OpenTelemetryTracingBridge::new(
                &logger_provider,
            )
            .with_filter(filter_fn(OtelProvider::log_export_filter)),
        )
        .with(
            tracing_opentelemetry::layer()
                .with_tracer(tracer)
                .with_filter(filter_fn(OtelProvider::trace_export_filter)),
        );

    tracing::subscriber::with_default(subscriber, || {
        tracing::callsite::rebuild_interest_cache();
        let manager = SessionTelemetry::new(
            ThreadId::new(),
            "gpt-5.1",
            "gpt-5.1",
            Some("account-id".to_string()),
            Some("[email protected]".to_string()),
            Some(TelemetryAuthMode::Chatgpt),
            "codex_exec".to_string(),
            /*log_user_prompts*/ true,
            "tty".to_string(),
            SessionSource::Cli,
        )
        .with_auth_env(auth_env_metadata());
        let root_span = tracing::info_span!("root");
        let _root_guard = root_span.enter();
        let agent_identity_telemetry = AgentIdentityTelemetry {
            agent_id: "agent-runtime-ws-request".to_string(),
            task_id: "task-run-ws-request".to_string(),
        };
        manager.record_websocket_request(
            std::time::Duration::from_millis(23),
            Some("stream error"),
            /*connection_reused*/ true,
            Some(&agent_identity_telemetry),
        );
    });

    logger_provider.force_flush().expect("flush logs");
    tracer_provider.force_flush().expect("flush traces");

    let logs = log_exporter.get_emitted_logs().expect("log export");
    let request_log = find_log_by_event_name(&logs, "codex.websocket_request");
    let request_log_attrs = log_attributes(&request_log.record);
    assert_eq!(
        request_log_attrs
            .get("auth.connection_reused")
            .map(String::as_str),
        Some("true")
    );
    assert_eq!(
        request_log_attrs.get("error.message").map(String::as_str),
        Some("stream error")
    );
    assert_eq!(
        request_log_attrs
            .get("auth.env_openai_api_key_present")
            .map(String::as_str),
        Some("true")
    );
    assert_eq!(
        request_log_attrs.get("auth.agent_id").map(String::as_str),
        Some("agent-runtime-ws-request")
    );
    assert_eq!(
        request_log_attrs.get("auth.task_id").map(String::as_str),
        Some("task-run-ws-request")
    );

    let spans = span_exporter.get_finished_spans().expect("span export");
    let request_trace_event =
        find_span_event_by_name_attr(&spans[0].events.events, "codex.websocket_request");
    let request_trace_attrs = span_event_attributes(request_trace_event);
    assert_eq!(
        request_trace_attrs
            .get("auth.connection_reused")
            .map(String::as_str),
        Some("true")
    );
    assert_eq!(
        request_trace_attrs
            .get("auth.env_provider_key_present")
            .map(String::as_str),
        Some("true")
    );
    assert_eq!(
        request_trace_attrs.get("auth.agent_id").map(String::as_str),
        Some("agent-runtime-ws-request")
    );
    assert_eq!(
        request_trace_attrs.get("auth.task_id").map(String::as_str),
        Some("task-run-ws-request")
    );
}