ImageApi/tests/otel_llm_spans.rs

//! Verifies the LLM-client tracing migration (#140 / map #135).
//!
//! The `LlmClient` methods no longer take an OpenTelemetry `Context`; the
//! `llm.request` span parents on the caller's ambient `tracing` span. Two
//! properties are pinned here, in the production shape of the streaming chat
//! path (`insight_chat::chat_turn_stream_inner`: a `tokio::spawn` carrying
//! the turn span via `.instrument`, on a multi-thread runtime):
//!
//! 1. **Nesting.** Both the blocking and the streaming `llm.request` spans
//!    nest under the caller's span (same trace_id, parent link) and keep the
//!    `llm.provider` / `llm.operation` / `llm.model` / `http.status_code`
//!    attributes.
//! 2. **Propagation.** The outgoing `traceparent` header names the
//!    `llm.request` span itself, so the downstream server parents onto it.
//!
//! One process-global subscriber per test binary, so this file holds a
//! single test.

use futures::StreamExt;
use image_api::ai::llamacpp::LlamaCppClient;
use image_api::ai::llm_client::{ChatMessage, LlmClient, LlmStreamEvent};
use opentelemetry::Value;
use opentelemetry::trace::TracerProvider;
use opentelemetry_sdk::trace::{InMemorySpanExporter, SdkTracerProvider, SpanData};
use std::sync::{Arc, Mutex};
use tokio::io::{AsyncReadExt, AsyncWriteExt};
use tracing::Instrument;
use tracing_subscriber::layer::SubscriberExt;
use tracing_subscriber::registry::Registry;
use tracing_subscriber::util::SubscriberInitExt;

const CHAT_BODY: &str = r#"{"choices":[{"message":{"role":"assistant","content":"hi"}}],"usage":{"prompt_tokens":3,"completion_tokens":1}}"#;
const STREAM_BODY: &str = "data: {\"choices\":[{\"delta\":{\"role\":\"assistant\",\"content\":\"hi\"}}]}\n\n\
data: {\"choices\":[],\"usage\":{\"prompt_tokens\":3,\"completion_tokens\":1}}\n\n\
data: [DONE]\n\n";

/// Minimal OpenAI-compatible server: records each request's `traceparent`
/// and answers with an SSE stream when the body asks for one, JSON otherwise.
async fn serve(listener: tokio::net::TcpListener, seen: Arc<Mutex<Vec<String>>>) {
    loop {
        let (mut sock, _) = listener.accept().await.expect("accept");
        let seen = seen.clone();
        tokio::spawn(async move {
            let mut buf = Vec::new();
            let mut chunk = [0u8; 4096];
            let header_end = loop {
                let n = sock.read(&mut chunk).await.expect("read");
                buf.extend_from_slice(&chunk[..n]);
                if let Some(i) = buf.windows(4).position(|w| w == b"\r\n\r\n") {
                    break i + 4;
                }
            };
            let head = String::from_utf8_lossy(&buf[..header_end]).to_lowercase();
            let content_length: usize = head
                .lines()
                .find_map(|l| l.strip_prefix("content-length:"))
                .map(|v| v.trim().parse().unwrap())
                .unwrap_or(0);
            while buf.len() < header_end + content_length {
                let n = sock.read(&mut chunk).await.expect("read body");
                buf.extend_from_slice(&chunk[..n]);
            }
            if let Some(tp) = head.lines().find_map(|l| l.strip_prefix("traceparent:")) {
                seen.lock().unwrap().push(tp.trim().to_string());
            }
            let body = String::from_utf8_lossy(&buf[header_end..]);
            let (ctype, payload) = if body.contains("\"stream\":true") {
                ("text/event-stream", STREAM_BODY)
            } else {
                ("application/json", CHAT_BODY)
            };
            let resp = format!(
                "HTTP/1.1 200 OK\r\ncontent-type: {ctype}\r\ncontent-length: {}\r\nconnection: close\r\n\r\n{payload}",
                payload.len()
            );
            sock.write_all(resp.as_bytes()).await.expect("write");
        });
    }
}

fn attr(span: &SpanData, key: &str) -> Option<Value> {
    span.attributes
        .iter()
        .find(|kv| kv.key.as_str() == key)
        .map(|kv| kv.value.clone())
}

#[test]
fn llm_spans_nest_under_caller_across_spawn() {
    let exporter = InMemorySpanExporter::default();
    let provider = SdkTracerProvider::builder()
        .with_simple_exporter(exporter.clone())
        .build();
    Registry::default()
        .with(tracing_opentelemetry::OpenTelemetryLayer::new(
            provider.tracer("llm-span-test"),
        ))
        .try_init()
        .expect("global subscriber not yet set");

    let rt = tokio::runtime::Builder::new_multi_thread()
        .worker_threads(2)
        .enable_all()
        .build()
        .unwrap();

    let seen = Arc::new(Mutex::new(Vec::new()));
    rt.block_on(async {
        let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap();
        let base = format!("http://{}", listener.local_addr().unwrap());
        tokio::spawn(serve(listener, seen.clone()));

        let client = Arc::new(LlamaCppClient::new(Some(base), Some("chat".into())));
        let turn = tracing::info_span!("test.turn");
        tokio::spawn(
            async move {
                let (reply, _, _) = client
                    .chat_with_tools(vec![ChatMessage::user("hello")], vec![])
                    .await
                    .expect("chat");
                assert_eq!(reply.content, "hi");

                let mut stream = client
                    .chat_with_tools_stream(vec![ChatMessage::user("hello")], vec![])
                    .await
                    .expect("stream");
                let mut done = false;
                while let Some(ev) = stream.next().await {
                    if let LlmStreamEvent::Done { .. } = ev.expect("stream event") {
                        done = true;
                    }
                }
                assert!(done, "stream ended without Done");
            }
            .instrument(turn),
        )
        .await
        .expect("spawned turn");
    });
    let _ = provider.force_flush();

    let spans = exporter.get_finished_spans().unwrap();
    let turn = spans
        .iter()
        .find(|s| s.name == "test.turn")
        .expect("turn span exported");
    let llm: Vec<&SpanData> = spans.iter().filter(|s| s.name == "llm.request").collect();
    assert_eq!(llm.len(), 2, "one llm.request per HTTP call");

    let seen = seen.lock().unwrap().clone();
    assert_eq!(seen.len(), 2, "every LLM request carries a traceparent");

    for (span, op) in llm.iter().zip(["chat", "chat_stream"]) {
        assert_eq!(
            span.span_context.trace_id(),
            turn.span_context.trace_id(),
            "{op}: llm.request must share the caller's trace"
        );
        assert_eq!(
            span.parent_span_id,
            turn.span_context.span_id(),
            "{op}: llm.request must be a child of the caller span"
        );
        assert_eq!(span.span_kind, opentelemetry::trace::SpanKind::Client);
        assert_eq!(attr(span, "llm.provider"), Some("llamacpp".into()));
        assert_eq!(attr(span, "llm.operation"), Some(op.into()));
        assert_eq!(attr(span, "llm.model"), Some("chat".into()));
        assert_eq!(attr(span, "http.status_code"), Some(200i64.into()));

        // Downstream parents onto the llm.request span itself.
        let expected = format!(
            "00-{}-{}-01",
            span.span_context.trace_id(),
            span.span_context.span_id()
        );
        assert!(
            seen.contains(&expected),
            "{op}: traceparent {expected} not sent (saw {seen:?})"
        );
    }
}