Genarrative/server-rs/crates/platform-llm/src/observability_tests.rs

use std::{
    collections::BTreeMap,
    io::{Read, Write},
    net::TcpListener,
    sync::{Arc, Mutex, mpsc},
    thread,
    time::Duration,
};

use tokio::sync::oneshot;
use tracing::{
    Instrument, Subscriber,
    field::{Field, Visit},
    instrument::WithSubscriber,
    span::{Attributes, Id, Record},
};
use tracing_subscriber::{Layer, layer::Context, prelude::*, registry::LookupSpan};

use super::{LlmClient, LlmConfig, LlmError, LlmProvider, LlmRunRequest};

const PRIVATE_INPUT: &str = "PRIVATE_MESSAGE_MUST_NOT_ENTER_SPAN";
const PRIVATE_KEY: &str = "PRIVATE_API_KEY_MUST_NOT_ENTER_SPAN";

#[derive(Clone, Debug, Default)]
struct CapturedSpan {
    name: String,
    parent: Option<u64>,
    fields: BTreeMap<String, String>,
    closed: bool,
}

#[derive(Clone, Default)]
struct Capture {
    spans: Arc<Mutex<BTreeMap<u64, CapturedSpan>>>,
    delta_parents: Arc<Mutex<Vec<u64>>>,
}

struct Fields<'a>(&'a mut BTreeMap<String, String>);

impl Visit for Fields<'_> {
    fn record_debug(&mut self, field: &Field, value: &dyn std::fmt::Debug) {
        self.0.insert(field.name().into(), format!("{value:?}"));
    }

    fn record_str(&mut self, field: &Field, value: &str) {
        self.0.insert(field.name().into(), value.into());
    }
}

impl<S: Subscriber + for<'a> LookupSpan<'a>> Layer<S> for Capture {
    fn on_new_span(&self, attributes: &Attributes<'_>, id: &Id, context: Context<'_, S>) {
        let span = context.span(id).expect("registered span");
        let mut captured = CapturedSpan {
            name: attributes.metadata().name().into(),
            parent: span.parent().map(|parent| parent.id().into_u64()),
            ..CapturedSpan::default()
        };
        attributes.record(&mut Fields(&mut captured.fields));
        self.spans.lock().unwrap().insert(id.into_u64(), captured);
    }

    fn on_record(&self, id: &Id, values: &Record<'_>, _: Context<'_, S>) {
        let mut spans = self.spans.lock().unwrap();
        values.record(&mut Fields(
            &mut spans.get_mut(&id.into_u64()).unwrap().fields,
        ));
    }

    fn on_close(&self, id: Id, _: Context<'_, S>) {
        self.spans
            .lock()
            .unwrap()
            .get_mut(&id.into_u64())
            .unwrap()
            .closed = true;
    }

    fn on_event(&self, event: &tracing::Event<'_>, context: Context<'_, S>) {
        if event.metadata().target() == "llm_observability_test_delta" {
            let parent = context
                .event_span(event)
                .expect("delta should have a parent");
            self.delta_parents
                .lock()
                .unwrap()
                .push(parent.id().into_u64());
        }
    }
}

impl Capture {
    fn provider_span(&self) -> (u64, CapturedSpan) {
        let spans = self.spans.lock().unwrap();
        let provider_spans = spans
            .iter()
            .filter(|(_, span)| span.name == "llm.request")
            .collect::<Vec<_>>();
        assert_eq!(provider_spans.len(), 1);
        let (id, span) = provider_spans[0];
        (*id, span.clone())
    }

