attributes() {
+ return attributes;
+ }
+ }
+}
diff --git a/java/traceai-spring-ai/src/main/java/ai/traceai/spring/StreamOutputCapture.java b/java/traceai-spring-ai/src/main/java/ai/traceai/spring/StreamOutputCapture.java
new file mode 100644
index 00000000..0f1fe2dc
--- /dev/null
+++ b/java/traceai-spring-ai/src/main/java/ai/traceai/spring/StreamOutputCapture.java
@@ -0,0 +1,82 @@
+package ai.traceai.spring;
+
+import ai.traceai.FITracer;
+import ai.traceai.TraceConfig;
+import io.opentelemetry.api.trace.Span;
+import org.springframework.ai.chat.messages.AssistantMessage;
+import org.springframework.ai.chat.model.ChatResponse;
+import org.springframework.ai.chat.model.Generation;
+
+import java.util.Collections;
+
+/**
+ * Bounded per-subscription output capture for streamed chat responses.
+ *
+ * Allocates no text buffer when both output switches ({@code hideOutputs} and {@code hideOutputMessages})
+ * hide the content. Otherwise retains at most a 32,000 character prefix plus an overflow flag; the delivered
+ * stream elements are never modified or truncated. Token usage is never aggregated from chunks.
+ */
+final class StreamOutputCapture {
+
+ /** Matches the OpenTelemetry attribute budget used by {@link FITracer}. */
+ static final int MAX_CHARS = 32_000;
+ static final String TRUNCATION_SUFFIX = "...[truncated]";
+
+ private final FITracer tracer;
+ private final StringBuilder buffer;
+ private boolean overflow;
+
+ StreamOutputCapture(FITracer tracer) {
+ this.tracer = tracer;
+ TraceConfig config = tracer.getConfig();
+ boolean hideAll = config.isHideOutputs() && config.isHideOutputMessages();
+ this.buffer = hideAll ? null : new StringBuilder();
+ }
+
+ boolean isCapturing() {
+ return buffer != null;
+ }
+
+ boolean hasOverflow() {
+ return overflow;
+ }
+
+ void accept(ChatResponse response) {
+ if (buffer == null) {
+ return;
+ }
+ Generation generation = response.getResult();
+ if (generation == null) {
+ return;
+ }
+ AssistantMessage output = generation.getOutput();
+ if (output == null) {
+ return;
+ }
+ String text = output.getText();
+ if (text == null || text.isEmpty()) {
+ return;
+ }
+ int remaining = MAX_CHARS - buffer.length();
+ if (text.length() <= remaining) {
+ buffer.append(text);
+ } else {
+ if (remaining > 0) {
+ buffer.append(text, 0, remaining);
+ }
+ overflow = true;
+ }
+ }
+
+ /** Applies the captured output to the span; privacy switches are enforced by {@link FITracer}. */
+ void complete(Span span) {
+ if (buffer == null) {
+ return;
+ }
+ String text = overflow
+ ? buffer.substring(0, MAX_CHARS - TRUNCATION_SUFFIX.length()) + TRUNCATION_SUFFIX
+ : buffer.toString();
+ tracer.setOutputValue(span, text);
+ tracer.setOutputMessages(span, Collections.singletonList(FITracer.message("assistant", text)));
+ }
+}
diff --git a/java/traceai-spring-ai/src/main/java/ai/traceai/spring/TracedChatModel.java b/java/traceai-spring-ai/src/main/java/ai/traceai/spring/TracedChatModel.java
index 8df591a1..1e2acdfc 100644
--- a/java/traceai-spring-ai/src/main/java/ai/traceai/spring/TracedChatModel.java
+++ b/java/traceai-spring-ai/src/main/java/ai/traceai/spring/TracedChatModel.java
@@ -1,35 +1,55 @@
package ai.traceai.spring;
-import ai.traceai.*;
+import ai.traceai.FISpanKind;
+import ai.traceai.FITracer;
+import ai.traceai.SemanticConventions;
+import ai.traceai.TraceAI;
import io.opentelemetry.api.trace.Span;
import io.opentelemetry.api.trace.StatusCode;
+import io.opentelemetry.context.Context;
import io.opentelemetry.context.Scope;
+import org.springframework.ai.chat.messages.AssistantMessage;
+import org.springframework.ai.chat.messages.Message;
+import org.springframework.ai.chat.metadata.ChatResponseMetadata;
+import org.springframework.ai.chat.metadata.Usage;
import org.springframework.ai.chat.model.ChatModel;
import org.springframework.ai.chat.model.ChatResponse;
+import org.springframework.ai.chat.model.Generation;
import org.springframework.ai.chat.prompt.ChatOptions;
import org.springframework.ai.chat.prompt.Prompt;
-import org.springframework.ai.chat.messages.Message;
+import reactor.core.CoreSubscriber;
import reactor.core.publisher.Flux;
import java.util.ArrayList;
import java.util.Collections;
import java.util.List;
import java.util.Map;
+import java.util.Objects;
/**
- * Traced wrapper for Spring AI ChatModel.
- * Implements the ChatModel interface to provide automatic tracing.
+ * Traced wrapper for a Spring AI 1.1.x {@link ChatModel}.
+ *
+ * The wrapper passes the exact {@link Prompt} to the delegate and returns the delegate's exact
+ * {@link ChatResponse} (or stream elements). It never retries, substitutes or caches responses.
*
* Usage:
*
- * ChatModel model = new OpenAiChatModel(api);
- * ChatModel tracedModel = new TracedChatModel(model, "openai");
+ * ChatModel model = OllamaChatModel.builder()...build();
+ * ChatModel tracedModel = new TracedChatModel(model, tracer, "ollama");
*
* ChatResponse response = tracedModel.call(new Prompt("Hello!"));
*
+ *
+ * Streaming: {@link #stream(Prompt)} starts no span and does not call the delegate until a subscriber
+ * subscribes. Each subscription owns one span that ends exactly once on complete, error or cancel. Streamed
+ * token usage is not aggregated. To bind a parent span and TraceAI session/user attributes to a stream that
+ * is subscribed on another thread, use {@link SpringAITracingContext#withContext}.
*/
public class TracedChatModel implements ChatModel {
+ static final String SPAN_NAME_CALL = "Spring AI Chat";
+ static final String SPAN_NAME_STREAM = "Spring AI Chat (Stream)";
+
private final ChatModel delegate;
private final FITracer tracer;
private final String provider;
@@ -42,8 +62,8 @@ public class TracedChatModel implements ChatModel {
* @param provider the provider name (e.g., "openai", "anthropic")
*/
public TracedChatModel(ChatModel delegate, FITracer tracer, String provider) {
- this.delegate = delegate;
- this.tracer = tracer;
+ this.delegate = Objects.requireNonNull(delegate, "delegate ChatModel must not be null");
+ this.tracer = Objects.requireNonNull(tracer, "FITracer must not be null");
this.provider = provider;
}
@@ -52,6 +72,7 @@ public TracedChatModel(ChatModel delegate, FITracer tracer, String provider) {
*
* @param delegate the underlying model to wrap
* @param provider the provider name
+ * @throws IllegalStateException if {@link TraceAI#init} has not been called
*/
public TracedChatModel(ChatModel delegate, String provider) {
this(delegate, TraceAI.getTracer(), provider);
@@ -59,174 +80,201 @@ public TracedChatModel(ChatModel delegate, String provider) {
@Override
public ChatResponse call(Prompt prompt) {
- Span span = tracer.startSpan("Spring AI Chat", FISpanKind.LLM);
-
- try (Scope scope = span.makeCurrent()) {
- // Set system attributes
- span.setAttribute(SemanticConventions.LLM_SYSTEM, "spring-ai");
- span.setAttribute(SemanticConventions.LLM_PROVIDER, provider);
-
- // Capture input messages and input.value
- if (prompt != null && prompt.getInstructions() != null) {
- List messages = prompt.getInstructions();
- List> inputMessages = new ArrayList<>();
- StringBuilder inputValueBuilder = new StringBuilder();
- for (Message msg : messages) {
- String role = msg.getMessageType().getValue();
- String content = msg.getContent();
- inputMessages.add(FITracer.message(role, content));
- if (content != null) {
- if (inputValueBuilder.length() > 0) {
- inputValueBuilder.append("\n");
- }
- inputValueBuilder.append(content);
- }
- }
- tracer.setInputMessages(span, inputMessages);
- tracer.setInputValue(span, inputValueBuilder.toString());
- }
+ Span span = SpringAITracingContext.startSpan(tracer, SPAN_NAME_CALL, FISpanKind.LLM, SpringAITracingContext.capture());
+ boolean completed = false;
+ boolean errorRecorded = false;
+ try (Scope ignored = span.makeCurrent()) {
+ captureRequest(span, prompt);
- // Capture prompt options if available
- if (prompt != null && prompt.getOptions() != null) {
- var options = prompt.getOptions();
- if (options.getModel() != null) {
- span.setAttribute(SemanticConventions.LLM_REQUEST_MODEL, options.getModel());
- }
- if (options.getTemperature() != null) {
- span.setAttribute(SemanticConventions.LLM_REQUEST_TEMPERATURE, options.getTemperature());
- }
- if (options.getTopP() != null) {
- span.setAttribute(SemanticConventions.LLM_REQUEST_TOP_P, options.getTopP());
- }
- }
-
- // Execute call
ChatResponse response = delegate.call(prompt);
- // Extract model from response metadata (fallback when not in prompt options)
- if (response != null && response.getMetadata() != null) {
- String responseModel = response.getMetadata().getModel();
- if (responseModel != null && !responseModel.isEmpty()) {
- span.setAttribute(SemanticConventions.LLM_RESPONSE_MODEL, responseModel);
- // Also set request model if it wasn't set from options
- if (prompt == null || prompt.getOptions() == null || prompt.getOptions().getModel() == null) {
- span.setAttribute(SemanticConventions.LLM_REQUEST_MODEL, responseModel);
- }
- }
- }
-
- // Capture output
- if (response != null && response.getResult() != null) {
- var output = response.getResult().getOutput();
- if (output != null) {
- tracer.setOutputValue(span, output.getContent());
- tracer.setOutputMessages(span, Collections.singletonList(FITracer.message("assistant", output.getContent())));
- }
- }
-
- // Token usage
- if (response != null && response.getMetadata() != null && response.getMetadata().getUsage() != null) {
- var usage = response.getMetadata().getUsage();
- tracer.setTokenCounts(
- span,
- usage.getPromptTokens(),
- usage.getGenerationTokens(),
- usage.getTotalTokens()
- );
- }
-
+ captureResponse(span, prompt, response);
span.setStatus(StatusCode.OK);
+ completed = true;
return response;
-
- } catch (Exception e) {
+ } catch (RuntimeException e) {
+ errorRecorded = true;
tracer.setError(span, e);
throw e;
} finally {
+ if (!completed && !errorRecorded) {
+ // Fatal (non-RuntimeException) termination: never report success, let the Error propagate.
+ span.setStatus(StatusCode.ERROR, "abnormal termination");
+ }
span.end();
}
}
@Override
public Flux stream(Prompt prompt) {
- Span span = tracer.startSpan("Spring AI Chat (Stream)", FISpanKind.LLM);
+ SpringAITracingContext.Snapshot assembly = SpringAITracingContext.capture();
+ return new TracedStream(prompt, assembly);
+ }
- try {
- // Set system attributes
- span.setAttribute(SemanticConventions.LLM_SYSTEM, "spring-ai");
- span.setAttribute(SemanticConventions.LLM_PROVIDER, provider);
+ @Override
+ public ChatOptions getDefaultOptions() {
+ return delegate.getDefaultOptions();
+ }
- // Capture input messages and input.value
- if (prompt != null && prompt.getInstructions() != null) {
- List messages = prompt.getInstructions();
- List> inputMessages = new ArrayList<>();
- StringBuilder inputValueBuilder = new StringBuilder();
- for (Message msg : messages) {
- String role = msg.getMessageType().getValue();
- String content = msg.getContent();
- inputMessages.add(FITracer.message(role, content));
- if (content != null) {
- if (inputValueBuilder.length() > 0) {
- inputValueBuilder.append("\n");
- }
- inputValueBuilder.append(content);
+ /**
+ * Gets the underlying model.
+ *
+ * @return the wrapped ChatModel
+ */
+ public ChatModel unwrap() {
+ return delegate;
+ }
+
+ // ---------------------------------------------------------------------------------------------
+ // Capture helpers (null-safe for every Spring AI accessor the wrapper reads)
+ // ---------------------------------------------------------------------------------------------
+
+ private void captureRequest(Span span, Prompt prompt) {
+ span.setAttribute(SemanticConventions.LLM_SYSTEM, "spring-ai");
+ if (provider != null) {
+ span.setAttribute(SemanticConventions.LLM_PROVIDER, provider);
+ }
+ if (prompt == null) {
+ return;
+ }
+ List messages = prompt.getInstructions();
+ if (messages != null && !messages.isEmpty()) {
+ List> inputMessages = new ArrayList<>(messages.size());
+ StringBuilder inputValue = new StringBuilder();
+ for (Message message : messages) {
+ if (message == null) {
+ continue;
+ }
+ String role = message.getMessageType() != null ? message.getMessageType().getValue() : null;
+ String text = message.getText();
+ inputMessages.add(FITracer.message(role, text));
+ if (text != null) {
+ if (inputValue.length() > 0) {
+ inputValue.append('\n');
}
+ inputValue.append(text);
}
- tracer.setInputMessages(span, inputMessages);
- tracer.setInputValue(span, inputValueBuilder.toString());
}
+ tracer.setInputMessages(span, inputMessages);
+ tracer.setInputValue(span, inputValue.toString());
+ }
+ ChatOptions options = prompt.getOptions();
+ if (options != null) {
+ if (options.getModel() != null) {
+ span.setAttribute(SemanticConventions.LLM_REQUEST_MODEL, options.getModel());
+ }
+ if (options.getTemperature() != null) {
+ span.setAttribute(SemanticConventions.LLM_REQUEST_TEMPERATURE, options.getTemperature());
+ }
+ if (options.getTopP() != null) {
+ span.setAttribute(SemanticConventions.LLM_REQUEST_TOP_P, options.getTopP());
+ }
+ if (options.getMaxTokens() != null) {
+ span.setAttribute(SemanticConventions.LLM_REQUEST_MAX_TOKENS, options.getMaxTokens().longValue());
+ }
+ }
+ }
- // Capture prompt options if available
- if (prompt != null && prompt.getOptions() != null) {
- var options = prompt.getOptions();
- if (options.getModel() != null) {
- span.setAttribute(SemanticConventions.LLM_REQUEST_MODEL, options.getModel());
+ private void captureResponse(Span span, Prompt prompt, ChatResponse response) {
+ if (response == null) {
+ return;
+ }
+ ChatResponseMetadata metadata = response.getMetadata();
+ if (metadata != null) {
+ String responseModel = metadata.getModel();
+ if (responseModel != null && !responseModel.isEmpty()) {
+ span.setAttribute(SemanticConventions.LLM_RESPONSE_MODEL, responseModel);
+ boolean requestModelKnown = prompt != null && prompt.getOptions() != null && prompt.getOptions().getModel() != null;
+ if (!requestModelKnown) {
+ span.setAttribute(SemanticConventions.LLM_REQUEST_MODEL, responseModel);
}
}
+ }
- // Get the stream
- Flux stream = delegate.stream(prompt);
-
- // Wrap the stream to capture output
- StringBuilder content = new StringBuilder();
-
- return stream
- .doOnNext(response -> {
- if (response != null && response.getResult() != null) {
- var output = response.getResult().getOutput();
- if (output != null && output.getContent() != null) {
- content.append(output.getContent());
- }
- }
- })
- .doOnComplete(() -> {
- tracer.setOutputValue(span, content.toString());
- tracer.setOutputMessages(span, Collections.singletonList(FITracer.message("assistant", content.toString())));
- span.setStatus(StatusCode.OK);
- span.end();
- })
- .doOnError(e -> {
- tracer.setError(span, e);
- span.end();
- });
-
- } catch (Exception e) {
- tracer.setError(span, e);
- span.end();
- throw e;
+ Generation generation = response.getResult();
+ if (generation != null) {
+ AssistantMessage output = generation.getOutput();
+ if (output != null) {
+ String text = output.getText();
+ if (text != null) {
+ tracer.setOutputValue(span, text);
+ tracer.setOutputMessages(span, Collections.singletonList(FITracer.message("assistant", text)));
+ }
+ }
}
- }
- @Override
- public ChatOptions getDefaultOptions() {
- return delegate.getDefaultOptions();
+ if (metadata != null) {
+ applyUsage(span, metadata.getUsage());
+ }
}
/**
- * Gets the underlying model.
+ * Reads the three boxed counts once and sets each attribute only when the provider supplied it.
+ * Zero is preserved; missing counts stay absent (never synthesized).
*
- * @return the wrapped ChatModel
+ * Spring AI's {@code EmptyUsage} (the default when metadata carries no usage) returns zero for every
+ * count and an empty map from {@code getNativeUsage()}. That is "no usage reported", not a measured zero,
+ * so it is recorded as absent. A real measured zero keeps a provider-native usage object.
*/
- public ChatModel unwrap() {
- return delegate;
+ static void applyUsage(Span span, Usage usage) {
+ if (usage == null) {
+ return;
+ }
+ Integer prompt = usage.getPromptTokens();
+ Integer completion = usage.getCompletionTokens();
+ Integer total = usage.getTotalTokens();
+ Object nativeUsage = usage.getNativeUsage();
+ boolean emptyPlaceholder = nativeUsage instanceof java.util.Map && ((java.util.Map) nativeUsage).isEmpty()
+ && Integer.valueOf(0).equals(prompt) && Integer.valueOf(0).equals(completion) && Integer.valueOf(0).equals(total);
+ if (emptyPlaceholder) {
+ return;
+ }
+ if (prompt != null) {
+ span.setAttribute(SemanticConventions.LLM_TOKEN_COUNT_PROMPT, prompt.longValue());
+ }
+ if (completion != null) {
+ span.setAttribute(SemanticConventions.LLM_TOKEN_COUNT_COMPLETION, completion.longValue());
+ }
+ if (total != null) {
+ span.setAttribute(SemanticConventions.LLM_TOKEN_COUNT_TOTAL, total.longValue());
+ }
+ }
+
+ // ---------------------------------------------------------------------------------------------
+ // Deferred, per-subscription traced stream
+ // ---------------------------------------------------------------------------------------------
+
+ private final class TracedStream extends Flux {
+ private final Prompt prompt;
+ private final SpringAITracingContext.Snapshot assembly;
+
+ TracedStream(Prompt prompt, SpringAITracingContext.Snapshot assembly) {
+ this.prompt = prompt;
+ this.assembly = assembly;
+ }
+
+ @Override
+ public void subscribe(CoreSubscriber actual) {
+ SpringAITracingContext.Snapshot selected = SpringAITracingContext.select(actual.currentContext(), assembly);
+ Span span = SpringAITracingContext.startSpan(tracer, SPAN_NAME_STREAM, FISpanKind.LLM, selected);
+ TracingSubscriber subscriber = new TracingSubscriber(actual, tracer, span, new StreamOutputCapture(tracer));
+ try (Scope ignored = Context.current().with(span).makeCurrent()) {
+ captureRequest(span, prompt);
+ Flux upstream = delegate.stream(prompt);
+ if (upstream == null) {
+ throw new NullPointerException("delegate.stream(prompt) returned null");
+ }
+ upstream.subscribe(subscriber);
+ } catch (RuntimeException e) {
+ subscriber.failBeforeSubscription(e);
+ } catch (Error e) {
+ // Fatal creation/subscribe failure: close the span without claiming success, then propagate.
+ if (subscriber.state() == TracingSubscriber.OPEN) {
+ span.setStatus(StatusCode.ERROR, "abnormal termination");
+ span.end();
+ }
+ throw e;
+ }
+ }
}
}
diff --git a/java/traceai-spring-ai/src/main/java/ai/traceai/spring/TracedEmbeddingModel.java b/java/traceai-spring-ai/src/main/java/ai/traceai/spring/TracedEmbeddingModel.java
index 205016b1..3bc5ccfb 100644
--- a/java/traceai-spring-ai/src/main/java/ai/traceai/spring/TracedEmbeddingModel.java
+++ b/java/traceai-spring-ai/src/main/java/ai/traceai/spring/TracedEmbeddingModel.java
@@ -1,30 +1,46 @@
package ai.traceai.spring;
-import ai.traceai.*;
+import ai.traceai.FISpanKind;
+import ai.traceai.FITracer;
+import ai.traceai.SemanticConventions;
+import ai.traceai.TraceAI;
import io.opentelemetry.api.trace.Span;
import io.opentelemetry.api.trace.StatusCode;
import io.opentelemetry.context.Scope;
import org.springframework.ai.document.Document;
+import org.springframework.ai.embedding.BatchingStrategy;
+import org.springframework.ai.embedding.Embedding;
import org.springframework.ai.embedding.EmbeddingModel;
+import org.springframework.ai.embedding.EmbeddingOptions;
import org.springframework.ai.embedding.EmbeddingRequest;
import org.springframework.ai.embedding.EmbeddingResponse;
+import org.springframework.ai.embedding.EmbeddingResponseMetadata;
import java.util.List;
+import java.util.Objects;
+import java.util.function.Supplier;
/**
- * Traced wrapper for Spring AI EmbeddingModel.
- * Implements the EmbeddingModel interface to provide automatic tracing.
+ * Traced wrapper for a Spring AI 1.1.x {@link EmbeddingModel}.
+ *
+ * Every overload delegates directly to the same overload of the wrapped model with the original
+ * arguments, options and batching strategy, under exactly one span. Return values are the delegate's
+ * exact objects. Token usage is recorded only when the provider response carries it; bare vector
+ * overloads never get inferred usage.
*
* Usage:
*
* EmbeddingModel model = new OpenAiEmbeddingModel(api);
- * EmbeddingModel tracedModel = new TracedEmbeddingModel(model, "openai");
+ * EmbeddingModel tracedModel = new TracedEmbeddingModel(model, tracer, "openai");
*
- * EmbeddingResponse response = tracedModel.embed(List.of("Hello world"));
+ * List<float[]> vectors = tracedModel.embed(List.of("Hello world"));
*
*/
public class TracedEmbeddingModel implements EmbeddingModel {
+ static final String SPAN_NAME = "Spring AI Embedding";
+ static final int INPUT_PREVIEW_TEXTS = 5;
+
private final EmbeddingModel delegate;
private final FITracer tracer;
private final String provider;
@@ -37,8 +53,8 @@ public class TracedEmbeddingModel implements EmbeddingModel {
* @param provider the provider name (e.g., "openai", "cohere")
*/
public TracedEmbeddingModel(EmbeddingModel delegate, FITracer tracer, String provider) {
- this.delegate = delegate;
- this.tracer = tracer;
+ this.delegate = Objects.requireNonNull(delegate, "delegate EmbeddingModel must not be null");
+ this.tracer = Objects.requireNonNull(tracer, "FITracer must not be null");
this.provider = provider;
}
@@ -47,6 +63,7 @@ public TracedEmbeddingModel(EmbeddingModel delegate, FITracer tracer, String pro
*
* @param delegate the underlying model to wrap
* @param provider the provider name
+ * @throws IllegalStateException if {@link TraceAI#init} has not been called
*/
public TracedEmbeddingModel(EmbeddingModel delegate, String provider) {
this(delegate, TraceAI.getTracer(), provider);
@@ -54,106 +71,80 @@ public TracedEmbeddingModel(EmbeddingModel delegate, String provider) {
@Override
public EmbeddingResponse call(EmbeddingRequest request) {
- Span span = tracer.startSpan("Spring AI Embedding", FISpanKind.EMBEDDING);
-
- try (Scope scope = span.makeCurrent()) {
- // Set system attributes
- span.setAttribute(SemanticConventions.LLM_SYSTEM, "spring-ai");
- span.setAttribute(SemanticConventions.LLM_PROVIDER, provider);
-
- // Capture input
- if (request != null && request.getInstructions() != null) {
- List instructions = request.getInstructions();
- span.setAttribute(SemanticConventions.EMBEDDING_VECTOR_COUNT, (long) instructions.size());
-
- // Capture first few texts
- StringBuilder inputBuilder = new StringBuilder();
- int maxTexts = Math.min(instructions.size(), 5);
- for (int i = 0; i < maxTexts; i++) {
- if (i > 0) inputBuilder.append("\n---\n");
- inputBuilder.append(instructions.get(i));
- }
- if (instructions.size() > 5) {
- inputBuilder.append("\n... and ").append(instructions.size() - 5).append(" more");
+ return traced(span -> {
+ if (request != null) {
+ captureInputTexts(span, request.getInstructions());
+ EmbeddingOptions options = request.getOptions();
+ if (options != null && options.getModel() != null) {
+ span.setAttribute(SemanticConventions.EMBEDDING_MODEL_NAME, options.getModel());
+ span.setAttribute(SemanticConventions.LLM_REQUEST_MODEL, options.getModel());
}
- tracer.setInputValue(span, inputBuilder.toString());
- }
-
- // Capture model if available
- if (request != null && request.getOptions() != null && request.getOptions().getModel() != null) {
- span.setAttribute(SemanticConventions.EMBEDDING_MODEL_NAME, request.getOptions().getModel());
- span.setAttribute(SemanticConventions.LLM_REQUEST_MODEL, request.getOptions().getModel());
}
-
- // Execute embedding
- EmbeddingResponse response = delegate.call(request);
-
- // Extract model from response metadata as fallback
- if (response != null && response.getMetadata() != null && response.getMetadata().getModel() != null) {
- String responseModel = response.getMetadata().getModel();
- span.setAttribute(SemanticConventions.LLM_RESPONSE_MODEL, responseModel);
- if (request == null || request.getOptions() == null || request.getOptions().getModel() == null) {
- span.setAttribute(SemanticConventions.LLM_REQUEST_MODEL, responseModel);
- span.setAttribute(SemanticConventions.EMBEDDING_MODEL_NAME, responseModel);
- }
- }
-
- // Capture output
- if (response != null && response.getResults() != null && !response.getResults().isEmpty()) {
- span.setAttribute(SemanticConventions.EMBEDDING_VECTOR_COUNT,
- (long) response.getResults().size());
-
- // Get dimensions from first embedding
- var firstEmbedding = response.getResults().get(0);
- if (firstEmbedding != null && firstEmbedding.getOutput() != null) {
- span.setAttribute(SemanticConventions.EMBEDDING_DIMENSIONS,
- (long) firstEmbedding.getOutput().length);
- }
- }
-
- // Token usage
- if (response != null && response.getMetadata() != null && response.getMetadata().getUsage() != null) {
- var usage = response.getMetadata().getUsage();
- span.setAttribute(SemanticConventions.LLM_TOKEN_COUNT_PROMPT, usage.getPromptTokens());
- span.setAttribute(SemanticConventions.LLM_TOKEN_COUNT_TOTAL, usage.getTotalTokens());
- }
-
- span.setStatus(StatusCode.OK);
- return response;
-
- } catch (Exception e) {
- tracer.setError(span, e);
- throw e;
- } finally {
- span.end();
- }
+ }, () -> delegate.call(request), (span, response) -> {
+ boolean requestModelKnown = request != null && request.getOptions() != null && request.getOptions().getModel() != null;
+ captureResponse(span, response, requestModelKnown);
+ });
}
@Override
public float[] embed(Document document) {
- return delegate.embed(document);
+ return traced(span -> {
+ if (document != null) {
+ span.setAttribute(SemanticConventions.EMBEDDING_VECTOR_COUNT, 1L);
+ }
+ }, () -> delegate.embed(document), (span, vector) -> captureVector(span, vector));
}
@Override
public float[] embed(String text) {
- return delegate.embed(text);
+ return traced(span -> captureInputTexts(span, text == null ? null : List.of(text)),
+ () -> delegate.embed(text), (span, vector) -> captureVector(span, vector));
}
+ /**
+ * {@inheritDoc}
+ *
+ * @return one vector per input text, in input order (the delegate's exact list)
+ */
@Override
public List embed(List texts) {
- return delegate.embed(texts);
+ return traced(span -> captureInputTexts(span, texts), () -> delegate.embed(texts), (span, vectors) -> captureVectors(span, vectors));
+ }
+
+ @Override
+ public List embed(List documents, EmbeddingOptions options, BatchingStrategy batchingStrategy) {
+ return traced(span -> {
+ if (documents != null) {
+ span.setAttribute(SemanticConventions.EMBEDDING_VECTOR_COUNT, (long) documents.size());
+ }
+ if (options != null && options.getModel() != null) {
+ span.setAttribute(SemanticConventions.EMBEDDING_MODEL_NAME, options.getModel());
+ span.setAttribute(SemanticConventions.LLM_REQUEST_MODEL, options.getModel());
+ }
+ }, () -> delegate.embed(documents, options, batchingStrategy), (span, vectors) -> captureVectors(span, vectors));
}
@Override
public EmbeddingResponse embedForResponse(List texts) {
- return delegate.embedForResponse(texts);
+ return traced(span -> captureInputTexts(span, texts), () -> delegate.embedForResponse(texts),
+ (span, response) -> captureResponse(span, response, false));
}
+ /** Delegates directly; no probe embedding is issued by the wrapper. */
@Override
public int dimensions() {
return delegate.dimensions();
}
+ /**
+ * Delegates directly without a model-operation span. This accessor exists in Spring AI 1.1.8 and later
+ * only; an application on an older 1.1.x line simply never calls it.
+ */
+ @Override
+ public String getEmbeddingContent(Document document) {
+ return delegate.getEmbeddingContent(document);
+ }
+
/**
* Gets the underlying model.
*
@@ -162,4 +153,130 @@ public int dimensions() {
public EmbeddingModel unwrap() {
return delegate;
}
+
+ // ---------------------------------------------------------------------------------------------
+
+ private interface SpanConsumer {
+ void accept(Span span);
+ }
+
+ private interface SpanResultConsumer {
+ void accept(Span span, T result);
+ }
+
+ private T traced(SpanConsumer before, Supplier operation, SpanResultConsumer after) {
+ Span span = SpringAITracingContext.startSpan(tracer, SPAN_NAME, FISpanKind.EMBEDDING, SpringAITracingContext.capture());
+ boolean completed = false;
+ boolean errorRecorded = false;
+ try (Scope ignored = span.makeCurrent()) {
+ span.setAttribute(SemanticConventions.LLM_SYSTEM, "spring-ai");
+ if (provider != null) {
+ span.setAttribute(SemanticConventions.LLM_PROVIDER, provider);
+ }
+ before.accept(span);
+ T result = operation.get();
+ after.accept(span, result);
+ span.setStatus(StatusCode.OK);
+ completed = true;
+ return result;
+ } catch (RuntimeException e) {
+ errorRecorded = true;
+ tracer.setError(span, e);
+ throw e;
+ } finally {
+ if (!completed && !errorRecorded) {
+ span.setStatus(StatusCode.ERROR, "abnormal termination");
+ }
+ span.end();
+ }
+ }
+
+ private void captureInputTexts(Span span, List texts) {
+ if (texts == null) {
+ return;
+ }
+ span.setAttribute(SemanticConventions.EMBEDDING_VECTOR_COUNT, (long) texts.size());
+ StringBuilder preview = new StringBuilder();
+ int shown = Math.min(texts.size(), INPUT_PREVIEW_TEXTS);
+ for (int i = 0; i < shown; i++) {
+ if (i > 0) {
+ preview.append("\n---\n");
+ }
+ String text = texts.get(i);
+ if (text != null) {
+ preview.append(text);
+ }
+ }
+ if (texts.size() > INPUT_PREVIEW_TEXTS) {
+ preview.append("\n... and ").append(texts.size() - INPUT_PREVIEW_TEXTS).append(" more");
+ }
+ tracer.setInputValue(span, preview.toString());
+ }
+
+ private void captureResponse(Span span, EmbeddingResponse response, boolean requestModelKnown) {
+ if (response == null) {
+ return;
+ }
+ EmbeddingResponseMetadata metadata = response.getMetadata();
+ if (metadata != null) {
+ String responseModel = metadata.getModel();
+ if (responseModel != null && !responseModel.isEmpty()) {
+ span.setAttribute(SemanticConventions.LLM_RESPONSE_MODEL, responseModel);
+ if (!requestModelKnown) {
+ span.setAttribute(SemanticConventions.LLM_REQUEST_MODEL, responseModel);
+ span.setAttribute(SemanticConventions.EMBEDDING_MODEL_NAME, responseModel);
+ }
+ }
+ }
+ List results = response.getResults();
+ if (results != null && !results.isEmpty()) {
+ span.setAttribute(SemanticConventions.EMBEDDING_VECTOR_COUNT, (long) results.size());
+ Embedding first = results.get(0);
+ if (first != null && first.getOutput() != null) {
+ span.setAttribute(SemanticConventions.EMBEDDING_DIMENSIONS, (long) first.getOutput().length);
+ }
+ }
+ if (metadata != null) {
+ applyUsage(span, metadata.getUsage());
+ }
+ }
+
+ private static void captureVector(Span span, float[] vector) {
+ if (vector != null) {
+ span.setAttribute(SemanticConventions.EMBEDDING_VECTOR_COUNT, 1L);
+ span.setAttribute(SemanticConventions.EMBEDDING_DIMENSIONS, (long) vector.length);
+ }
+ }
+
+ private static void captureVectors(Span span, List vectors) {
+ if (vectors == null) {
+ return;
+ }
+ span.setAttribute(SemanticConventions.EMBEDDING_VECTOR_COUNT, (long) vectors.size());
+ if (!vectors.isEmpty() && vectors.get(0) != null) {
+ span.setAttribute(SemanticConventions.EMBEDDING_DIMENSIONS, (long) vectors.get(0).length);
+ }
+ }
+
+ /** Prompt and total counts are read once and set only when present; zero is preserved. */
+ private static void applyUsage(Span span, org.springframework.ai.chat.metadata.Usage usage) {
+ if (usage == null) {
+ return;
+ }
+ Integer prompt = usage.getPromptTokens();
+ Integer total = usage.getTotalTokens();
+ Object nativeUsage = usage.getNativeUsage();
+ // EmptyUsage (no usage reported) returns zeros plus an empty native map; that is absence, not zero.
+ if (nativeUsage instanceof java.util.Map && ((java.util.Map) nativeUsage).isEmpty()
+ && Integer.valueOf(0).equals(prompt) && Integer.valueOf(0).equals(usage.getCompletionTokens())
+ && Integer.valueOf(0).equals(total)) {
+ return;
+ }
+ if (prompt != null) {
+ span.setAttribute(SemanticConventions.LLM_TOKEN_COUNT_PROMPT, prompt.longValue());
+ }
+ if (total != null) {
+ span.setAttribute(SemanticConventions.LLM_TOKEN_COUNT_TOTAL, total.longValue());
+ }
+ }
}
diff --git a/java/traceai-spring-ai/src/main/java/ai/traceai/spring/TracingSubscriber.java b/java/traceai-spring-ai/src/main/java/ai/traceai/spring/TracingSubscriber.java
new file mode 100644
index 00000000..1b2e3f19
--- /dev/null
+++ b/java/traceai-spring-ai/src/main/java/ai/traceai/spring/TracingSubscriber.java
@@ -0,0 +1,183 @@
+package ai.traceai.spring;
+
+import ai.traceai.FITracer;
+import io.opentelemetry.api.trace.Span;
+import io.opentelemetry.api.trace.StatusCode;
+import io.opentelemetry.context.Context;
+import io.opentelemetry.context.Scope;
+import org.reactivestreams.Subscription;
+import org.springframework.ai.chat.model.ChatResponse;
+import reactor.core.CoreSubscriber;
+import reactor.core.publisher.Operators;
+
+import java.util.Objects;
+import java.util.concurrent.atomic.AtomicInteger;
+
+/**
+ * Package-private bounded adapter that owns exactly one span for one subscription of a traced stream.
+ *
+ * State machine: {@code OPEN -> COMPLETE | ERROR | CANCEL}, decided once by compare-and-set. The winning
+ * terminal path records the span outcome and ends the span in its {@code finally}; later signals are dropped.
+ * Demand is forwarded exactly; nothing is buffered, cached, collected or requested eagerly.
+ */
+final class TracingSubscriber implements CoreSubscriber, Subscription {
+
+ static final int OPEN = 0;
+ static final int COMPLETE = 1;
+ static final int ERROR = 2;
+ static final int CANCEL = 3;
+
+ static final String CANCELLED_DESCRIPTION = "cancelled";
+
+ private final CoreSubscriber actual;
+ private final FITracer tracer;
+ private final Span span;
+ private final StreamOutputCapture capture;
+ private final AtomicInteger state = new AtomicInteger(OPEN);
+ private volatile Subscription upstream;
+ private volatile boolean downstreamSubscribed;
+
+ TracingSubscriber(CoreSubscriber actual, FITracer tracer, Span span, StreamOutputCapture capture) {
+ this.actual = Objects.requireNonNull(actual, "actual");
+ this.tracer = Objects.requireNonNull(tracer, "tracer");
+ this.span = Objects.requireNonNull(span, "span");
+ this.capture = Objects.requireNonNull(capture, "capture");
+ }
+
+ boolean isDownstreamSubscribed() {
+ return downstreamSubscribed;
+ }
+
+ int state() {
+ return state.get();
+ }
+
+ @Override
+ public reactor.util.context.Context currentContext() {
+ return actual.currentContext();
+ }
+
+ @Override
+ public void onSubscribe(Subscription s) {
+ if (s == null) {
+ throw new NullPointerException("upstream Subscription must not be null");
+ }
+ if (this.upstream != null) {
+ s.cancel();
+ return;
+ }
+ this.upstream = s;
+ downstreamSubscribed = true;
+ actual.onSubscribe(this);
+ }
+
+ @Override
+ public void request(long n) {
+ Subscription s = upstream;
+ if (s == null || state.get() != OPEN) {
+ return;
+ }
+ try (Scope ignored = Context.current().with(span).makeCurrent()) {
+ s.request(n);
+ } catch (RuntimeException e) {
+ if (state.compareAndSet(OPEN, ERROR)) {
+ try {
+ tracer.setError(span, e);
+ } finally {
+ span.end();
+ }
+ try {
+ s.cancel();
+ } finally {
+ actual.onError(e);
+ }
+ } else {
+ Operators.onErrorDropped(e, actual.currentContext());
+ }
+ }
+ }
+
+ @Override
+ public void cancel() {
+ if (!state.compareAndSet(OPEN, CANCEL)) {
+ return;
+ }
+ try {
+ Subscription s = upstream;
+ if (s != null) {
+ try (Scope ignored = Context.current().with(span).makeCurrent()) {
+ s.cancel();
+ }
+ }
+ } finally {
+ try {
+ span.setStatus(StatusCode.ERROR, CANCELLED_DESCRIPTION);
+ } finally {
+ span.end();
+ }
+ }
+ }
+
+ @Override
+ public void onNext(ChatResponse response) {
+ if (state.get() != OPEN) {
+ Operators.onNextDropped(response, actual.currentContext());
+ return;
+ }
+ if (response != null) {
+ capture.accept(response);
+ }
+ actual.onNext(response);
+ }
+
+ @Override
+ public void onError(Throwable t) {
+ if (!state.compareAndSet(OPEN, ERROR)) {
+ Operators.onErrorDropped(t, actual.currentContext());
+ return;
+ }
+ try {
+ tracer.setError(span, t);
+ } catch (RuntimeException instrumentationFailure) {
+ // recording the error must not replace the original failure or skip ending the span
+ } finally {
+ span.end();
+ }
+ actual.onError(t);
+ }
+
+ @Override
+ public void onComplete() {
+ if (!state.compareAndSet(OPEN, COMPLETE)) {
+ return;
+ }
+ try {
+ capture.complete(span);
+ span.setStatus(StatusCode.OK);
+ } finally {
+ span.end();
+ }
+ actual.onComplete();
+ }
+
+ /**
+ * Terminates a subscription whose upstream never delivered {@code onSubscribe} (creation/subscribe failure).
+ * Records the error, ends the span once and signals the downstream subscriber.
+ */
+ void failBeforeSubscription(Throwable error) {
+ if (!state.compareAndSet(OPEN, ERROR)) {
+ Operators.onErrorDropped(error, actual.currentContext());
+ return;
+ }
+ try {
+ tracer.setError(span, error);
+ } finally {
+ span.end();
+ }
+ if (downstreamSubscribed) {
+ actual.onError(error);
+ } else {
+ Operators.error(actual, error);
+ }
+ }
+}
diff --git a/java/traceai-spring-ai/src/test/java/ai/traceai/spring/SpanTestSupport.java b/java/traceai-spring-ai/src/test/java/ai/traceai/spring/SpanTestSupport.java
new file mode 100644
index 00000000..f0ecf182
--- /dev/null
+++ b/java/traceai-spring-ai/src/test/java/ai/traceai/spring/SpanTestSupport.java
@@ -0,0 +1,177 @@
+package ai.traceai.spring;
+
+import ai.traceai.FITracer;
+import ai.traceai.TraceConfig;
+import io.opentelemetry.context.Context;
+import io.opentelemetry.sdk.testing.exporter.InMemorySpanExporter;
+import io.opentelemetry.sdk.trace.ReadWriteSpan;
+import io.opentelemetry.sdk.trace.ReadableSpan;
+import io.opentelemetry.sdk.trace.SdkTracerProvider;
+import io.opentelemetry.sdk.trace.SpanProcessor;
+import io.opentelemetry.sdk.trace.data.SpanData;
+import io.opentelemetry.sdk.trace.export.SimpleSpanProcessor;
+import org.springframework.ai.chat.messages.AssistantMessage;
+import org.springframework.ai.chat.metadata.ChatResponseMetadata;
+import org.springframework.ai.chat.metadata.DefaultUsage;
+import org.springframework.ai.chat.metadata.Usage;
+import org.springframework.ai.chat.model.ChatModel;
+import org.springframework.ai.chat.model.ChatResponse;
+import org.springframework.ai.chat.model.Generation;
+import org.springframework.ai.chat.prompt.ChatOptions;
+import org.springframework.ai.chat.prompt.Prompt;
+import reactor.core.publisher.Flux;
+
+import java.util.List;
+import java.util.concurrent.atomic.AtomicInteger;
+import java.util.function.Function;
+import java.util.function.Supplier;
+
+/**
+ * Deterministic in-memory tracing fixture. Uses an explicit SdkTracerProvider; never the global registration.
+ */
+final class SpanTestSupport implements AutoCloseable {
+
+ final InMemorySpanExporter exporter = InMemorySpanExporter.create();
+ final AtomicInteger spansStarted = new AtomicInteger();
+ final AtomicInteger spansEnded = new AtomicInteger();
+ final SdkTracerProvider provider;
+ final FITracer tracer;
+
+ SpanTestSupport() {
+ this(TraceConfig.builder().build());
+ }
+
+ SpanTestSupport(TraceConfig config) {
+ SpanProcessor counter = new SpanProcessor() {
+ @Override
+ public void onStart(Context parentContext, ReadWriteSpan span) {
+ spansStarted.incrementAndGet();
+ }
+
+ @Override
+ public boolean isStartRequired() {
+ return true;
+ }
+
+ @Override
+ public void onEnd(ReadableSpan span) {
+ spansEnded.incrementAndGet();
+ }
+
+ @Override
+ public boolean isEndRequired() {
+ return true;
+ }
+ };
+ this.provider = SdkTracerProvider.builder()
+ .addSpanProcessor(counter)
+ .addSpanProcessor(SimpleSpanProcessor.create(exporter))
+ .build();
+ this.tracer = new FITracer(provider.get("spring-ai-unit-test"), config);
+ }
+
+ List finished() {
+ return exporter.getFinishedSpanItems();
+ }
+
+ SpanData single() {
+ List spans = finished();
+ if (spans.size() != 1) {
+ throw new AssertionError("expected exactly one exported span but found " + spans.size() + ": " + spans);
+ }
+ return spans.get(0);
+ }
+
+ @Override
+ public void close() {
+ provider.close();
+ }
+
+ // ------------------------------------------------------------------ real Spring AI objects
+
+ static ChatResponse response(String text, Usage usage, String model) {
+ ChatResponseMetadata.Builder metadata = ChatResponseMetadata.builder();
+ if (usage != null) {
+ metadata.usage(usage);
+ }
+ if (model != null) {
+ metadata.model(model);
+ }
+ return new ChatResponse(List.of(new Generation(new AssistantMessage(text))), metadata.build());
+ }
+
+ static ChatResponse response(String text) {
+ return new ChatResponse(List.of(new Generation(new AssistantMessage(text))));
+ }
+
+ static Usage usage(int prompt, int completion, int total) {
+ return new DefaultUsage(prompt, completion, total);
+ }
+
+ /** Usage whose three counts are independently nullable (constructors cannot express a null total). */
+ static Usage usage(Integer prompt, Integer completion, Integer total, Object nativeUsage) {
+ return new Usage() {
+ @Override
+ public Integer getPromptTokens() {
+ return prompt;
+ }
+
+ @Override
+ public Integer getCompletionTokens() {
+ return completion;
+ }
+
+ @Override
+ public Integer getTotalTokens() {
+ return total;
+ }
+
+ @Override
+ public Object getNativeUsage() {
+ return nativeUsage;
+ }
+ };
+ }
+
+ static ChatModel syncModel(Function call) {
+ return new ChatModel() {
+ @Override
+ public ChatResponse call(Prompt prompt) {
+ return call.apply(prompt);
+ }
+
+ @Override
+ public Flux stream(Prompt prompt) {
+ throw new AssertionError("stream must not be used by this test");
+ }
+
+ @Override
+ public ChatOptions getDefaultOptions() {
+ return null;
+ }
+ };
+ }
+
+ static ChatModel streamModel(Function> stream) {
+ return new ChatModel() {
+ @Override
+ public ChatResponse call(Prompt prompt) {
+ throw new AssertionError("call must not be used by this test");
+ }
+
+ @Override
+ public Flux stream(Prompt prompt) {
+ return stream.apply(prompt);
+ }
+
+ @Override
+ public ChatOptions getDefaultOptions() {
+ return null;
+ }
+ };
+ }
+
+ static ChatModel streamModel(Supplier> stream) {
+ return streamModel(prompt -> stream.get());
+ }
+}
diff --git a/java/traceai-spring-ai/src/test/java/ai/traceai/spring/SpringAITracingContextTest.java b/java/traceai-spring-ai/src/test/java/ai/traceai/spring/SpringAITracingContextTest.java
new file mode 100644
index 00000000..adfc93fb
--- /dev/null
+++ b/java/traceai-spring-ai/src/test/java/ai/traceai/spring/SpringAITracingContextTest.java
@@ -0,0 +1,185 @@
+package ai.traceai.spring;
+
+import ai.traceai.ContextAttributes;
+import ai.traceai.SemanticConventions;
+import io.opentelemetry.api.common.AttributeKey;
+import io.opentelemetry.api.trace.Span;
+import io.opentelemetry.context.Context;
+import io.opentelemetry.sdk.trace.data.SpanData;
+import org.junit.jupiter.api.Test;
+import org.springframework.ai.chat.model.ChatResponse;
+import org.springframework.ai.chat.prompt.Prompt;
+import reactor.core.publisher.Flux;
+import reactor.core.scheduler.Schedulers;
+import reactor.test.StepVerifier;
+
+import java.time.Duration;
+import java.util.LinkedHashMap;
+import java.util.List;
+import java.util.Map;
+import java.util.concurrent.CountDownLatch;
+import java.util.concurrent.TimeUnit;
+import java.util.concurrent.atomic.AtomicReference;
+
+import static org.assertj.core.api.Assertions.assertThat;
+import static org.assertj.core.api.Assertions.assertThatThrownBy;
+
+class SpringAITracingContextTest {
+
+ private static final class QuietScope implements AutoCloseable {
+ private final AutoCloseable delegate;
+ QuietScope(AutoCloseable delegate) { this.delegate = delegate; }
+ @Override public void close() {
+ try { delegate.close(); } catch (RuntimeException e) { throw e; } catch (Exception e) { throw new IllegalStateException(e); }
+ }
+ }
+
+
+ private static final Duration TIMEOUT = Duration.ofSeconds(5);
+
+ private static SpanData streamSpan(SpanTestSupport support, Flux stream) {
+ StepVerifier.create(stream).expectNextCount(1).expectComplete().verify(TIMEOUT);
+ return support.single();
+ }
+
+ @Test
+ void explicitRootAndEmptyMapOverridePopulatedAssemblySnapshot() {
+ try (SpanTestSupport support = new SpanTestSupport();
+ var session = new QuietScope(ContextAttributes.usingSession("assembly-session"));
+ var user = new QuietScope(ContextAttributes.usingUser("assembly-user"))) {
+ Flux stream = new TracedChatModel(SpanTestSupport.streamModel(() -> Flux.just(SpanTestSupport.response("x"))),
+ support.tracer, "synthetic")
+ .stream(new Prompt("hi"))
+ .contextWrite(SpringAITracingContext.withContext(Context.root(), Map.of()));
+
+ SpanData span = streamSpan(support, stream);
+ assertThat(span.getAttributes().get(AttributeKey.stringKey(SemanticConventions.SESSION_ID))).isNull();
+ assertThat(span.getAttributes().get(AttributeKey.stringKey(SemanticConventions.USER_ID))).isNull();
+ assertThat(span.getParentSpanContext().isValid()).isFalse();
+ }
+ }
+
+ @Test
+ void serializedMetadataAndTagsAreAppliedVerbatim() {
+ try (SpanTestSupport support = new SpanTestSupport()) {
+ Map attributes = new LinkedHashMap<>();
+ attributes.put(SemanticConventions.METADATA, "{\"n\":7,\"flag\":true}");
+ attributes.put(SemanticConventions.TAG_TAGS, "[\"a\",\"b\"]");
+ attributes.put(SemanticConventions.SESSION_ID, "s-1");
+ attributes.put(SemanticConventions.GEN_AI_CONVERSATION_ID, "s-1");
+ attributes.put(SemanticConventions.USER_ID, "u-1");
+
+ Flux stream = new TracedChatModel(SpanTestSupport.streamModel(() -> Flux.just(SpanTestSupport.response("x"))),
+ support.tracer, "synthetic")
+ .stream(new Prompt("hi"))
+ .contextWrite(SpringAITracingContext.withContext(Context.root(), attributes));
+
+ SpanData span = streamSpan(support, stream);
+ assertThat(span.getAttributes().get(AttributeKey.stringKey(SemanticConventions.METADATA))).isEqualTo("{\"n\":7,\"flag\":true}");
+ assertThat(span.getAttributes().get(AttributeKey.stringKey(SemanticConventions.TAG_TAGS))).isEqualTo("[\"a\",\"b\"]");
+ assertThat(span.getAttributes().get(AttributeKey.stringKey(SemanticConventions.SESSION_ID))).isEqualTo("s-1");
+ assertThat(span.getAttributes().get(AttributeKey.stringKey(SemanticConventions.USER_ID))).isEqualTo("u-1");
+ }
+ }
+
+ @Test
+ void unknownKeysAndNullValuesAreRejectedBeforeSubscription() {
+ assertThatThrownBy(() -> SpringAITracingContext.withContext(Context.root(), Map.of("not.a.key", "v")))
+ .isInstanceOf(IllegalArgumentException.class);
+ Map nullValue = new LinkedHashMap<>();
+ nullValue.put(SemanticConventions.USER_ID, null);
+ assertThatThrownBy(() -> SpringAITracingContext.withContext(Context.root(), nullValue))
+ .isInstanceOf(IllegalArgumentException.class);
+ assertThatThrownBy(() -> SpringAITracingContext.withContext(null, Map.of()))
+ .isInstanceOf(NullPointerException.class);
+ }
+
+ @Test
+ void staleWorkerThreadAttributesDoNotLeakOntoTheSpanAndAreRestored() throws InterruptedException {
+ try (SpanTestSupport support = new SpanTestSupport()) {
+ AtomicReference seenDuringDelegate = new AtomicReference<>("unset");
+ CountDownLatch done = new CountDownLatch(1);
+ Schedulers.boundedElastic().schedule(() -> {
+ try (var ignored = new QuietScope(ContextAttributes.usingUser("stale-worker-user"))) {
+ Flux stream = new TracedChatModel(SpanTestSupport.streamModel(() -> {
+ seenDuringDelegate.set(String.valueOf(ContextAttributes.getAttributesFromContext().get(SemanticConventions.USER_ID)));
+ return Flux.just(SpanTestSupport.response("x"));
+ }), support.tracer, "synthetic")
+ .stream(new Prompt("hi"))
+ .contextWrite(SpringAITracingContext.withContext(Context.root(),
+ Map.of(SemanticConventions.USER_ID, "explicit-user")));
+ StepVerifier.create(stream).expectNextCount(1).expectComplete().verify(TIMEOUT);
+ assertThat(ContextAttributes.getAttributesFromContext().get(SemanticConventions.USER_ID))
+ .as("worker thread attributes restored after the span").isEqualTo("stale-worker-user");
+ } finally {
+ done.countDown();
+ }
+ });
+ assertThat(done.await(10, TimeUnit.SECONDS)).isTrue();
+
+ SpanData span = support.single();
+ assertThat(span.getAttributes().get(AttributeKey.stringKey(SemanticConventions.USER_ID))).isEqualTo("explicit-user");
+ assertThat(seenDuringDelegate.get()).isEqualTo("stale-worker-user");
+ }
+ }
+
+ @Test
+ void schedulerHopKeepsTheExplicitSnapshot() {
+ try (SpanTestSupport support = new SpanTestSupport()) {
+ Flux stream = new TracedChatModel(SpanTestSupport.streamModel(() ->
+ Flux.just(SpanTestSupport.response("x")).subscribeOn(Schedulers.boundedElastic())), support.tracer, "synthetic")
+ .stream(new Prompt("hi"))
+ .contextWrite(SpringAITracingContext.withContext(Context.root(),
+ Map.of(SemanticConventions.SESSION_ID, "hop-session")));
+
+ SpanData span = streamSpan(support, stream);
+ assertThat(span.getAttributes().get(AttributeKey.stringKey(SemanticConventions.SESSION_ID))).isEqualTo("hop-session");
+ }
+ }
+
+ @Test
+ void twoSyntheticTenantsStayIsolatedOnOnePublisher() {
+ try (SpanTestSupport support = new SpanTestSupport()) {
+ Flux shared = new TracedChatModel(SpanTestSupport.streamModel(() -> Flux.just(SpanTestSupport.response("x"))),
+ support.tracer, "synthetic").stream(new Prompt("hi"));
+
+ StepVerifier.create(shared.contextWrite(SpringAITracingContext.withContext(Context.root(),
+ Map.of(SemanticConventions.USER_ID, "tenant-a")))).expectNextCount(1).expectComplete().verify(TIMEOUT);
+ StepVerifier.create(shared.contextWrite(SpringAITracingContext.withContext(Context.root(),
+ Map.of(SemanticConventions.USER_ID, "tenant-b")))).expectNextCount(1).expectComplete().verify(TIMEOUT);
+
+ List spans = support.finished();
+ assertThat(spans).hasSize(2);
+ assertThat(spans.get(0).getAttributes().get(AttributeKey.stringKey(SemanticConventions.USER_ID))).isEqualTo("tenant-a");
+ assertThat(spans.get(1).getAttributes().get(AttributeKey.stringKey(SemanticConventions.USER_ID))).isEqualTo("tenant-b");
+ }
+ }
+
+ @Test
+ void subscriptionTimeAttributesWinWhenNoExplicitSnapshotExists() {
+ try (SpanTestSupport support = new SpanTestSupport();
+ var user = new QuietScope(ContextAttributes.usingUser("subscription-user"))) {
+ Flux stream = new TracedChatModel(SpanTestSupport.streamModel(() -> Flux.just(SpanTestSupport.response("x"))),
+ support.tracer, "synthetic").stream(new Prompt("hi"));
+ SpanData span = streamSpan(support, stream);
+ assertThat(span.getAttributes().get(AttributeKey.stringKey(SemanticConventions.USER_ID))).isEqualTo("subscription-user");
+ }
+ }
+
+ @Test
+ void explicitParentSpanIsTheSpanParent() {
+ try (SpanTestSupport support = new SpanTestSupport()) {
+ Span parent = support.tracer.startSpan("parent", ai.traceai.FISpanKind.AGENT);
+ try {
+ Flux stream = new TracedChatModel(SpanTestSupport.streamModel(() -> Flux.just(SpanTestSupport.response("x"))),
+ support.tracer, "synthetic")
+ .stream(new Prompt("hi"))
+ .contextWrite(SpringAITracingContext.withContext(Context.current().with(parent), Map.of()));
+ SpanData span = streamSpan(support, stream);
+ assertThat(span.getParentSpanContext().getSpanId()).isEqualTo(parent.getSpanContext().getSpanId());
+ } finally {
+ parent.end();
+ }
+ }
+ }
+}
diff --git a/java/traceai-spring-ai/src/test/java/ai/traceai/spring/TracedChatModelStreamTest.java b/java/traceai-spring-ai/src/test/java/ai/traceai/spring/TracedChatModelStreamTest.java
new file mode 100644
index 00000000..ab84ba07
--- /dev/null
+++ b/java/traceai-spring-ai/src/test/java/ai/traceai/spring/TracedChatModelStreamTest.java
@@ -0,0 +1,198 @@
+package ai.traceai.spring;
+
+import ai.traceai.SemanticConventions;
+import io.opentelemetry.api.common.AttributeKey;
+import io.opentelemetry.api.trace.StatusCode;
+import io.opentelemetry.sdk.trace.data.SpanData;
+import org.junit.jupiter.api.Test;
+import org.springframework.ai.chat.messages.AssistantMessage;
+import org.springframework.ai.chat.metadata.ChatResponseMetadata;
+import org.springframework.ai.chat.model.ChatModel;
+import org.springframework.ai.chat.model.ChatResponse;
+import org.springframework.ai.chat.model.Generation;
+import org.springframework.ai.chat.prompt.Prompt;
+import reactor.core.publisher.Flux;
+import reactor.test.StepVerifier;
+
+import java.time.Duration;
+import java.util.List;
+import java.util.concurrent.CountDownLatch;
+import java.util.concurrent.TimeUnit;
+import java.util.concurrent.atomic.AtomicInteger;
+import java.util.concurrent.atomic.AtomicLong;
+
+import static org.assertj.core.api.Assertions.assertThat;
+
+class TracedChatModelStreamTest {
+
+ private static final Duration TIMEOUT = Duration.ofSeconds(5);
+ private static final AttributeKey OUTPUT_VALUE = AttributeKey.stringKey(SemanticConventions.OUTPUT_VALUE);
+ private static final AttributeKey OUTPUT_MESSAGES = AttributeKey.stringKey(SemanticConventions.LLM_OUTPUT_MESSAGES);
+ private static final AttributeKey PROMPT_TOKENS = AttributeKey.longKey(SemanticConventions.LLM_TOKEN_COUNT_PROMPT);
+ private static final AttributeKey COMPLETION_TOKENS = AttributeKey.longKey(SemanticConventions.LLM_TOKEN_COUNT_COMPLETION);
+ private static final AttributeKey TOTAL_TOKENS = AttributeKey.longKey(SemanticConventions.LLM_TOKEN_COUNT_TOTAL);
+
+ @Test
+ void streamAssemblyDoesNotStartSpanOrInvokeDelegate() {
+ try (SpanTestSupport support = new SpanTestSupport()) {
+ AtomicInteger delegateCalls = new AtomicInteger();
+ ChatModel delegate = SpanTestSupport.streamModel(() -> {
+ delegateCalls.incrementAndGet();
+ return Flux.empty();
+ });
+ TracedChatModel model = new TracedChatModel(delegate, support.tracer, "synthetic");
+
+ Flux stream = model.stream(new Prompt("synthetic assembly input"));
+
+ assertThat(stream).isNotNull();
+ assertThat(support.spansStarted.get()).as("stream(prompt) must not start a span before subscription").isZero();
+ assertThat(delegateCalls.get()).as("stream(prompt) must not call the delegate before subscription").isZero();
+ assertThat(support.finished()).isEmpty();
+ }
+ }
+
+ @Test
+ void threeChunksAreDeliveredIdenticallyAndOutputIsConcatenated() {
+ try (SpanTestSupport support = new SpanTestSupport()) {
+ ChatResponse c1 = SpanTestSupport.response("Hel");
+ ChatResponse c2 = SpanTestSupport.response("lo ");
+ ChatResponse c3 = SpanTestSupport.response("world");
+ TracedChatModel model = new TracedChatModel(SpanTestSupport.streamModel(() -> Flux.just(c1, c2, c3)), support.tracer, "synthetic");
+
+ StepVerifier.create(model.stream(new Prompt("hi")))
+ .expectNextMatches(r -> r == c1)
+ .expectNextMatches(r -> r == c2)
+ .expectNextMatches(r -> r == c3)
+ .expectComplete()
+ .verify(TIMEOUT);
+
+ SpanData span = support.single();
+ assertThat(span.getName()).isEqualTo("Spring AI Chat (Stream)");
+ assertThat(span.getStatus().getStatusCode()).isEqualTo(StatusCode.OK);
+ assertThat(span.getAttributes().get(OUTPUT_VALUE)).isEqualTo("Hello world");
+ assertThat(span.getAttributes().get(OUTPUT_MESSAGES)).isEqualTo("[{\"role\":\"assistant\",\"content\":\"Hello world\"}]");
+ assertThat(span.getAttributes().get(AttributeKey.stringKey(SemanticConventions.INPUT_VALUE))).isEqualTo("hi");
+ assertThat(support.spansEnded.get()).isEqualTo(1);
+ }
+ }
+
+ @Test
+ void usageOnlyChunkIsForwardedWithoutTokenAggregation() {
+ try (SpanTestSupport support = new SpanTestSupport()) {
+ ChatResponse text = SpanTestSupport.response("done");
+ ChatResponse usageOnly = new ChatResponse(List.of(new Generation(new AssistantMessage(""))),
+ ChatResponseMetadata.builder().usage(SpanTestSupport.usage(3, 4, 7)).build());
+ ChatResponse metadataOnly = new ChatResponse(List.of(), ChatResponseMetadata.builder().model("m").build());
+ TracedChatModel model = new TracedChatModel(SpanTestSupport.streamModel(() -> Flux.just(text, usageOnly, metadataOnly)),
+ support.tracer, "synthetic");
+
+ StepVerifier.create(model.stream(new Prompt("hi")))
+ .expectNextMatches(r -> r == text)
+ .expectNextMatches(r -> r == usageOnly)
+ .expectNextMatches(r -> r == metadataOnly)
+ .expectComplete()
+ .verify(TIMEOUT);
+
+ SpanData span = support.single();
+ assertThat(span.getAttributes().get(OUTPUT_VALUE)).isEqualTo("done");
+ assertThat(span.getAttributes().get(PROMPT_TOKENS)).isNull();
+ assertThat(span.getAttributes().get(COMPLETION_TOKENS)).isNull();
+ assertThat(span.getAttributes().get(TOTAL_TOKENS)).isNull();
+ }
+ }
+
+ @Test
+ void emptyStreamCompletesWithOneOkSpan() {
+ try (SpanTestSupport support = new SpanTestSupport()) {
+ TracedChatModel model = new TracedChatModel(SpanTestSupport.streamModel(Flux::empty), support.tracer, "synthetic");
+
+ StepVerifier.create(model.stream(new Prompt("hi"))).expectComplete().verify(TIMEOUT);
+
+ SpanData span = support.single();
+ assertThat(span.getStatus().getStatusCode()).isEqualTo(StatusCode.OK);
+ assertThat(span.getAttributes().get(OUTPUT_VALUE)).isEqualTo("");
+ assertThat(support.spansEnded.get()).isEqualTo(1);
+ }
+ }
+
+ @Test
+ void repeatedSubscriptionsOwnSeparateSpans() {
+ try (SpanTestSupport support = new SpanTestSupport()) {
+ AtomicInteger delegateCalls = new AtomicInteger();
+ TracedChatModel model = new TracedChatModel(SpanTestSupport.streamModel(() -> {
+ delegateCalls.incrementAndGet();
+ return Flux.just(SpanTestSupport.response("a"));
+ }), support.tracer, "synthetic");
+ Flux stream = model.stream(new Prompt("reuse"));
+
+ StepVerifier.create(stream).expectNextCount(1).expectComplete().verify(TIMEOUT);
+ StepVerifier.create(stream).expectNextCount(1).expectComplete().verify(TIMEOUT);
+
+ assertThat(delegateCalls.get()).isEqualTo(2);
+ List spans = support.finished();
+ assertThat(spans).hasSize(2);
+ assertThat(spans.get(0).getSpanContext().getSpanId()).isNotEqualTo(spans.get(1).getSpanContext().getSpanId());
+ assertThat(support.spansEnded.get()).isEqualTo(2);
+ }
+ }
+
+ @Test
+ void concurrentSubscriptionsEachEndExactlyOnce() throws InterruptedException {
+ try (SpanTestSupport support = new SpanTestSupport()) {
+ CountDownLatch release = new CountDownLatch(1);
+ TracedChatModel model = new TracedChatModel(SpanTestSupport.streamModel(() ->
+ Flux.defer(() -> {
+ try {
+ release.await(5, TimeUnit.SECONDS);
+ } catch (InterruptedException e) {
+ Thread.currentThread().interrupt();
+ }
+ return Flux.just(SpanTestSupport.response("x"), SpanTestSupport.response("y"));
+ })), support.tracer, "synthetic");
+ Flux stream = model.stream(new Prompt("parallel"));
+
+ CountDownLatch done = new CountDownLatch(2);
+ Runnable subscriber = () -> stream.subscribe(r -> { }, e -> done.countDown(), done::countDown);
+ Thread t1 = new Thread(subscriber);
+ Thread t2 = new Thread(subscriber);
+ t1.start();
+ t2.start();
+ release.countDown();
+ assertThat(done.await(5, TimeUnit.SECONDS)).isTrue();
+ t1.join(5_000);
+ t2.join(5_000);
+
+ assertThat(support.finished()).hasSize(2);
+ assertThat(support.spansStarted.get()).isEqualTo(2);
+ assertThat(support.spansEnded.get()).isEqualTo(2);
+ assertThat(support.finished()).allSatisfy(s -> {
+ assertThat(s.getStatus().getStatusCode()).isEqualTo(StatusCode.OK);
+ assertThat(s.getAttributes().get(OUTPUT_VALUE)).isEqualTo("xy");
+ });
+ }
+ }
+
+ @Test
+ void demandIsForwardedExactlyWithoutEagerRequests() {
+ try (SpanTestSupport support = new SpanTestSupport()) {
+ AtomicLong requested = new AtomicLong();
+ Flux upstream = Flux.just(SpanTestSupport.response("1"), SpanTestSupport.response("2"), SpanTestSupport.response("3"))
+ .doOnRequest(requested::addAndGet);
+ TracedChatModel model = new TracedChatModel(SpanTestSupport.streamModel(() -> upstream), support.tracer, "synthetic");
+
+ StepVerifier.create(model.stream(new Prompt("demand")), 0)
+ .expectSubscription()
+ .then(() -> assertThat(requested.get()).isZero())
+ .thenRequest(1)
+ .expectNextCount(1)
+ .then(() -> assertThat(requested.get()).isEqualTo(1))
+ .thenRequest(2)
+ .expectNextCount(2)
+ .expectComplete()
+ .verify(TIMEOUT);
+
+ assertThat(requested.get()).isEqualTo(3);
+ assertThat(support.single().getStatus().getStatusCode()).isEqualTo(StatusCode.OK);
+ }
+ }
+}
diff --git a/java/traceai-spring-ai/src/test/java/ai/traceai/spring/TracedChatModelTest.java b/java/traceai-spring-ai/src/test/java/ai/traceai/spring/TracedChatModelTest.java
new file mode 100644
index 00000000..05d9cce8
--- /dev/null
+++ b/java/traceai-spring-ai/src/test/java/ai/traceai/spring/TracedChatModelTest.java
@@ -0,0 +1,221 @@
+package ai.traceai.spring;
+
+import ai.traceai.SemanticConventions;
+import io.opentelemetry.api.common.AttributeKey;
+import io.opentelemetry.api.trace.StatusCode;
+import io.opentelemetry.sdk.trace.data.SpanData;
+import org.junit.jupiter.api.Test;
+import org.springframework.ai.chat.messages.SystemMessage;
+import org.springframework.ai.chat.messages.UserMessage;
+import org.springframework.ai.chat.metadata.ChatResponseMetadata;
+import org.springframework.ai.chat.metadata.DefaultUsage;
+import org.springframework.ai.chat.metadata.EmptyUsage;
+import org.springframework.ai.chat.model.ChatModel;
+import org.springframework.ai.chat.model.ChatResponse;
+import org.springframework.ai.chat.model.Generation;
+import org.springframework.ai.chat.prompt.ChatOptions;
+import org.springframework.ai.chat.prompt.Prompt;
+
+import java.util.List;
+import java.util.concurrent.atomic.AtomicReference;
+
+import static org.assertj.core.api.Assertions.assertThat;
+import static org.assertj.core.api.Assertions.assertThatThrownBy;
+
+class TracedChatModelTest {
+
+ private static final AttributeKey PROMPT_TOKENS = AttributeKey.longKey(SemanticConventions.LLM_TOKEN_COUNT_PROMPT);
+ private static final AttributeKey COMPLETION_TOKENS = AttributeKey.longKey(SemanticConventions.LLM_TOKEN_COUNT_COMPLETION);
+ private static final AttributeKey TOTAL_TOKENS = AttributeKey.longKey(SemanticConventions.LLM_TOKEN_COUNT_TOTAL);
+ private static final AttributeKey INPUT_VALUE = AttributeKey.stringKey(SemanticConventions.INPUT_VALUE);
+ private static final AttributeKey OUTPUT_VALUE = AttributeKey.stringKey(SemanticConventions.OUTPUT_VALUE);
+ private static final AttributeKey INPUT_MESSAGES = AttributeKey.stringKey(SemanticConventions.LLM_INPUT_MESSAGES);
+ private static final AttributeKey OUTPUT_MESSAGES = AttributeKey.stringKey(SemanticConventions.LLM_OUTPUT_MESSAGES);
+
+ @Test
+ void callRecordsRealInputsOutputsAndUsageWithExactIdentity() {
+ try (SpanTestSupport support = new SpanTestSupport()) {
+ Prompt prompt = new Prompt(List.of(new SystemMessage("be terse"), new UserMessage("hello 1.1.x")),
+ ChatOptions.builder().model("synthetic-model").temperature(0.2).topP(0.9).maxTokens(64).build());
+ ChatResponse expected = SpanTestSupport.response("world", SpanTestSupport.usage(11, 7, 18), "synthetic-model-v2");
+ AtomicReference seen = new AtomicReference<>();
+ ChatModel delegate = SpanTestSupport.syncModel(p -> {
+ seen.set(p);
+ return expected;
+ });
+
+ ChatResponse actual = new TracedChatModel(delegate, support.tracer, "synthetic").call(prompt);
+
+ assertThat(actual).isSameAs(expected);
+ assertThat(seen.get()).isSameAs(prompt);
+ SpanData span = support.single();
+ assertThat(span.getName()).isEqualTo("Spring AI Chat");
+ assertThat(span.getStatus().getStatusCode()).isEqualTo(StatusCode.OK);
+ assertThat(span.getAttributes().get(AttributeKey.stringKey(SemanticConventions.FI_SPAN_KIND))).isEqualTo("LLM");
+ assertThat(span.getAttributes().get(AttributeKey.stringKey(SemanticConventions.LLM_PROVIDER))).isEqualTo("synthetic");
+ assertThat(span.getAttributes().get(INPUT_VALUE)).isEqualTo("be terse\nhello 1.1.x");
+ assertThat(span.getAttributes().get(INPUT_MESSAGES))
+ .isEqualTo("[{\"role\":\"system\",\"content\":\"be terse\"},{\"role\":\"user\",\"content\":\"hello 1.1.x\"}]");
+ assertThat(span.getAttributes().get(OUTPUT_VALUE)).isEqualTo("world");
+ assertThat(span.getAttributes().get(OUTPUT_MESSAGES)).isEqualTo("[{\"role\":\"assistant\",\"content\":\"world\"}]");
+ assertThat(span.getAttributes().get(PROMPT_TOKENS)).isEqualTo(11L);
+ assertThat(span.getAttributes().get(COMPLETION_TOKENS)).isEqualTo(7L);
+ assertThat(span.getAttributes().get(TOTAL_TOKENS)).isEqualTo(18L);
+ assertThat(span.getAttributes().get(AttributeKey.stringKey(SemanticConventions.LLM_REQUEST_MODEL))).isEqualTo("synthetic-model");
+ assertThat(span.getAttributes().get(AttributeKey.stringKey(SemanticConventions.LLM_RESPONSE_MODEL))).isEqualTo("synthetic-model-v2");
+ assertThat(span.getAttributes().get(AttributeKey.doubleKey(SemanticConventions.LLM_REQUEST_TEMPERATURE))).isEqualTo(0.2);
+ assertThat(span.getAttributes().get(AttributeKey.doubleKey(SemanticConventions.LLM_REQUEST_TOP_P))).isEqualTo(0.9);
+ assertThat(span.getAttributes().get(AttributeKey.longKey(SemanticConventions.LLM_REQUEST_MAX_TOKENS))).isEqualTo(64L);
+ assertThat(support.spansEnded.get()).isEqualTo(1);
+ }
+ }
+
+ @Test
+ void responseModelBackfillsRequestModelOnlyWhenOptionsDoNotName() {
+ try (SpanTestSupport support = new SpanTestSupport()) {
+ ChatModel delegate = SpanTestSupport.syncModel(p -> SpanTestSupport.response("x", null, "from-response"));
+ new TracedChatModel(delegate, support.tracer, "synthetic").call(new Prompt("no options"));
+
+ SpanData span = support.single();
+ assertThat(span.getAttributes().get(AttributeKey.stringKey(SemanticConventions.LLM_REQUEST_MODEL))).isEqualTo("from-response");
+ assertThat(span.getAttributes().get(AttributeKey.stringKey(SemanticConventions.LLM_RESPONSE_MODEL))).isEqualTo("from-response");
+ }
+ }
+
+ @Test
+ void nullResponseResultMetadataAndUsageAreTolerated() {
+ try (SpanTestSupport support = new SpanTestSupport()) {
+ TracedChatModel nullResponse = new TracedChatModel(SpanTestSupport.syncModel(p -> null), support.tracer, "synthetic");
+ assertThat(nullResponse.call(new Prompt("a"))).isNull();
+
+ ChatResponse emptyGenerations = new ChatResponse(List.of());
+ TracedChatModel noResult = new TracedChatModel(SpanTestSupport.syncModel(p -> emptyGenerations), support.tracer, "synthetic");
+ assertThat(noResult.call(new Prompt("b"))).isSameAs(emptyGenerations);
+
+ ChatResponse nullOutput = new ChatResponse(List.of(new Generation(null)),
+ ChatResponseMetadata.builder().usage(null).build());
+ TracedChatModel noOutput = new TracedChatModel(SpanTestSupport.syncModel(p -> nullOutput), support.tracer, "synthetic");
+ assertThat(noOutput.call(new Prompt("c"))).isSameAs(nullOutput);
+
+ List spans = support.finished();
+ assertThat(spans).hasSize(3);
+ for (SpanData span : spans) {
+ assertThat(span.getStatus().getStatusCode()).isEqualTo(StatusCode.OK);
+ assertThat(span.getAttributes().get(OUTPUT_VALUE)).isNull();
+ assertThat(span.getAttributes().get(PROMPT_TOKENS)).isNull();
+ assertThat(span.getAttributes().get(COMPLETION_TOKENS)).isNull();
+ assertThat(span.getAttributes().get(TOTAL_TOKENS)).isNull();
+ }
+ assertThat(support.spansEnded.get()).isEqualTo(3);
+ }
+ }
+
+ @Test
+ void tokenCountsAreIndependentlyNullableAndZeroIsPreserved() {
+ try (SpanTestSupport support = new SpanTestSupport()) {
+ ChatResponse zeroCompletion = SpanTestSupport.response("z", SpanTestSupport.usage(5, 0, 5), null);
+ ChatResponse onlyPrompt = SpanTestSupport.response("p", SpanTestSupport.usage(9, null, null, null), null);
+ ChatResponse onlyTotal = SpanTestSupport.response("t", SpanTestSupport.usage(null, null, 42, null), null);
+ TracedChatModel model = new TracedChatModel(SpanTestSupport.syncModel(p -> {
+ String text = p.getContents();
+ return text.equals("zero") ? zeroCompletion : text.equals("prompt") ? onlyPrompt : onlyTotal;
+ }), support.tracer, "synthetic");
+
+ model.call(new Prompt("zero"));
+ model.call(new Prompt("prompt"));
+ model.call(new Prompt("total"));
+
+ List spans = support.finished();
+ assertThat(spans.get(0).getAttributes().get(PROMPT_TOKENS)).isEqualTo(5L);
+ assertThat(spans.get(0).getAttributes().get(COMPLETION_TOKENS)).isEqualTo(0L);
+ assertThat(spans.get(0).getAttributes().get(TOTAL_TOKENS)).isEqualTo(5L);
+
+ assertThat(spans.get(1).getAttributes().get(PROMPT_TOKENS)).isEqualTo(9L);
+ assertThat(spans.get(1).getAttributes().get(COMPLETION_TOKENS)).isNull();
+ assertThat(spans.get(1).getAttributes().get(TOTAL_TOKENS)).isNull();
+
+ assertThat(spans.get(2).getAttributes().get(PROMPT_TOKENS)).isNull();
+ assertThat(spans.get(2).getAttributes().get(COMPLETION_TOKENS)).isNull();
+ assertThat(spans.get(2).getAttributes().get(TOTAL_TOKENS)).isEqualTo(42L);
+ }
+ }
+
+ @Test
+ void realEmptyUsageIsRecordedAsAbsentWhileMeasuredZeroIsPreserved() {
+ try (SpanTestSupport support = new SpanTestSupport()) {
+ ChatResponse placeholder = SpanTestSupport.response("e", new EmptyUsage(), null);
+ ChatResponse measuredZero = SpanTestSupport.response("z", new DefaultUsage(0, 0, 0), null);
+ TracedChatModel model = new TracedChatModel(SpanTestSupport.syncModel(p ->
+ p.getContents().equals("empty") ? placeholder : measuredZero), support.tracer, "synthetic");
+
+ model.call(new Prompt("empty"));
+ model.call(new Prompt("zero"));
+
+ java.util.List spans = support.finished();
+ assertThat(spans.get(0).getAttributes().get(PROMPT_TOKENS)).as("EmptyUsage is absence").isNull();
+ assertThat(spans.get(0).getAttributes().get(COMPLETION_TOKENS)).isNull();
+ assertThat(spans.get(0).getAttributes().get(TOTAL_TOKENS)).isNull();
+ assertThat(spans.get(1).getAttributes().get(PROMPT_TOKENS)).as("measured zero is kept").isEqualTo(0L);
+ assertThat(spans.get(1).getAttributes().get(COMPLETION_TOKENS)).isEqualTo(0L);
+ assertThat(spans.get(1).getAttributes().get(TOTAL_TOKENS)).isEqualTo(0L);
+ }
+ }
+
+ @Test
+ void runtimeExceptionKeepsIdentityAndEndsSpanWithError() {
+ try (SpanTestSupport support = new SpanTestSupport()) {
+ IllegalStateException failure = new IllegalStateException("synthetic provider failure");
+ TracedChatModel model = new TracedChatModel(SpanTestSupport.syncModel(p -> {
+ throw failure;
+ }), support.tracer, "synthetic");
+
+ assertThatThrownBy(() -> model.call(new Prompt("boom"))).isSameAs(failure);
+
+ SpanData span = support.single();
+ assertThat(span.getStatus().getStatusCode()).isEqualTo(StatusCode.ERROR);
+ assertThat(span.getAttributes().get(AttributeKey.stringKey(SemanticConventions.ERROR_TYPE)))
+ .isEqualTo(IllegalStateException.class.getName());
+ assertThat(span.getEvents()).anySatisfy(event -> assertThat(event.getName()).isEqualTo("exception"));
+ assertThat(support.spansEnded.get()).isEqualTo(1);
+ }
+ }
+
+ @Test
+ void syntheticLinkageErrorEscapesUnchangedAfterSpanCleanup() {
+ try (SpanTestSupport support = new SpanTestSupport()) {
+ NoSuchMethodError failure = new NoSuchMethodError("synthetic Message.getContent()");
+ TracedChatModel model = new TracedChatModel(SpanTestSupport.syncModel(p -> {
+ throw failure;
+ }), support.tracer, "synthetic");
+
+ assertThatThrownBy(() -> model.call(new Prompt("link"))).isSameAs(failure);
+
+ SpanData span = support.single();
+ assertThat(span.getStatus().getStatusCode()).isNotEqualTo(StatusCode.OK);
+ assertThat(span.hasEnded()).isTrue();
+ assertThat(support.spansEnded.get()).isEqualTo(1);
+ }
+ }
+
+ @Test
+ void defaultOptionsAndUnwrapDelegateExactly() {
+ try (SpanTestSupport support = new SpanTestSupport()) {
+ ChatOptions defaults = ChatOptions.builder().model("d").build();
+ ChatModel delegate = new ChatModel() {
+ @Override
+ public ChatResponse call(Prompt prompt) {
+ return null;
+ }
+
+ @Override
+ public ChatOptions getDefaultOptions() {
+ return defaults;
+ }
+ };
+ TracedChatModel model = new TracedChatModel(delegate, support.tracer, "synthetic");
+ assertThat(model.getDefaultOptions()).isSameAs(defaults);
+ assertThat(model.unwrap()).isSameAs(delegate);
+ assertThat(support.spansStarted.get()).isZero();
+ }
+ }
+}
diff --git a/java/traceai-spring-ai/src/test/java/ai/traceai/spring/TracedEmbeddingModelTest.java b/java/traceai-spring-ai/src/test/java/ai/traceai/spring/TracedEmbeddingModelTest.java
new file mode 100644
index 00000000..9f9ea4d0
--- /dev/null
+++ b/java/traceai-spring-ai/src/test/java/ai/traceai/spring/TracedEmbeddingModelTest.java
@@ -0,0 +1,164 @@
+package ai.traceai.spring;
+
+import ai.traceai.SemanticConventions;
+import io.opentelemetry.api.common.AttributeKey;
+import io.opentelemetry.api.trace.StatusCode;
+import io.opentelemetry.sdk.trace.data.SpanData;
+import org.junit.jupiter.api.Test;
+import org.springframework.ai.chat.metadata.Usage;
+import org.springframework.ai.document.Document;
+import org.springframework.ai.embedding.BatchingStrategy;
+import org.springframework.ai.embedding.Embedding;
+import org.springframework.ai.embedding.EmbeddingModel;
+import org.springframework.ai.embedding.EmbeddingOptions;
+import org.springframework.ai.embedding.EmbeddingRequest;
+import org.springframework.ai.embedding.EmbeddingResponse;
+import org.springframework.ai.embedding.EmbeddingResponseMetadata;
+
+import java.util.List;
+import java.util.concurrent.atomic.AtomicReference;
+
+import static org.assertj.core.api.Assertions.assertThat;
+import static org.assertj.core.api.Assertions.assertThatThrownBy;
+
+class TracedEmbeddingModelTest {
+
+ private static final AttributeKey VECTOR_COUNT = AttributeKey.longKey(SemanticConventions.EMBEDDING_VECTOR_COUNT);
+ private static final AttributeKey DIMENSIONS = AttributeKey.longKey(SemanticConventions.EMBEDDING_DIMENSIONS);
+ private static final AttributeKey PROMPT_TOKENS = AttributeKey.longKey(SemanticConventions.LLM_TOKEN_COUNT_PROMPT);
+ private static final AttributeKey TOTAL_TOKENS = AttributeKey.longKey(SemanticConventions.LLM_TOKEN_COUNT_TOTAL);
+ private static final AttributeKey INPUT_VALUE = AttributeKey.stringKey(SemanticConventions.INPUT_VALUE);
+
+ /** Records which overload was invoked and returns fixed values. */
+ static class RecordingModel implements EmbeddingModel {
+ final AtomicReference invoked = new AtomicReference<>();
+ EmbeddingRequest seenRequest;
+ EmbeddingOptions seenOptions;
+ BatchingStrategy seenStrategy;
+ final EmbeddingResponse response;
+
+ RecordingModel() {
+ EmbeddingResponseMetadata metadata = new EmbeddingResponseMetadata("embed-model",
+ new Usage() {
+ @Override public Integer getPromptTokens() { return 6; }
+ @Override public Integer getCompletionTokens() { return null; }
+ @Override public Integer getTotalTokens() { return 6; }
+ @Override public Object getNativeUsage() { return "native"; }
+ });
+ this.response = new EmbeddingResponse(List.of(
+ new Embedding(new float[]{1f, 2f, 3f}, 0),
+ new Embedding(new float[]{4f, 5f, 6f}, 1)), metadata);
+ }
+
+ @Override public EmbeddingResponse call(EmbeddingRequest request) { invoked.set("call"); seenRequest = request; return response; }
+ @Override public float[] embed(Document document) { invoked.set("embed-document"); return new float[]{1f, 2f, 3f}; }
+ @Override public float[] embed(String text) { invoked.set("embed-string"); return new float[]{1f, 2f, 3f}; }
+ @Override public List embed(List texts) { invoked.set("embed-list"); return List.of(new float[]{1f, 2f, 3f}, new float[]{4f, 5f, 6f}); }
+ @Override public List embed(List documents, EmbeddingOptions options, BatchingStrategy batchingStrategy) {
+ invoked.set("embed-batch"); seenOptions = options; seenStrategy = batchingStrategy;
+ return List.of(new float[]{1f, 2f, 3f});
+ }
+ @Override public EmbeddingResponse embedForResponse(List texts) { invoked.set("embedForResponse"); return response; }
+ @Override public int dimensions() { invoked.set("dimensions"); return 3; }
+ @Override public String getEmbeddingContent(Document document) { invoked.set("getEmbeddingContent"); return "doc-text"; }
+ }
+
+ private static TracedEmbeddingModel traced(SpanTestSupport support, RecordingModel delegate) {
+ return new TracedEmbeddingModel(delegate, support.tracer, "synthetic");
+ }
+
+ @Test
+ void callDelegatesExactlyAndRecordsVectorsAndUsage() {
+ try (SpanTestSupport support = new SpanTestSupport()) {
+ RecordingModel delegate = new RecordingModel();
+ EmbeddingRequest request = new EmbeddingRequest(List.of("alpha", "beta"), EmbeddingOptions.builder().model("m").build());
+ EmbeddingResponse actual = traced(support, delegate).call(request);
+
+ assertThat(actual).isSameAs(delegate.response);
+ assertThat(delegate.invoked.get()).isEqualTo("call");
+ assertThat(delegate.seenRequest).isSameAs(request);
+ SpanData span = support.single();
+ assertThat(span.getName()).isEqualTo("Spring AI Embedding");
+ assertThat(span.getAttributes().get(VECTOR_COUNT)).isEqualTo(2L);
+ assertThat(span.getAttributes().get(DIMENSIONS)).isEqualTo(3L);
+ assertThat(span.getAttributes().get(PROMPT_TOKENS)).isEqualTo(6L);
+ assertThat(span.getAttributes().get(TOTAL_TOKENS)).isEqualTo(6L);
+ assertThat(span.getAttributes().get(AttributeKey.stringKey(SemanticConventions.LLM_REQUEST_MODEL))).isEqualTo("m");
+ assertThat(span.getAttributes().get(AttributeKey.stringKey(SemanticConventions.LLM_RESPONSE_MODEL))).isEqualTo("embed-model");
+ assertThat(span.getAttributes().get(INPUT_VALUE)).isEqualTo("alpha\n---\nbeta");
+ }
+ }
+
+ @Test
+ void everyOverloadDelegatesDirectlyWithOriginalArguments() {
+ try (SpanTestSupport support = new SpanTestSupport()) {
+ RecordingModel delegate = new RecordingModel();
+ TracedEmbeddingModel model = traced(support, delegate);
+ Document document = new Document("doc-text");
+ EmbeddingOptions options = EmbeddingOptions.builder().model("batch-model").build();
+ BatchingStrategy strategy = documents -> List.of(documents);
+
+ assertThat(model.embed(document)).containsExactly(1f, 2f, 3f);
+ assertThat(delegate.invoked.get()).isEqualTo("embed-document");
+ assertThat(model.embed("hello")).containsExactly(1f, 2f, 3f);
+ assertThat(delegate.invoked.get()).isEqualTo("embed-string");
+ List list = model.embed(List.of("a", "b"));
+ assertThat(delegate.invoked.get()).isEqualTo("embed-list");
+ assertThat(list).hasSize(2);
+ List batch = model.embed(List.of(document), options, strategy);
+ assertThat(delegate.invoked.get()).isEqualTo("embed-batch");
+ assertThat(delegate.seenOptions).isSameAs(options);
+ assertThat(delegate.seenStrategy).isSameAs(strategy);
+ assertThat(batch).hasSize(1);
+ assertThat(model.embedForResponse(List.of("a"))).isSameAs(delegate.response);
+ assertThat(delegate.invoked.get()).isEqualTo("embedForResponse");
+ assertThat(model.dimensions()).isEqualTo(3);
+ assertThat(delegate.invoked.get()).isEqualTo("dimensions");
+ assertThat(model.getEmbeddingContent(document)).isEqualTo("doc-text");
+ assertThat(delegate.invoked.get()).isEqualTo("getEmbeddingContent");
+ assertThat(model.unwrap()).isSameAs(delegate);
+
+ // dimensions() and getEmbeddingContent() open no model-operation span
+ assertThat(support.finished()).hasSize(5);
+ assertThat(support.finished()).allSatisfy(span -> assertThat(span.getStatus().getStatusCode()).isEqualTo(StatusCode.OK));
+ }
+ }
+
+ @Test
+ void inputPreviewIsBoundedToFiveTexts() {
+ try (SpanTestSupport support = new SpanTestSupport()) {
+ RecordingModel delegate = new RecordingModel();
+ traced(support, delegate).embed(List.of("1", "2", "3", "4", "5", "6", "7"));
+ assertThat(support.single().getAttributes().get(INPUT_VALUE)).isEqualTo("1\n---\n2\n---\n3\n---\n4\n---\n5\n... and 2 more");
+ assertThat(support.single().getAttributes().get(VECTOR_COUNT)).isEqualTo(2L);
+ }
+ }
+
+ @Test
+ void nullUsageAndNullResponseAreTolerated() {
+ try (SpanTestSupport support = new SpanTestSupport()) {
+ EmbeddingModel delegate = new RecordingModel() {
+ @Override public EmbeddingResponse call(EmbeddingRequest request) { invoked.set("call"); return null; }
+ };
+ assertThat(new TracedEmbeddingModel(delegate, support.tracer, "synthetic").call(new EmbeddingRequest(List.of("a"), null))).isNull();
+ SpanData span = support.single();
+ assertThat(span.getStatus().getStatusCode()).isEqualTo(StatusCode.OK);
+ assertThat(span.getAttributes().get(PROMPT_TOKENS)).isNull();
+ assertThat(span.getAttributes().get(TOTAL_TOKENS)).isNull();
+ }
+ }
+
+ @Test
+ void runtimeExceptionKeepsIdentity() {
+ try (SpanTestSupport support = new SpanTestSupport()) {
+ IllegalStateException failure = new IllegalStateException("embedding failed");
+ EmbeddingModel delegate = new RecordingModel() {
+ @Override public EmbeddingResponse call(EmbeddingRequest request) { throw failure; }
+ };
+ assertThatThrownBy(() -> new TracedEmbeddingModel(delegate, support.tracer, "synthetic")
+ .call(new EmbeddingRequest(List.of("a"), null))).isSameAs(failure);
+ assertThat(support.single().getStatus().getStatusCode()).isEqualTo(StatusCode.ERROR);
+ assertThat(support.spansEnded.get()).isEqualTo(1);
+ }
+ }
+}
diff --git a/java/traceai-spring-ai/src/test/java/ai/traceai/spring/TracingLifecycleTest.java b/java/traceai-spring-ai/src/test/java/ai/traceai/spring/TracingLifecycleTest.java
new file mode 100644
index 00000000..a4bf1e09
--- /dev/null
+++ b/java/traceai-spring-ai/src/test/java/ai/traceai/spring/TracingLifecycleTest.java
@@ -0,0 +1,250 @@
+package ai.traceai.spring;
+
+import ai.traceai.SemanticConventions;
+import io.opentelemetry.api.common.AttributeKey;
+import io.opentelemetry.api.trace.StatusCode;
+import io.opentelemetry.sdk.trace.data.SpanData;
+import org.junit.jupiter.api.Test;
+import org.reactivestreams.Subscription;
+import org.springframework.ai.chat.model.ChatResponse;
+import org.springframework.ai.chat.prompt.Prompt;
+import reactor.core.CoreSubscriber;
+import reactor.core.publisher.BaseSubscriber;
+import reactor.core.publisher.Flux;
+import reactor.core.publisher.Sinks;
+import reactor.test.StepVerifier;
+
+import java.time.Duration;
+import java.util.List;
+import java.util.concurrent.CountDownLatch;
+import java.util.concurrent.TimeUnit;
+import java.util.concurrent.atomic.AtomicInteger;
+import java.util.concurrent.atomic.AtomicReference;
+
+import static org.assertj.core.api.Assertions.assertThat;
+import static org.assertj.core.api.Assertions.assertThatThrownBy;
+
+/** Adversarial lifecycle suite: exactly one span end and one export per subscription, whatever terminates it. */
+class TracingLifecycleTest {
+
+ private static final Duration TIMEOUT = Duration.ofSeconds(5);
+ private static final AttributeKey OUTPUT_VALUE = AttributeKey.stringKey(SemanticConventions.OUTPUT_VALUE);
+
+ @Test
+ void cancelAfterFirstChunkEndsSpanOnceAsCancelled() {
+ try (SpanTestSupport support = new SpanTestSupport()) {
+ Sinks.Many sink = Sinks.many().unicast().onBackpressureBuffer();
+ AtomicInteger upstreamCancels = new AtomicInteger();
+ TracedChatModel model = new TracedChatModel(SpanTestSupport.streamModel(() ->
+ sink.asFlux().doOnCancel(upstreamCancels::incrementAndGet)), support.tracer, "synthetic");
+
+ StepVerifier.create(model.stream(new Prompt("cancel")))
+ .then(() -> sink.tryEmitNext(SpanTestSupport.response("partial")))
+ .expectNextCount(1)
+ .thenCancel()
+ .verify(TIMEOUT);
+
+ SpanData span = support.single();
+ assertThat(span.getStatus().getStatusCode()).isEqualTo(StatusCode.ERROR);
+ assertThat(span.getStatus().getDescription()).isEqualTo("cancelled");
+ assertThat(span.getAttributes().get(OUTPUT_VALUE)).as("partial output is not recorded as success").isNull();
+ assertThat(upstreamCancels.get()).isEqualTo(1);
+ assertThat(support.spansEnded.get()).isEqualTo(1);
+
+ // late signals after cancel must not reopen or re-end the span
+ sink.tryEmitNext(SpanTestSupport.response("late"));
+ sink.tryEmitComplete();
+ assertThat(support.spansEnded.get()).isEqualTo(1);
+ assertThat(support.finished()).hasSize(1);
+ }
+ }
+
+ @Test
+ void immediateCancelFromOnSubscribePreventsDeliveryAndClosesSpan() {
+ try (SpanTestSupport support = new SpanTestSupport()) {
+ AtomicInteger delivered = new AtomicInteger();
+ TracedChatModel model = new TracedChatModel(SpanTestSupport.streamModel(() ->
+ Flux.just(SpanTestSupport.response("never"))), support.tracer, "synthetic");
+
+ model.stream(new Prompt("immediate")).subscribe(new BaseSubscriber<>() {
+ @Override
+ protected void hookOnSubscribe(Subscription subscription) {
+ subscription.cancel();
+ }
+
+ @Override
+ protected void hookOnNext(ChatResponse value) {
+ delivered.incrementAndGet();
+ }
+ });
+
+ assertThat(delivered.get()).isZero();
+ SpanData span = support.single();
+ assertThat(span.getStatus().getStatusCode()).isEqualTo(StatusCode.ERROR);
+ assertThat(span.getStatus().getDescription()).isEqualTo("cancelled");
+ assertThat(support.spansEnded.get()).isEqualTo(1);
+ }
+ }
+
+ @Test
+ void delegateStreamThrowingBeforeReturnSignalsErrorAndEndsSpanOnce() {
+ try (SpanTestSupport support = new SpanTestSupport()) {
+ IllegalStateException failure = new IllegalStateException("synthetic pre-return failure");
+ TracedChatModel model = new TracedChatModel(SpanTestSupport.streamModel(() -> {
+ throw failure;
+ }), support.tracer, "synthetic");
+
+ StepVerifier.create(model.stream(new Prompt("pre-return")))
+ .expectErrorSatisfies(e -> assertThat(e).isSameAs(failure))
+ .verify(TIMEOUT);
+
+ SpanData span = support.single();
+ assertThat(span.getStatus().getStatusCode()).isEqualTo(StatusCode.ERROR);
+ assertThat(span.getAttributes().get(AttributeKey.stringKey(SemanticConventions.ERROR_TYPE)))
+ .isEqualTo(IllegalStateException.class.getName());
+ assertThat(support.spansEnded.get()).isEqualTo(1);
+ }
+ }
+
+ @Test
+ void subscribeThrowingAfterReturnSignalsErrorAndEndsSpanOnce() {
+ try (SpanTestSupport support = new SpanTestSupport()) {
+ IllegalArgumentException failure = new IllegalArgumentException("synthetic subscribe failure");
+ Flux hostile = new Flux<>() {
+ @Override
+ public void subscribe(CoreSubscriber actual) {
+ throw failure;
+ }
+ };
+ TracedChatModel model = new TracedChatModel(SpanTestSupport.streamModel(() -> hostile), support.tracer, "synthetic");
+
+ StepVerifier.create(model.stream(new Prompt("subscribe")))
+ .expectErrorSatisfies(e -> assertThat(e).isSameAs(failure))
+ .verify(TIMEOUT);
+
+ assertThat(support.single().getStatus().getStatusCode()).isEqualTo(StatusCode.ERROR);
+ assertThat(support.spansEnded.get()).isEqualTo(1);
+ }
+ }
+
+ @Test
+ void requestThrowingIsRecordedOnceAndForwardedAsError() {
+ try (SpanTestSupport support = new SpanTestSupport()) {
+ IllegalStateException failure = new IllegalStateException("synthetic request failure");
+ Flux hostile = new Flux<>() {
+ @Override
+ public void subscribe(CoreSubscriber actual) {
+ actual.onSubscribe(new Subscription() {
+ @Override
+ public void request(long n) {
+ throw failure;
+ }
+
+ @Override
+ public void cancel() {
+ }
+ });
+ }
+ };
+ TracedChatModel model = new TracedChatModel(SpanTestSupport.streamModel(() -> hostile), support.tracer, "synthetic");
+
+ StepVerifier.create(model.stream(new Prompt("request")))
+ .expectErrorSatisfies(e -> assertThat(e).isSameAs(failure))
+ .verify(TIMEOUT);
+
+ assertThat(support.single().getStatus().getStatusCode()).isEqualTo(StatusCode.ERROR);
+ assertThat(support.spansEnded.get()).isEqualTo(1);
+ }
+ }
+
+ @Test
+ void upstreamErrorKeepsIdentityAndDoesNotRecordPartialOutputAsSuccess() {
+ try (SpanTestSupport support = new SpanTestSupport()) {
+ RuntimeException failure = new RuntimeException("synthetic mid-stream failure");
+ TracedChatModel model = new TracedChatModel(SpanTestSupport.streamModel(() ->
+ Flux.concat(Flux.just(SpanTestSupport.response("partial")), Flux.error(failure))), support.tracer, "synthetic");
+
+ StepVerifier.create(model.stream(new Prompt("error")))
+ .expectNextCount(1)
+ .expectErrorSatisfies(e -> assertThat(e).isSameAs(failure))
+ .verify(TIMEOUT);
+
+ SpanData span = support.single();
+ assertThat(span.getStatus().getStatusCode()).isEqualTo(StatusCode.ERROR);
+ assertThat(span.getAttributes().get(OUTPUT_VALUE)).isNull();
+ assertThat(span.getEvents()).anySatisfy(event -> assertThat(event.getName()).isEqualTo("exception"));
+ assertThat(support.spansEnded.get()).isEqualTo(1);
+ }
+ }
+
+ @Test
+ void syntheticLinkageErrorFromDelegateStreamPropagatesAfterCleanup() {
+ try (SpanTestSupport support = new SpanTestSupport()) {
+ NoSuchMethodError failure = new NoSuchMethodError("synthetic linkage");
+ TracedChatModel model = new TracedChatModel(SpanTestSupport.streamModel(() -> {
+ throw failure;
+ }), support.tracer, "synthetic");
+
+ Flux stream = model.stream(new Prompt("linkage"));
+ assertThatThrownBy(() -> stream.subscribe()).isSameAs(failure);
+
+ SpanData span = support.single();
+ assertThat(span.getStatus().getStatusCode()).isNotEqualTo(StatusCode.OK);
+ assertThat(support.spansEnded.get()).isEqualTo(1);
+ }
+ }
+
+ @Test
+ void concurrentCompleteAndCancelEndTheSpanExactlyOnce() throws InterruptedException {
+ int decided = 0;
+ for (int round = 0; round < 60 && decided < 8; round++) {
+ try (SpanTestSupport support = new SpanTestSupport()) {
+ CountDownLatch release = new CountDownLatch(1);
+ TracedChatModel model = new TracedChatModel(SpanTestSupport.streamModel(() ->
+ Flux.just(SpanTestSupport.response("x")).delaySubscription(Duration.ofMillis(5))
+ .doOnSubscribe(s -> release.countDown())), support.tracer, "synthetic");
+ AtomicReference subscription = new AtomicReference<>();
+ CountDownLatch downstreamSubscribed = new CountDownLatch(1);
+ model.stream(new Prompt("race")).subscribe(new BaseSubscriber<>() {
+ @Override
+ protected void hookOnSubscribe(Subscription s) {
+ subscription.set(s);
+ downstreamSubscribed.countDown();
+ s.request(Long.MAX_VALUE);
+ }
+ });
+ assertThat(downstreamSubscribed.await(5, TimeUnit.SECONDS)).isTrue();
+ release.await(5, TimeUnit.SECONDS);
+
+ CountDownLatch go = new CountDownLatch(1);
+ Thread canceller = new Thread(() -> {
+ await(go);
+ subscription.get().cancel();
+ });
+ canceller.start();
+ go.countDown();
+ canceller.join(5_000);
+
+ // whichever terminal wins, exactly one span is exported and it has ended
+ long deadline = System.currentTimeMillis() + 5_000;
+ while (support.finished().isEmpty() && System.currentTimeMillis() < deadline) {
+ Thread.sleep(10);
+ }
+ List spans = support.finished();
+ decided++;
+ assertThat(spans).as("round " + round).hasSize(1);
+ assertThat(support.spansEnded.get()).as("round " + round).isEqualTo(1);
+ assertThat(spans.get(0).hasEnded()).isTrue();
+ }
+ }
+ assertThat(decided).isGreaterThanOrEqualTo(8);
+ }
+
+ private static void await(CountDownLatch latch) {
+ try {
+ latch.await(5, TimeUnit.SECONDS);
+ } catch (InterruptedException e) {
+ Thread.currentThread().interrupt();
+ }
+ }
+}
diff --git a/java/traceai-spring-ai/src/test/java/ai/traceai/spring/TracingPrivacyTest.java b/java/traceai-spring-ai/src/test/java/ai/traceai/spring/TracingPrivacyTest.java
new file mode 100644
index 00000000..10bb006b
--- /dev/null
+++ b/java/traceai-spring-ai/src/test/java/ai/traceai/spring/TracingPrivacyTest.java
@@ -0,0 +1,106 @@
+package ai.traceai.spring;
+
+import ai.traceai.SemanticConventions;
+import ai.traceai.TraceConfig;
+import io.opentelemetry.api.common.AttributeKey;
+import io.opentelemetry.sdk.trace.data.SpanData;
+import org.junit.jupiter.api.Test;
+import org.springframework.ai.chat.model.ChatResponse;
+import org.springframework.ai.chat.prompt.Prompt;
+import reactor.core.publisher.Flux;
+import reactor.test.StepVerifier;
+
+import java.time.Duration;
+
+import static org.assertj.core.api.Assertions.assertThat;
+
+/** The four privacy switches are independent; both must be set to hide a channel completely. */
+class TracingPrivacyTest {
+
+ private static final Duration TIMEOUT = Duration.ofSeconds(5);
+ private static final AttributeKey INPUT_VALUE = AttributeKey.stringKey(SemanticConventions.INPUT_VALUE);
+ private static final AttributeKey INPUT_MESSAGES = AttributeKey.stringKey(SemanticConventions.LLM_INPUT_MESSAGES);
+ private static final AttributeKey OUTPUT_VALUE = AttributeKey.stringKey(SemanticConventions.OUTPUT_VALUE);
+ private static final AttributeKey OUTPUT_MESSAGES = AttributeKey.stringKey(SemanticConventions.LLM_OUTPUT_MESSAGES);
+
+ private static SpanData traced(TraceConfig config, Flux upstream) {
+ try (SpanTestSupport support = new SpanTestSupport(config)) {
+ StepVerifier.create(new TracedChatModel(SpanTestSupport.streamModel(() -> upstream), support.tracer, "synthetic")
+ .stream(new Prompt("secret input")))
+ .expectNextCount(1)
+ .expectComplete()
+ .verify(TIMEOUT);
+ return support.single();
+ }
+ }
+
+ @Test
+ void hideInputsAloneKeepsMessages() {
+ SpanData span = traced(TraceConfig.builder().hideInputs(true).build(), Flux.just(SpanTestSupport.response("visible")));
+ assertThat(span.getAttributes().get(INPUT_VALUE)).isNull();
+ assertThat(span.getAttributes().get(INPUT_MESSAGES)).contains("secret input");
+ assertThat(span.getAttributes().get(OUTPUT_VALUE)).isEqualTo("visible");
+ }
+
+ @Test
+ void hideInputMessagesAloneKeepsInputValue() {
+ SpanData span = traced(TraceConfig.builder().hideInputMessages(true).build(), Flux.just(SpanTestSupport.response("visible")));
+ assertThat(span.getAttributes().get(INPUT_VALUE)).isEqualTo("secret input");
+ assertThat(span.getAttributes().get(INPUT_MESSAGES)).isNull();
+ }
+
+ @Test
+ void hideOutputsAloneKeepsOutputMessages() {
+ SpanData span = traced(TraceConfig.builder().hideOutputs(true).build(), Flux.just(SpanTestSupport.response("visible")));
+ assertThat(span.getAttributes().get(OUTPUT_VALUE)).isNull();
+ assertThat(span.getAttributes().get(OUTPUT_MESSAGES)).contains("visible");
+ }
+
+ @Test
+ void bothOutputSwitchesAllocateNoBufferAndHideEverything() {
+ SpanData span = traced(TraceConfig.builder().hideOutputs(true).hideOutputMessages(true).build(),
+ Flux.just(SpanTestSupport.response("should-not-be-captured")));
+ assertThat(span.getAttributes().get(OUTPUT_VALUE)).isNull();
+ assertThat(span.getAttributes().get(OUTPUT_MESSAGES)).isNull();
+ }
+
+ @Test
+ void allFourSwitchesTogetherHideBothChannels() {
+ SpanData span = traced(TraceConfig.builder().hideInputs(true).hideInputMessages(true)
+ .hideOutputs(true).hideOutputMessages(true).build(),
+ Flux.just(SpanTestSupport.response("hidden")));
+ assertThat(span.getAttributes().get(INPUT_VALUE)).isNull();
+ assertThat(span.getAttributes().get(INPUT_MESSAGES)).isNull();
+ assertThat(span.getAttributes().get(OUTPUT_VALUE)).isNull();
+ assertThat(span.getAttributes().get(OUTPUT_MESSAGES)).isNull();
+ }
+
+ @Test
+ void allowedMetadataSurvivesEveryPrivacySwitch() {
+ try (SpanTestSupport support = new SpanTestSupport(TraceConfig.builder().hideInputs(true).hideInputMessages(true)
+ .hideOutputs(true).hideOutputMessages(true).build())) {
+ StepVerifier.create(new TracedChatModel(SpanTestSupport.streamModel(() -> Flux.just(SpanTestSupport.response("x"))),
+ support.tracer, "synthetic")
+ .stream(new Prompt("hi"))
+ .contextWrite(SpringAITracingContext.withContext(io.opentelemetry.context.Context.root(),
+ java.util.Map.of(SemanticConventions.USER_ID, "kept-user"))))
+ .expectNextCount(1).expectComplete().verify(TIMEOUT);
+ assertThat(support.single().getAttributes().get(AttributeKey.stringKey(SemanticConventions.USER_ID))).isEqualTo("kept-user");
+ }
+ }
+
+ @Test
+ void streamOutputIsTruncatedAtTheAttributeBudgetWithoutChangingDeliveredChunks() {
+ try (SpanTestSupport support = new SpanTestSupport()) {
+ String huge = "x".repeat(40_000);
+ ChatResponse chunk = SpanTestSupport.response(huge);
+ StepVerifier.create(new TracedChatModel(SpanTestSupport.streamModel(() -> Flux.just(chunk)), support.tracer, "synthetic")
+ .stream(new Prompt("hi")))
+ .expectNextMatches(delivered -> delivered == chunk && delivered.getResult().getOutput().getText().length() == 40_000)
+ .expectComplete()
+ .verify(TIMEOUT);
+ String captured = support.single().getAttributes().get(OUTPUT_VALUE);
+ assertThat(captured).hasSize(32_000).endsWith("...[truncated]");
+ }
+ }
+}
diff --git a/java/traceai-spring-boot-starter/pom.xml b/java/traceai-spring-boot-starter/pom.xml
index e823d672..126a0058 100644
--- a/java/traceai-spring-boot-starter/pom.xml
+++ b/java/traceai-spring-boot-starter/pom.xml
@@ -12,7 +12,36 @@
traceai-spring-boot-starter
TraceAI Spring Boot Starter
- Spring Boot auto-configuration for TraceAI Java SDK
+ Spring Boot auto-configuration for TraceAI Java SDK (Spring Boot 3.5.x / Spring AI 1.1.x)
+
+
+ 5.12.2
+
+
+
+
+
+ org.springframework.ai
+ spring-ai-bom
+ ${spring-ai.version}
+ pom
+ import
+
+
+ org.springframework.boot
+ spring-boot-dependencies
+ ${spring-boot.version}
+ pom
+ import
+
+
+ org.junit.jupiter
+ junit-jupiter
+ ${spring.junit.version}
+ test
+
+
+
@@ -26,24 +55,21 @@
${project.version}
-
+
org.springframework.boot
spring-boot-autoconfigure
- ${spring-boot.version}
org.springframework.boot
spring-boot-configuration-processor
- ${spring-boot.version}
true
-
+
org.springframework.ai
- spring-ai-core
- ${spring-ai.version}
+ spring-ai-model
true
@@ -53,10 +79,19 @@
junit-jupiter
test
+
+ org.assertj
+ assertj-core
+ test
+
org.springframework.boot
- spring-boot-starter-test
- ${spring-boot.version}
+ spring-boot-test
+ test
+
+
+ io.opentelemetry
+ opentelemetry-sdk-testing
test
@@ -64,9 +99,17 @@
- org.springframework.boot
- spring-boot-maven-plugin
- ${spring-boot.version}
+ org.apache.maven.plugins
+ maven-surefire-plugin
+
+
+ 1
+ false
+
+ --add-opens java.base/java.lang=ALL-UNNAMED
+ --add-opens java.base/java.util=ALL-UNNAMED
+
+
diff --git a/java/traceai-spring-boot-starter/src/test/java/ai/traceai/spring/autoconfigure/TraceAIAutoConfigurationTest.java b/java/traceai-spring-boot-starter/src/test/java/ai/traceai/spring/autoconfigure/TraceAIAutoConfigurationTest.java
new file mode 100644
index 00000000..bd571807
--- /dev/null
+++ b/java/traceai-spring-boot-starter/src/test/java/ai/traceai/spring/autoconfigure/TraceAIAutoConfigurationTest.java
@@ -0,0 +1,82 @@
+package ai.traceai.spring.autoconfigure;
+
+import ai.traceai.FITracer;
+import ai.traceai.TraceAI;
+import ai.traceai.spring.TracedChatModel;
+import org.junit.jupiter.api.Test;
+import org.springframework.ai.chat.model.ChatModel;
+import org.springframework.ai.chat.model.ChatResponse;
+import org.springframework.ai.chat.prompt.Prompt;
+import org.springframework.boot.autoconfigure.AutoConfigurations;
+import org.springframework.boot.test.context.runner.ApplicationContextRunner;
+import org.springframework.context.annotation.Bean;
+import org.springframework.context.annotation.Configuration;
+
+import static org.assertj.core.api.Assertions.assertThat;
+import static org.assertj.core.api.Assertions.assertThatThrownBy;
+
+/**
+ * Starter behaviour that does not touch the process-wide OpenTelemetry registration.
+ * Each case uses a fresh context; no exporter endpoint or credentials are configured.
+ */
+class TraceAIAutoConfigurationTest {
+
+ private final java.util.concurrent.atomic.AtomicInteger contexts = new java.util.concurrent.atomic.AtomicInteger();
+
+ /** A unique property per case defeats ApplicationContextRunner's context cache, so no case reuses another's initialized tracer. */
+ private ApplicationContextRunner runner() {
+ return new ApplicationContextRunner()
+ .withConfiguration(AutoConfigurations.of(TraceAIAutoConfiguration.class))
+ .withPropertyValues("traceai.project-name=unit-" + contexts.incrementAndGet());
+ }
+
+ @Test
+ void enabledByDefaultCreatesATracerBean() {
+ runner().run(context -> {
+ assertThat(context).hasSingleBean(FITracer.class);
+ assertThat(TraceAI.isInitialized()).isTrue();
+ });
+ }
+
+ @Test
+ void explicitEnabledTrueCreatesATracerBean() {
+ runner().withPropertyValues("traceai.enabled=true", "traceai.project-name=unit").run(context -> {
+ assertThat(context).hasSingleBean(FITracer.class);
+ assertThat(context.getBean(FITracer.class)).isSameAs(TraceAI.getTracer());
+ });
+ }
+
+ @Test
+ void autoConfigurationDoesNotWrapUserModels() {
+ runner().withUserConfiguration(ChatModelConfiguration.class).run(context -> {
+ assertThat(context).hasSingleBean(ChatModel.class);
+ assertThat(context.getBean(ChatModel.class)).isInstanceOf(PlainChatModel.class);
+ });
+ }
+
+ @Test
+ void manualWrappingUsesTheExplicitTracer() {
+ runner().run(context -> {
+ FITracer tracer = context.getBean(FITracer.class);
+ PlainChatModel delegate = new PlainChatModel();
+ TracedChatModel wrapped = new TracedChatModel(delegate, tracer, "synthetic");
+ assertThat(wrapped.unwrap()).isSameAs(delegate);
+ assertThat(wrapped.call(new Prompt("hi"))).isNull();
+ });
+ }
+
+ @Configuration
+ static class ChatModelConfiguration {
+ @Bean
+ ChatModel chatModel() {
+ return new PlainChatModel();
+ }
+ }
+
+ static final class PlainChatModel implements ChatModel {
+ @Override
+ public ChatResponse call(Prompt prompt) {
+ return null;
+ }
+ }
+}
diff --git a/java/traceai-spring-boot-starter/src/test/java/ai/traceai/spring/autoconfigure/TraceAICustomTracerTest.java b/java/traceai-spring-boot-starter/src/test/java/ai/traceai/spring/autoconfigure/TraceAICustomTracerTest.java
new file mode 100644
index 00000000..6d5e0242
--- /dev/null
+++ b/java/traceai-spring-boot-starter/src/test/java/ai/traceai/spring/autoconfigure/TraceAICustomTracerTest.java
@@ -0,0 +1,37 @@
+package ai.traceai.spring.autoconfigure;
+
+import ai.traceai.FITracer;
+import ai.traceai.TraceAI;
+import org.junit.jupiter.api.Test;
+import org.springframework.boot.autoconfigure.AutoConfigurations;
+import org.springframework.boot.test.context.runner.ApplicationContextRunner;
+import org.springframework.context.annotation.Bean;
+import org.springframework.context.annotation.Configuration;
+
+import static org.assertj.core.api.Assertions.assertThat;
+
+/** Fresh JVM: an application-supplied FITracer must win and the global tracer must stay uninitialized. */
+class TraceAICustomTracerTest {
+
+ @Test
+ void existingCustomTracerBeanBacksOffWithoutInitializingTheGlobalTracer() {
+ new ApplicationContextRunner()
+ .withConfiguration(AutoConfigurations.of(TraceAIAutoConfiguration.class))
+ .withUserConfiguration(CustomTracerConfiguration.class)
+ .run(context -> {
+ assertThat(context).hasSingleBean(FITracer.class);
+ assertThat(context.getBean(FITracer.class)).isSameAs(CustomTracerConfiguration.CUSTOM);
+ assertThat(TraceAI.isInitialized()).isFalse();
+ });
+ }
+
+ @Configuration
+ static class CustomTracerConfiguration {
+ static final FITracer CUSTOM = new FITracer(io.opentelemetry.api.OpenTelemetry.noop().getTracer("custom"));
+
+ @Bean
+ FITracer fiTracer() {
+ return CUSTOM;
+ }
+ }
+}
diff --git a/java/traceai-spring-boot-starter/src/test/java/ai/traceai/spring/autoconfigure/TraceAIDisabledTest.java b/java/traceai-spring-boot-starter/src/test/java/ai/traceai/spring/autoconfigure/TraceAIDisabledTest.java
new file mode 100644
index 00000000..b7c88ab3
--- /dev/null
+++ b/java/traceai-spring-boot-starter/src/test/java/ai/traceai/spring/autoconfigure/TraceAIDisabledTest.java
@@ -0,0 +1,25 @@
+package ai.traceai.spring.autoconfigure;
+
+import ai.traceai.FITracer;
+import ai.traceai.TraceAI;
+import org.junit.jupiter.api.Test;
+import org.springframework.boot.autoconfigure.AutoConfigurations;
+import org.springframework.boot.test.context.runner.ApplicationContextRunner;
+
+import static org.assertj.core.api.Assertions.assertThat;
+
+/** Fresh JVM: disabled auto-configuration must not create a tracer or touch the global registration. */
+class TraceAIDisabledTest {
+
+ @Test
+ void disabledCreatesNoTracerBeanAndDoesNotInitializeTheGlobalTracer() {
+ new ApplicationContextRunner()
+ .withConfiguration(AutoConfigurations.of(TraceAIAutoConfiguration.class))
+ .withPropertyValues("traceai.enabled=false")
+ .run(context -> {
+ assertThat(context).doesNotHaveBean(FITracer.class);
+ assertThat(TraceAI.isInitialized()).isFalse();
+ assertThat(context).hasNotFailed();
+ });
+ }
+}
diff --git a/java/traceai-spring-boot-starter/src/test/java/ai/traceai/spring/autoconfigure/TraceAIUninitializedTest.java b/java/traceai-spring-boot-starter/src/test/java/ai/traceai/spring/autoconfigure/TraceAIUninitializedTest.java
new file mode 100644
index 00000000..d54e35c0
--- /dev/null
+++ b/java/traceai-spring-boot-starter/src/test/java/ai/traceai/spring/autoconfigure/TraceAIUninitializedTest.java
@@ -0,0 +1,28 @@
+package ai.traceai.spring.autoconfigure;
+
+import ai.traceai.TraceAI;
+import ai.traceai.spring.TracedChatModel;
+import org.junit.jupiter.api.Test;
+import org.springframework.ai.chat.model.ChatModel;
+import org.springframework.ai.chat.model.ChatResponse;
+import org.springframework.ai.chat.prompt.Prompt;
+
+import static org.assertj.core.api.Assertions.assertThat;
+import static org.assertj.core.api.Assertions.assertThatThrownBy;
+
+/** Fresh JVM: the two-argument wrapper must fail closed when nobody initialized TraceAI. */
+class TraceAIUninitializedTest {
+
+ @Test
+ void twoArgConstructorFailsWhenTheGlobalTracerWasNeverInitialized() {
+ assertThat(TraceAI.isInitialized()).isFalse();
+ assertThatThrownBy(() -> new TracedChatModel(new ChatModel() {
+ @Override
+ public ChatResponse call(Prompt prompt) {
+ return null;
+ }
+ }, "synthetic"))
+ .isInstanceOf(IllegalStateException.class)
+ .hasMessageContaining("not initialized");
+ }
+}