codex/codex-rs/code-mode-host/src/main.rs

use std::collections::BTreeMap;
use std::collections::HashMap;
use std::path::PathBuf;
use std::time::Duration;

use anyhow::Context;
use clap::Parser;
use codex_otel::OtelExporter;
use codex_otel::OtelHttpProtocol;
use codex_otel::OtelProvider;
use codex_otel::OtelSettings;
use tracing_subscriber::Layer;
use tracing_subscriber::layer::SubscriberExt;
use tracing_subscriber::util::SubscriberInitExt;

const OTEL_SHUTDOWN_TIMEOUT: Duration = Duration::from_secs(/*secs*/ 5);

#[derive(Debug, Parser)]
struct Cli {
    /// Transport endpoint: `stdio`, `stdio://`, or `grpc://IP:PORT`.
    #[arg(
        long,
        value_name = "URL",
        default_value = codex_code_mode_host::DEFAULT_LISTEN_URL
    )]
    listen: String,

    /// Optional WebSocket endpoint that streams only raw OTLP trace batches.
    #[arg(long, value_name = "URL")]
    otel_trace_listen: Option<String>,

    /// Optional OTLP/HTTP JSON trace exporter endpoint, analogous to
    /// `otel.trace_exporter` in app-server configuration.
    #[arg(long, value_name = "URL", conflicts_with = "otel_trace_listen")]
    otel_trace_exporter: Option<String>,
}

#[tokio::main(flavor = "current_thread")]
async fn main() -> anyhow::Result<()> {
    let cli = Cli::parse();
    let trace_transport = if let Some(trace_listen) = cli.otel_trace_listen.as_deref() {
        let sender = codex_code_mode_host::trace_batch_channel();
        let (receiver, exporter_endpoint) =
            codex_code_mode_host::bind_otlp_trace_receiver().await?;
        Some((
            trace_listen.to_string(),
            receiver,
            sender,
            exporter_endpoint,
        ))
    } else {
        None
    };
    let trace_exporter_endpoint = trace_transport
        .as_ref()
        .map(|(_, _, _, endpoint)| endpoint.as_str())
        .or(cli.otel_trace_exporter.as_deref());
    let otel = trace_exporter_endpoint
        .map(build_trace_provider)
        .transpose()?;
    let otel_layer = otel.as_ref().and_then(OtelProvider::tracing_layer);
    tracing_subscriber::registry()
        .with(
            tracing_subscriber::fmt::layer()
                .with_writer(std::io::stderr)
                .with_ansi(false)
                .with_filter(tracing_subscriber::filter::LevelFilter::INFO),
        )
        .with(otel_layer)
        .init();
    tracing::info_span!(
        "code_mode_host.startup",
        otel.name = "code_mode_host.startup"
    )
    .in_scope(|| {});

    let main_transport = codex_code_mode_host::run_main(&cli.listen);
    let (result, trace_tasks) = match trace_transport {
        Some((trace_listen, receiver, sender, _)) => {
            let mut trace_listener_task = tokio::spawn({
                let sender = sender.clone();
                async move { codex_code_mode_host::run_otel_trace_listener(&trace_listen, sender).await }
            });
            let mut trace_receiver_task = tokio::spawn(
                codex_code_mode_host::run_otlp_trace_receiver(receiver, sender),
            );
            let result = tokio::select! {
                result = main_transport => result,
                result = &mut trace_listener_task => result.context("OTEL trace websocket task failed").flatten(),
                result = &mut trace_receiver_task => result.context("OTLP trace receiver task failed").flatten(),
            };
            (result, Some((trace_listener_task, trace_receiver_task)))
        }
        None => (main_transport.await, None),
    };
    if let Some(otel) = otel
        && let Err(error) = otel.shutdown_with_timeout(OTEL_SHUTDOWN_TIMEOUT).await
    {
        tracing::warn!(%error, "failed to finish code-mode host telemetry shutdown");
    }
    if let Some((trace_listener_task, trace_receiver_task)) = trace_tasks {
        trace_listener_task.abort();
        trace_receiver_task.abort();
    }
    result
}

fn build_trace_provider(endpoint: &str) -> anyhow::Result<OtelProvider> {
    OtelProvider::try_new(&OtelSettings {
        environment: "code-mode-host".to_string(),
        service_name: "codex-code-mode-host".to_string(),
        service_version: env!("CARGO_PKG_VERSION").to_string(),
        codex_home: PathBuf::from("/tmp"),
        exporter: OtelExporter::None,
        trace_exporter: OtelExporter::OtlpHttp {
            endpoint: endpoint.to_string(),
            headers: HashMap::new(),
            protocol: OtelHttpProtocol::Json,
            tls: None,
        },
        metrics_exporter: OtelExporter::None,
        runtime_metrics: false,
        span_attributes: BTreeMap::new(),
        tracestate: BTreeMap::new(),
    })
    .map_err(|error| anyhow::anyhow!("failed to build code-mode host OTEL provider: {error}"))?
    .context("code-mode host OTEL trace provider was unexpectedly disabled")
}