    fn assert_completed(&self, operation: &str) {
        let (_, span) = self.provider_span();
        assert!(
            span.closed,
            "provider span must close after completion or cancellation"
        );
        assert_eq!(
            span.fields.get("operation").map(String::as_str),
            Some(operation)
        );
        assert_eq!(
            span.fields.get("provider").map(String::as_str),
            Some("openai_compatible")
        );
        assert_eq!(
            span.fields.get("api_kind").map(String::as_str),
            Some("openai_chat")
        );
        assert_eq!(
            span.fields.get("model").map(String::as_str),
            Some("requested-model")
        );
        let spans = self.spans.lock().unwrap();
        assert_eq!(
            spans[&span.parent.expect("HTTP/application parent")].name,
            "test.request"
        );
        let captured = format!("{spans:?}");
        assert!(!captured.contains(PRIVATE_INPUT));
        assert!(!captured.contains(PRIVATE_KEY));
        assert!(!span.fields.contains_key("self"));
        assert!(!span.fields.contains_key("request"));
    }
}

struct ProviderFixture {
    client: LlmClient,
    entered: oneshot::Receiver<()>,
    release: mpsc::Sender<()>,
    server: thread::JoinHandle<()>,
}

// 使用本地替代上游和通道控制响应时点，不依赖实际 Provider 或计时猜测。
fn provider_fixture(status: &str, content_type: &str, body: &str) -> ProviderFixture {
    let listener = TcpListener::bind("127.0.0.1:0").unwrap();
    let address = listener.local_addr().unwrap();
    let response = format!(
        "HTTP/1.1 {status}\r\nContent-Type: {content_type}\r\nContent-Length: {}\r\nConnection: close\r\n\r\n{body}",
        body.len()
    );
    let (entered, entered_rx) = oneshot::channel();
    let (release, release_rx) = mpsc::channel();
    let server = thread::spawn(move || {
        listener.set_nonblocking(true).unwrap();
        let deadline = std::time::Instant::now() + Duration::from_secs(10);
        let mut stream = loop {
            match listener.accept() {
                Ok((stream, _)) => break stream,
                Err(error) if error.kind() == std::io::ErrorKind::WouldBlock => {
                    assert!(
                        std::time::Instant::now() < deadline,
                        "request did not arrive"
                    );
                    thread::sleep(Duration::from_millis(5));
                }
                Err(error) => panic!("fixture accept: {error}"),
            }
        };
        stream.set_nonblocking(false).unwrap();
        stream
            .set_read_timeout(Some(Duration::from_secs(5)))
            .unwrap();
        let mut bytes = Vec::new();
        let mut chunk = [0; 4096];
        loop {
            let count = stream.read(&mut chunk).unwrap();
            assert_ne!(count, 0);
            bytes.extend_from_slice(&chunk[..count]);
            if let Some(headers_end) = bytes.windows(4).position(|part| part == b"\r\n\r\n") {
                let headers = String::from_utf8_lossy(&bytes[..headers_end]);
                let length = headers
                    .lines()
                    .find_map(|line| {
                        let (name, value) = line.split_once(':')?;
                        name.eq_ignore_ascii_case("content-length")
                            .then(|| value.trim().parse::<usize>().unwrap())
                    })
                    .unwrap_or(0);
                if bytes.len() >= headers_end + 4 + length {
                    break;
                }
            }
        }
        entered.send(()).unwrap();
        release_rx.recv_timeout(Duration::from_secs(10)).unwrap();
        // 取消用例已关闭客户端连接，允许写回失败。
        let _ = stream.write_all(response.as_bytes());
    });
    let config = LlmConfig::new(
        LlmProvider::OpenAiCompatible,
        format!("http://{address}"),
        PRIVATE_KEY.into(),
        "default-model".into(),
        5_000,
        0,
        1,
    )
    .unwrap();
    ProviderFixture {
        client: LlmClient::new(config).unwrap(),
        entered: entered_rx,
        release,
        server,
    }
}

fn request() -> LlmRunRequest {
    LlmRunRequest::single_turn("system", PRIVATE_INPUT)
        .with_openai_chat()
        .with_model("requested-model")
}

#[tokio::test]
async fn provider_span_covers_awaited_execution_and_keeps_parent_without_arguments() {
    let fixture = provider_fixture(
        "200 OK",
        "application/json",
        r#"{"choices":[{"message":{"content":"completed"},"finish_reason":"stop"}]}"#,
    );
    let capture = Capture::default();
    let response =
        async {
            async {
            let mut future = Box::pin(fixture.client.run(request()));
            assert!(capture.spans.lock().unwrap().values().all(|span| span.name != "llm.request"));
            tokio::select! {
                response = &mut future => panic!("request returned before released: {response:?}"),
                entered = fixture.entered => entered.unwrap(),
            }
            assert!(!capture.provider_span().1.closed);
            // 挂起时不得把 Provider span 留在当前异步执行上下文。
            assert_eq!(tracing::Span::current().metadata().unwrap().name(), "test.request");
            fixture.release.send(()).unwrap();
            future.await.unwrap()
        }.instrument(tracing::info_span!("test.request")).await
        }
        .with_subscriber(tracing_subscriber::registry().with(capture.clone()))
        .await;
    fixture.server.join().unwrap();
    assert_eq!(response.text, "completed");
    capture.assert_completed("run");
}

#[tokio::test]
async fn stream_callbacks_inherit_provider_span_and_keep_result() {
    let fixture = provider_fixture(
        "200 OK",
        "text/event-stream",
        "data: {\"choices\":[{\"delta\":{\"content\":\"hello\"},\"finish_reason\":null}]}\n\ndata: {\"choices\":[{\"delta\":{},\"finish_reason\":\"stop\"}]}\n\ndata: [DONE]\n\n",
    );
    let capture = Capture::default();
    let mut deltas = Vec::new();
    let response =
        async {
            async {
            let mut future = Box::pin(fixture.client.stream_run(request(), |delta| {
                tracing::info!(target: "llm_observability_test_delta", "delta received");
                deltas.push(delta.delta_text.clone());
            }));
            tokio::select! {
                response = &mut future => panic!("stream returned before released: {response:?}"),
                entered = fixture.entered => entered.unwrap(),
            }
            fixture.release.send(()).unwrap();
            future.await.unwrap()
        }.instrument(tracing::info_span!("test.request")).await
        }
        .with_subscriber(tracing_subscriber::registry().with(capture.clone()))
        .await;
    fixture.server.join().unwrap();
    assert_eq!(response.text, "hello");
    assert_eq!(deltas.concat(), "hello");
    capture.assert_completed("stream_run");
    let id = capture.provider_span().0;
    let parents = capture.delta_parents.lock().unwrap();
    assert!(!parents.is_empty());
    assert!(parents.iter().all(|parent| *parent == id));
}

#[tokio::test]
async fn provider_span_preserves_upstream_errors_and_closes_on_cancellation() {
    let fixture = provider_fixture(
        "401 Unauthorized",
        "application/json",
        r#"{"error":{"message":"upstream-rejected"}}"#,
    );
    let capture = Capture::default();
    let error = async {
        async {
            fixture.release.send(()).unwrap();
            fixture.client.run(request()).await.unwrap_err()
        }
        .instrument(tracing::info_span!("test.request"))
        .await
    }
    .with_subscriber(tracing_subscriber::registry().with(capture.clone()))
    .await;
    fixture.server.join().unwrap();
    assert!(
        matches!(error, LlmError::Upstream { status_code: 401, message } if message.contains("upstream-rejected"))
    );
    capture.assert_completed("run");

    let fixture = provider_fixture("200 OK", "application/json", "{}");
    let capture = Capture::default();
    async {
        async {
            let mut future = Box::pin(fixture.client.run(request()));
            tokio::select! {
                response = &mut future => panic!("request returned before cancellation: {response:?}"),
                entered = fixture.entered => entered.unwrap(),
            }
            assert!(!capture.provider_span().1.closed);
            drop(future);
            fixture.release.send(()).unwrap();
        }.instrument(tracing::info_span!("test.request")).await
    }.with_subscriber(tracing_subscriber::registry().with(capture.clone())).await;
    fixture.server.join().unwrap();
    capture.assert_completed("run");
}