From 3c5d2a15a33e85acc71c812e6156e9dd10be8e21 Mon Sep 17 00:00:00 2001 From: Nikhil Pareek Date: Sat, 3 Oct 2026 20:45:18 -0700 Subject: [PATCH 1/7] fix(spring-ai): restore binary compatibility with Spring AI 1.1.x The wrapper was compiled against Spring AI 1.0.0-M4 and called APIs removed before GA (Message.getContent, Usage.getGenerationTokens), so every current release threw NoSuchMethodError on the first chat call. - Compile against Spring AI 1.1.8 / Spring Boot 3.5.15; depend on spring-ai-model - Read text with getText() and completion usage with getCompletionTokens() - Defer stream spans and delegate calls until subscription; one span per subscription, ended exactly once on complete, error or cancel - Record token counts only when the provider reports them; treat EmptyUsage zeros as absent and preserve measured zero - Delegate every embedding overload directly, including the 1.1.8-only getEmbeddingContent accessor (no span) - Add SpringAITracingContext.withContext for explicit parent/attribute binding - Bound stream capture and keep the four privacy switches independent - Prove it with a deterministic suite plus a same-binary consumer on 1.1.8 and 1.1.0, and a real local Ollama tool-agent (sync and stream) TH-7477 --- .../java-spring-ai-compatibility.yml | 51 +++ java/README.md | 4 +- java/compatibility/spring-ai/MIGRATION.md | 35 ++ .../spring-ai/baseline-cell/pom.xml | 24 ++ .../java/local/th7477/CommonConsumer.java | 59 +++ .../java/local/th7477/SpanTestSupport.java | 58 +++ .../java/local/th7477/CommonConsumer.java | 59 +++ .../java/local/th7477/SpanTestSupport.java | 58 +++ java/compatibility/spring-ai/common/pom.xml | 24 ++ .../java/local/th7477/CommonConsumer.java | 59 +++ .../java/local/th7477/SpanTestSupport.java | 58 +++ java/compatibility/spring-ai/primary/pom.xml | 26 ++ .../java/local/th7477/CommonConsumer.java | 59 +++ .../java/local/th7477/PrimaryConsumer.java | 23 ++ .../java/local/th7477/RealAgentCandidate.java | 90 +++++ .../java/local/th7477/SpanTestSupport.java | 58 +++ .../java/local/th7477/CommonConsumer.java | 59 +++ .../java/local/th7477/SpanTestSupport.java | 58 +++ java/examples/spring-ai-example/pom.xml | 35 +- .../examples/SpringAIExampleApplication.java | 8 +- java/pom.xml | 4 +- java/traceai-spring-ai/pom.xml | 72 +++- .../spring/SpringAITracingContext.java | 181 +++++++++ .../traceai/spring/StreamOutputCapture.java | 82 +++++ .../ai/traceai/spring/TracedChatModel.java | 342 ++++++++++-------- .../traceai/spring/TracedEmbeddingModel.java | 278 +++++++++----- .../ai/traceai/spring/TracingSubscriber.java | 178 +++++++++ .../ai/traceai/spring/SpanTestSupport.java | 177 +++++++++ .../spring/SpringAITracingContextTest.java | 185 ++++++++++ .../spring/TracedChatModelStreamTest.java | 198 ++++++++++ .../traceai/spring/TracedChatModelTest.java | 198 ++++++++++ .../spring/TracedEmbeddingModelTest.java | 164 +++++++++ .../traceai/spring/TracingLifecycleTest.java | 250 +++++++++++++ .../ai/traceai/spring/TracingPrivacyTest.java | 106 ++++++ java/traceai-spring-boot-starter/pom.xml | 67 +++- 35 files changed, 3103 insertions(+), 284 deletions(-) create mode 100644 .github/workflows/java-spring-ai-compatibility.yml create mode 100644 java/compatibility/spring-ai/MIGRATION.md create mode 100644 java/compatibility/spring-ai/baseline-cell/pom.xml create mode 100644 java/compatibility/spring-ai/baseline-cell/src/main/java/local/th7477/CommonConsumer.java create mode 100644 java/compatibility/spring-ai/baseline-cell/src/main/java/local/th7477/SpanTestSupport.java create mode 100644 java/compatibility/spring-ai/baseline-cell/src/src/main/java/local/th7477/CommonConsumer.java create mode 100644 java/compatibility/spring-ai/baseline-cell/src/src/main/java/local/th7477/SpanTestSupport.java create mode 100644 java/compatibility/spring-ai/common/pom.xml create mode 100644 java/compatibility/spring-ai/common/src/main/java/local/th7477/CommonConsumer.java create mode 100644 java/compatibility/spring-ai/common/src/main/java/local/th7477/SpanTestSupport.java create mode 100644 java/compatibility/spring-ai/primary/pom.xml create mode 100644 java/compatibility/spring-ai/primary/src/main/java/local/th7477/CommonConsumer.java create mode 100644 java/compatibility/spring-ai/primary/src/main/java/local/th7477/PrimaryConsumer.java create mode 100644 java/compatibility/spring-ai/primary/src/main/java/local/th7477/RealAgentCandidate.java create mode 100644 java/compatibility/spring-ai/primary/src/main/java/local/th7477/SpanTestSupport.java create mode 100644 java/compatibility/spring-ai/primary/src/src/main/java/local/th7477/CommonConsumer.java create mode 100644 java/compatibility/spring-ai/primary/src/src/main/java/local/th7477/SpanTestSupport.java create mode 100644 java/traceai-spring-ai/src/main/java/ai/traceai/spring/SpringAITracingContext.java create mode 100644 java/traceai-spring-ai/src/main/java/ai/traceai/spring/StreamOutputCapture.java create mode 100644 java/traceai-spring-ai/src/main/java/ai/traceai/spring/TracingSubscriber.java create mode 100644 java/traceai-spring-ai/src/test/java/ai/traceai/spring/SpanTestSupport.java create mode 100644 java/traceai-spring-ai/src/test/java/ai/traceai/spring/SpringAITracingContextTest.java create mode 100644 java/traceai-spring-ai/src/test/java/ai/traceai/spring/TracedChatModelStreamTest.java create mode 100644 java/traceai-spring-ai/src/test/java/ai/traceai/spring/TracedChatModelTest.java create mode 100644 java/traceai-spring-ai/src/test/java/ai/traceai/spring/TracedEmbeddingModelTest.java create mode 100644 java/traceai-spring-ai/src/test/java/ai/traceai/spring/TracingLifecycleTest.java create mode 100644 java/traceai-spring-ai/src/test/java/ai/traceai/spring/TracingPrivacyTest.java diff --git a/.github/workflows/java-spring-ai-compatibility.yml b/.github/workflows/java-spring-ai-compatibility.yml new file mode 100644 index 00000000..a4b256b9 --- /dev/null +++ b/.github/workflows/java-spring-ai-compatibility.yml @@ -0,0 +1,51 @@ +name: Java Spring AI compatibility + +# Deterministic Spring AI 1.1.x compatibility for the traceAI Java SDK. +# Builds the candidate once on Spring AI 1.1.8 / Spring Boot 3.5.15 and runs the +# same JARs against Spring AI 1.1.0 / Spring Boot 3.5.7. No secrets, no live model. + +on: + pull_request: + paths: + - 'java/pom.xml' + - 'java/traceai-spring-ai/**' + - 'java/traceai-spring-boot-starter/**' + - 'java/traceai-java-core/**' + - 'java/compatibility/spring-ai/**' + - 'java/examples/spring-ai-example/**' + - '.github/workflows/java-spring-ai-compatibility.yml' + +permissions: + contents: read + +jobs: + deterministic: + runs-on: ubuntu-latest + container: + image: maven@sha256:e4a7ace3dc0d645ed97f8d9ad0b0d3f0b14fa8d150138f27f116d7105a639b82 + steps: + - uses: actions/checkout@v4 + - name: Deterministic Spring module suite (Spring AI 1.1.8) + working-directory: java + run: mvn -B -ntp -pl traceai-spring-ai -am -Dsurefire.failIfNoSpecifiedTests=false test + - name: Package candidate JARs + working-directory: java + run: mvn -B -ntp -pl traceai-spring-ai,traceai-spring-boot-starter -am -DskipTests package + - name: Stage candidate artifacts + run: | + mkdir -p candidate-jars + cp java/traceai-java-core/target/traceai-java-core-1.0.0.jar candidate-jars/ + cp java/traceai-spring-ai/target/traceai-spring-ai-1.0.0.jar candidate-jars/ + sha256sum candidate-jars/*.jar + - name: Same-binary consumer on Spring AI 1.1.8 + working-directory: java/compatibility/spring-ai/common + run: mvn -B -ntp -q exec:java -Dexec.mainClass=local.th7477.CommonConsumer -Dexec.args=P -Dexec.classpathScope=compile + - name: Same-binary consumer on Spring AI 1.1.0 + working-directory: java/compatibility/spring-ai/baseline-cell + run: mvn -B -ntp -o -q exec:java -Dexec.mainClass=local.th7477.CommonConsumer -Dexec.args=B -Dexec.classpathScope=compile + - name: 1.1.8-only embedding accessor + working-directory: java/compatibility/spring-ai/primary + run: mvn -B -ntp -o -q exec:java -Dexec.mainClass=local.th7477.PrimaryConsumer -Dexec.classpathScope=compile + - name: Example compiles against 1.1.8 + working-directory: java/examples/spring-ai-example + run: mvn -B -ntp -DskipTests compile diff --git a/java/README.md b/java/README.md index 369b6fdf..2e8d2094 100644 --- a/java/README.md +++ b/java/README.md @@ -357,8 +357,8 @@ traceai: | Module | Description | Integration | |--------|-------------|-------------| | `traceai-langchain4j` | LangChain4j instrumentation | `dev.langchain4j:langchain4j` | -| `traceai-spring-ai` | Spring AI instrumentation | `org.springframework.ai:spring-ai-core` | -| `traceai-spring-boot-starter` | Spring Boot auto-configuration | Spring Boot 3.x | +| `traceai-spring-ai` | Spring AI 1.1.x instrumentation | `org.springframework.ai:spring-ai-model` (provided) | +| `traceai-spring-boot-starter` | Spring Boot auto-configuration | Spring Boot 3.5.x | | `traceai-java-semantic-kernel` | Microsoft Semantic Kernel instrumentation | `com.microsoft.semantic-kernel:semantickernel-api` | ## Configuration diff --git a/java/compatibility/spring-ai/MIGRATION.md b/java/compatibility/spring-ai/MIGRATION.md new file mode 100644 index 00000000..70760210 --- /dev/null +++ b/java/compatibility/spring-ai/MIGRATION.md @@ -0,0 +1,35 @@ +# Spring AI compatibility notes (TH-7477) + +The `traceai-spring-ai` wrapper is compiled and tested against Spring AI 1.1.x GA. Spring AI `1.0.0-M4` +is retired: the milestone API (`Message.getContent()`, `Usage.getGenerationTokens()`, `spring-ai-core`) +was removed before GA, so an artifact built on it fails at runtime on every current release. + +## Closed matrix + +| Cell | Spring AI | Spring Boot | Java | What runs | +|---|---|---|---|---| +| P | 1.1.8 | 3.5.15 | 17 | compile once, deterministic suite, example, real local agent | +| B | 1.1.0 | 3.5.7 | 17 | the same built JARs, common API only | + +Other versions (Spring AI 1.0 GA, 2.0, Spring Boot 3.4 or 4, newer JDKs) are not claimed. + +## What changed for callers + +- Message and assistant text is read with `getText()`, not `getContent()`. +- Completion usage is read with `getCompletionTokens()`, not `getGenerationTokens()`. +- Depend on `org.springframework.ai:spring-ai-model` (the wrapper declares it `provided`). +- `TracedEmbeddingModel.getEmbeddingContent(Document)` delegates directly and exists only on 1.1.8+. + Applications on 1.1.0 simply never call it; no reflection bridge is provided. + +## Streaming + +`stream(prompt)` starts no span and does not call the delegate until a subscriber subscribes. Each +subscription owns one span that ends exactly once. Streamed token usage is not aggregated. To bind a +parent span and TraceAI session/user attributes to a stream subscribed on another thread, use +`SpringAITracingContext.withContext(parent, attributes)` with `Flux.contextWrite`. Reusing one publisher +across tenants requires an explicit snapshot per subscription. + +## Privacy + +`hideInputs` and `hideInputMessages` are independent, as are the two output switches. Set both of a pair +to omit a channel completely. Metadata, session and user attributes are unaffected. diff --git a/java/compatibility/spring-ai/baseline-cell/pom.xml b/java/compatibility/spring-ai/baseline-cell/pom.xml new file mode 100644 index 00000000..4326e269 --- /dev/null +++ b/java/compatibility/spring-ai/baseline-cell/pom.xml @@ -0,0 +1,24 @@ + + + 4.0.0 + local.th7477common-consumer1.0-SNAPSHOT + + 17UTF-8 + 1.1.03.5.7 + + + + org.springframework.aispring-ai-bom${spring-ai.version}pomimport + org.springframework.bootspring-boot-dependencies${spring-boot.version}pomimport + + + + org.springframework.aispring-ai-model + org.springframework.aispring-ai-client-chat + ai.traceaitraceai-java-core1.0.0-candidatesystem/candidate-jars/traceai-java-core-1.0.0.jar + ai.traceaitraceai-spring-ai1.0.0-candidatesystem/candidate-jars/traceai-spring-ai-1.0.0.jar + io.opentelemetryopentelemetry-sdk-testing1.49.0 + com.google.code.gsongson2.10.1 + + org.apache.maven.pluginsmaven-compiler-plugin3.12.1 + diff --git a/java/compatibility/spring-ai/baseline-cell/src/main/java/local/th7477/CommonConsumer.java b/java/compatibility/spring-ai/baseline-cell/src/main/java/local/th7477/CommonConsumer.java new file mode 100644 index 00000000..22058daf --- /dev/null +++ b/java/compatibility/spring-ai/baseline-cell/src/main/java/local/th7477/CommonConsumer.java @@ -0,0 +1,59 @@ +package local.th7477; + +import ai.traceai.FITracer; +import ai.traceai.SemanticConventions; +import ai.traceai.spring.TracedChatModel; +import ai.traceai.spring.TracedEmbeddingModel; +import io.opentelemetry.api.common.AttributeKey; +import io.opentelemetry.sdk.testing.exporter.InMemorySpanExporter; +import io.opentelemetry.sdk.trace.SdkTracerProvider; +import io.opentelemetry.sdk.trace.data.SpanData; +import io.opentelemetry.sdk.trace.export.SimpleSpanProcessor; +import org.springframework.ai.chat.messages.UserMessage; +import org.springframework.ai.chat.model.ChatResponse; +import org.springframework.ai.chat.prompt.Prompt; +import org.springframework.ai.embedding.EmbeddingRequest; + +import java.security.CodeSource; +import java.security.MessageDigest; +import java.util.HexFormat; +import java.util.List; + +/** Same compiled classes, same candidate JARs, run on both Spring AI cells. Asserts common API only. */ +public final class CommonConsumer { + public static void main(String[] args) throws Exception { + String cell = args.length > 0 ? args[0] : "unknown"; + InMemorySpanExporter exporter = InMemorySpanExporter.create(); + SdkTracerProvider provider = SdkTracerProvider.builder().addSpanProcessor(SimpleSpanProcessor.create(exporter)).build(); + FITracer tracer = new FITracer(provider.get("th7477-matrix")); + + ChatResponse expected = SpanTestSupport.response("matrix-answer", SpanTestSupport.usage(4, 2, 6), "matrix-model"); + TracedChatModel traced = new TracedChatModel(SpanTestSupport.syncModel(p -> expected), tracer, "synthetic"); + ChatResponse actual = traced.call(new Prompt(List.of(new UserMessage("matrix input " + cell)))); + if (actual != expected) throw new AssertionError("call identity lost on " + cell); + + SpanData span = exporter.getFinishedSpanItems().get(exporter.getFinishedSpanItems().size() - 1); + eq(span.getAttributes().get(AttributeKey.stringKey(SemanticConventions.OUTPUT_VALUE)), "matrix-answer", "output"); + eq(span.getAttributes().get(AttributeKey.longKey(SemanticConventions.LLM_TOKEN_COUNT_PROMPT)), 4L, "prompt tokens"); + eq(span.getAttributes().get(AttributeKey.longKey(SemanticConventions.LLM_TOKEN_COUNT_COMPLETION)), 2L, "completion tokens"); + eq(span.getAttributes().get(AttributeKey.longKey(SemanticConventions.LLM_TOKEN_COUNT_TOTAL)), 6L, "total tokens"); + + traced.stream(new Prompt("stream " + cell)).blockLast(); + long streamSpans = exporter.getFinishedSpanItems().stream().filter(s -> s.getName().contains("Stream")).count(); + if (streamSpans != 1) throw new AssertionError("expected one stream span, got " + streamSpans); + + TracedEmbeddingModel tracedEmbeddings = new TracedEmbeddingModel(new SpanTestSupport.FixedEmbeddingModel(), tracer, "synthetic"); + List vectors = tracedEmbeddings.embed(List.of("a", "b")); + if (vectors.size() != 2 || vectors.get(0).length != 3) throw new AssertionError("batch identity lost"); + tracedEmbeddings.call(new EmbeddingRequest(List.of("a"), null)); + + CodeSource source = TracedChatModel.class.getProtectionDomain().getCodeSource(); + String digest = HexFormat.of().formatHex(MessageDigest.getInstance("SHA-256").digest(source.getLocation().openStream().readAllBytes())); + System.out.println("CELL=" + cell + " SHA256=" + digest + " SPANS=" + exporter.getFinishedSpanItems().size() + " COMMON_CONSUMER_PASS"); + provider.close(); + } + + private static void eq(Object actual, Object expected, String what) { + if (actual == null || !actual.equals(expected)) throw new AssertionError(what + " expected " + expected + " but was " + actual); + } +} diff --git a/java/compatibility/spring-ai/baseline-cell/src/main/java/local/th7477/SpanTestSupport.java b/java/compatibility/spring-ai/baseline-cell/src/main/java/local/th7477/SpanTestSupport.java new file mode 100644 index 00000000..af34a897 --- /dev/null +++ b/java/compatibility/spring-ai/baseline-cell/src/main/java/local/th7477/SpanTestSupport.java @@ -0,0 +1,58 @@ +package local.th7477; + +import org.springframework.ai.chat.messages.AssistantMessage; +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.document.Document; +import org.springframework.ai.embedding.Embedding; +import org.springframework.ai.embedding.EmbeddingModel; +import org.springframework.ai.embedding.EmbeddingRequest; +import org.springframework.ai.embedding.EmbeddingResponse; +import org.springframework.ai.embedding.EmbeddingResponseMetadata; +import reactor.core.publisher.Flux; + +import java.util.List; +import java.util.function.Function; + +/** Spring AI objects shared by both matrix cells. Only APIs present in 1.1.0 and 1.1.8. */ +final class SpanTestSupport { + private SpanTestSupport() {} + + 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 Usage usage(int prompt, int completion, int total) { + 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 "matrix"; } + }; + } + + static ChatModel syncModel(Function call) { + return new ChatModel() { + @Override public ChatResponse call(Prompt prompt) { return call.apply(prompt); } + @Override public Flux stream(Prompt prompt) { return Flux.just(response("streamed-" + prompt.getContents(), null, null)); } + @Override public ChatOptions getDefaultOptions() { return null; } + }; + } + + static final class FixedEmbeddingModel implements EmbeddingModel { + private final EmbeddingResponse response = new EmbeddingResponse( + List.of(new Embedding(new float[]{1f, 2f, 3f}, 0), new Embedding(new float[]{4f, 5f, 6f}, 1)), + new EmbeddingResponseMetadata("m", null)); + @Override public EmbeddingResponse call(EmbeddingRequest request) { return response; } + @Override public float[] embed(Document document) { return new float[]{1f, 2f, 3f}; } + @Override public List embed(List texts) { return List.of(new float[]{1f, 2f, 3f}, new float[]{4f, 5f, 6f}); } + } +} diff --git a/java/compatibility/spring-ai/baseline-cell/src/src/main/java/local/th7477/CommonConsumer.java b/java/compatibility/spring-ai/baseline-cell/src/src/main/java/local/th7477/CommonConsumer.java new file mode 100644 index 00000000..22058daf --- /dev/null +++ b/java/compatibility/spring-ai/baseline-cell/src/src/main/java/local/th7477/CommonConsumer.java @@ -0,0 +1,59 @@ +package local.th7477; + +import ai.traceai.FITracer; +import ai.traceai.SemanticConventions; +import ai.traceai.spring.TracedChatModel; +import ai.traceai.spring.TracedEmbeddingModel; +import io.opentelemetry.api.common.AttributeKey; +import io.opentelemetry.sdk.testing.exporter.InMemorySpanExporter; +import io.opentelemetry.sdk.trace.SdkTracerProvider; +import io.opentelemetry.sdk.trace.data.SpanData; +import io.opentelemetry.sdk.trace.export.SimpleSpanProcessor; +import org.springframework.ai.chat.messages.UserMessage; +import org.springframework.ai.chat.model.ChatResponse; +import org.springframework.ai.chat.prompt.Prompt; +import org.springframework.ai.embedding.EmbeddingRequest; + +import java.security.CodeSource; +import java.security.MessageDigest; +import java.util.HexFormat; +import java.util.List; + +/** Same compiled classes, same candidate JARs, run on both Spring AI cells. Asserts common API only. */ +public final class CommonConsumer { + public static void main(String[] args) throws Exception { + String cell = args.length > 0 ? args[0] : "unknown"; + InMemorySpanExporter exporter = InMemorySpanExporter.create(); + SdkTracerProvider provider = SdkTracerProvider.builder().addSpanProcessor(SimpleSpanProcessor.create(exporter)).build(); + FITracer tracer = new FITracer(provider.get("th7477-matrix")); + + ChatResponse expected = SpanTestSupport.response("matrix-answer", SpanTestSupport.usage(4, 2, 6), "matrix-model"); + TracedChatModel traced = new TracedChatModel(SpanTestSupport.syncModel(p -> expected), tracer, "synthetic"); + ChatResponse actual = traced.call(new Prompt(List.of(new UserMessage("matrix input " + cell)))); + if (actual != expected) throw new AssertionError("call identity lost on " + cell); + + SpanData span = exporter.getFinishedSpanItems().get(exporter.getFinishedSpanItems().size() - 1); + eq(span.getAttributes().get(AttributeKey.stringKey(SemanticConventions.OUTPUT_VALUE)), "matrix-answer", "output"); + eq(span.getAttributes().get(AttributeKey.longKey(SemanticConventions.LLM_TOKEN_COUNT_PROMPT)), 4L, "prompt tokens"); + eq(span.getAttributes().get(AttributeKey.longKey(SemanticConventions.LLM_TOKEN_COUNT_COMPLETION)), 2L, "completion tokens"); + eq(span.getAttributes().get(AttributeKey.longKey(SemanticConventions.LLM_TOKEN_COUNT_TOTAL)), 6L, "total tokens"); + + traced.stream(new Prompt("stream " + cell)).blockLast(); + long streamSpans = exporter.getFinishedSpanItems().stream().filter(s -> s.getName().contains("Stream")).count(); + if (streamSpans != 1) throw new AssertionError("expected one stream span, got " + streamSpans); + + TracedEmbeddingModel tracedEmbeddings = new TracedEmbeddingModel(new SpanTestSupport.FixedEmbeddingModel(), tracer, "synthetic"); + List vectors = tracedEmbeddings.embed(List.of("a", "b")); + if (vectors.size() != 2 || vectors.get(0).length != 3) throw new AssertionError("batch identity lost"); + tracedEmbeddings.call(new EmbeddingRequest(List.of("a"), null)); + + CodeSource source = TracedChatModel.class.getProtectionDomain().getCodeSource(); + String digest = HexFormat.of().formatHex(MessageDigest.getInstance("SHA-256").digest(source.getLocation().openStream().readAllBytes())); + System.out.println("CELL=" + cell + " SHA256=" + digest + " SPANS=" + exporter.getFinishedSpanItems().size() + " COMMON_CONSUMER_PASS"); + provider.close(); + } + + private static void eq(Object actual, Object expected, String what) { + if (actual == null || !actual.equals(expected)) throw new AssertionError(what + " expected " + expected + " but was " + actual); + } +} diff --git a/java/compatibility/spring-ai/baseline-cell/src/src/main/java/local/th7477/SpanTestSupport.java b/java/compatibility/spring-ai/baseline-cell/src/src/main/java/local/th7477/SpanTestSupport.java new file mode 100644 index 00000000..af34a897 --- /dev/null +++ b/java/compatibility/spring-ai/baseline-cell/src/src/main/java/local/th7477/SpanTestSupport.java @@ -0,0 +1,58 @@ +package local.th7477; + +import org.springframework.ai.chat.messages.AssistantMessage; +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.document.Document; +import org.springframework.ai.embedding.Embedding; +import org.springframework.ai.embedding.EmbeddingModel; +import org.springframework.ai.embedding.EmbeddingRequest; +import org.springframework.ai.embedding.EmbeddingResponse; +import org.springframework.ai.embedding.EmbeddingResponseMetadata; +import reactor.core.publisher.Flux; + +import java.util.List; +import java.util.function.Function; + +/** Spring AI objects shared by both matrix cells. Only APIs present in 1.1.0 and 1.1.8. */ +final class SpanTestSupport { + private SpanTestSupport() {} + + 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 Usage usage(int prompt, int completion, int total) { + 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 "matrix"; } + }; + } + + static ChatModel syncModel(Function call) { + return new ChatModel() { + @Override public ChatResponse call(Prompt prompt) { return call.apply(prompt); } + @Override public Flux stream(Prompt prompt) { return Flux.just(response("streamed-" + prompt.getContents(), null, null)); } + @Override public ChatOptions getDefaultOptions() { return null; } + }; + } + + static final class FixedEmbeddingModel implements EmbeddingModel { + private final EmbeddingResponse response = new EmbeddingResponse( + List.of(new Embedding(new float[]{1f, 2f, 3f}, 0), new Embedding(new float[]{4f, 5f, 6f}, 1)), + new EmbeddingResponseMetadata("m", null)); + @Override public EmbeddingResponse call(EmbeddingRequest request) { return response; } + @Override public float[] embed(Document document) { return new float[]{1f, 2f, 3f}; } + @Override public List embed(List texts) { return List.of(new float[]{1f, 2f, 3f}, new float[]{4f, 5f, 6f}); } + } +} diff --git a/java/compatibility/spring-ai/common/pom.xml b/java/compatibility/spring-ai/common/pom.xml new file mode 100644 index 00000000..8cdcc03e --- /dev/null +++ b/java/compatibility/spring-ai/common/pom.xml @@ -0,0 +1,24 @@ + + + 4.0.0 + local.th7477common-consumer1.0-SNAPSHOT + + 17UTF-8 + 1.1.83.5.15 + + + + org.springframework.aispring-ai-bom${spring-ai.version}pomimport + org.springframework.bootspring-boot-dependencies${spring-boot.version}pomimport + + + + org.springframework.aispring-ai-model + org.springframework.aispring-ai-client-chat + ai.traceaitraceai-java-core1.0.0-candidatesystem/candidate-jars/traceai-java-core-1.0.0.jar + ai.traceaitraceai-spring-ai1.0.0-candidatesystem/candidate-jars/traceai-spring-ai-1.0.0.jar + io.opentelemetryopentelemetry-sdk-testing1.49.0 + com.google.code.gsongson2.10.1 + + org.apache.maven.pluginsmaven-compiler-plugin3.12.1 + diff --git a/java/compatibility/spring-ai/common/src/main/java/local/th7477/CommonConsumer.java b/java/compatibility/spring-ai/common/src/main/java/local/th7477/CommonConsumer.java new file mode 100644 index 00000000..22058daf --- /dev/null +++ b/java/compatibility/spring-ai/common/src/main/java/local/th7477/CommonConsumer.java @@ -0,0 +1,59 @@ +package local.th7477; + +import ai.traceai.FITracer; +import ai.traceai.SemanticConventions; +import ai.traceai.spring.TracedChatModel; +import ai.traceai.spring.TracedEmbeddingModel; +import io.opentelemetry.api.common.AttributeKey; +import io.opentelemetry.sdk.testing.exporter.InMemorySpanExporter; +import io.opentelemetry.sdk.trace.SdkTracerProvider; +import io.opentelemetry.sdk.trace.data.SpanData; +import io.opentelemetry.sdk.trace.export.SimpleSpanProcessor; +import org.springframework.ai.chat.messages.UserMessage; +import org.springframework.ai.chat.model.ChatResponse; +import org.springframework.ai.chat.prompt.Prompt; +import org.springframework.ai.embedding.EmbeddingRequest; + +import java.security.CodeSource; +import java.security.MessageDigest; +import java.util.HexFormat; +import java.util.List; + +/** Same compiled classes, same candidate JARs, run on both Spring AI cells. Asserts common API only. */ +public final class CommonConsumer { + public static void main(String[] args) throws Exception { + String cell = args.length > 0 ? args[0] : "unknown"; + InMemorySpanExporter exporter = InMemorySpanExporter.create(); + SdkTracerProvider provider = SdkTracerProvider.builder().addSpanProcessor(SimpleSpanProcessor.create(exporter)).build(); + FITracer tracer = new FITracer(provider.get("th7477-matrix")); + + ChatResponse expected = SpanTestSupport.response("matrix-answer", SpanTestSupport.usage(4, 2, 6), "matrix-model"); + TracedChatModel traced = new TracedChatModel(SpanTestSupport.syncModel(p -> expected), tracer, "synthetic"); + ChatResponse actual = traced.call(new Prompt(List.of(new UserMessage("matrix input " + cell)))); + if (actual != expected) throw new AssertionError("call identity lost on " + cell); + + SpanData span = exporter.getFinishedSpanItems().get(exporter.getFinishedSpanItems().size() - 1); + eq(span.getAttributes().get(AttributeKey.stringKey(SemanticConventions.OUTPUT_VALUE)), "matrix-answer", "output"); + eq(span.getAttributes().get(AttributeKey.longKey(SemanticConventions.LLM_TOKEN_COUNT_PROMPT)), 4L, "prompt tokens"); + eq(span.getAttributes().get(AttributeKey.longKey(SemanticConventions.LLM_TOKEN_COUNT_COMPLETION)), 2L, "completion tokens"); + eq(span.getAttributes().get(AttributeKey.longKey(SemanticConventions.LLM_TOKEN_COUNT_TOTAL)), 6L, "total tokens"); + + traced.stream(new Prompt("stream " + cell)).blockLast(); + long streamSpans = exporter.getFinishedSpanItems().stream().filter(s -> s.getName().contains("Stream")).count(); + if (streamSpans != 1) throw new AssertionError("expected one stream span, got " + streamSpans); + + TracedEmbeddingModel tracedEmbeddings = new TracedEmbeddingModel(new SpanTestSupport.FixedEmbeddingModel(), tracer, "synthetic"); + List vectors = tracedEmbeddings.embed(List.of("a", "b")); + if (vectors.size() != 2 || vectors.get(0).length != 3) throw new AssertionError("batch identity lost"); + tracedEmbeddings.call(new EmbeddingRequest(List.of("a"), null)); + + CodeSource source = TracedChatModel.class.getProtectionDomain().getCodeSource(); + String digest = HexFormat.of().formatHex(MessageDigest.getInstance("SHA-256").digest(source.getLocation().openStream().readAllBytes())); + System.out.println("CELL=" + cell + " SHA256=" + digest + " SPANS=" + exporter.getFinishedSpanItems().size() + " COMMON_CONSUMER_PASS"); + provider.close(); + } + + private static void eq(Object actual, Object expected, String what) { + if (actual == null || !actual.equals(expected)) throw new AssertionError(what + " expected " + expected + " but was " + actual); + } +} diff --git a/java/compatibility/spring-ai/common/src/main/java/local/th7477/SpanTestSupport.java b/java/compatibility/spring-ai/common/src/main/java/local/th7477/SpanTestSupport.java new file mode 100644 index 00000000..af34a897 --- /dev/null +++ b/java/compatibility/spring-ai/common/src/main/java/local/th7477/SpanTestSupport.java @@ -0,0 +1,58 @@ +package local.th7477; + +import org.springframework.ai.chat.messages.AssistantMessage; +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.document.Document; +import org.springframework.ai.embedding.Embedding; +import org.springframework.ai.embedding.EmbeddingModel; +import org.springframework.ai.embedding.EmbeddingRequest; +import org.springframework.ai.embedding.EmbeddingResponse; +import org.springframework.ai.embedding.EmbeddingResponseMetadata; +import reactor.core.publisher.Flux; + +import java.util.List; +import java.util.function.Function; + +/** Spring AI objects shared by both matrix cells. Only APIs present in 1.1.0 and 1.1.8. */ +final class SpanTestSupport { + private SpanTestSupport() {} + + 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 Usage usage(int prompt, int completion, int total) { + 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 "matrix"; } + }; + } + + static ChatModel syncModel(Function call) { + return new ChatModel() { + @Override public ChatResponse call(Prompt prompt) { return call.apply(prompt); } + @Override public Flux stream(Prompt prompt) { return Flux.just(response("streamed-" + prompt.getContents(), null, null)); } + @Override public ChatOptions getDefaultOptions() { return null; } + }; + } + + static final class FixedEmbeddingModel implements EmbeddingModel { + private final EmbeddingResponse response = new EmbeddingResponse( + List.of(new Embedding(new float[]{1f, 2f, 3f}, 0), new Embedding(new float[]{4f, 5f, 6f}, 1)), + new EmbeddingResponseMetadata("m", null)); + @Override public EmbeddingResponse call(EmbeddingRequest request) { return response; } + @Override public float[] embed(Document document) { return new float[]{1f, 2f, 3f}; } + @Override public List embed(List texts) { return List.of(new float[]{1f, 2f, 3f}, new float[]{4f, 5f, 6f}); } + } +} diff --git a/java/compatibility/spring-ai/primary/pom.xml b/java/compatibility/spring-ai/primary/pom.xml new file mode 100644 index 00000000..f4eff421 --- /dev/null +++ b/java/compatibility/spring-ai/primary/pom.xml @@ -0,0 +1,26 @@ + + + 4.0.0 + local.th7477primary-consumer1.0-SNAPSHOT + + 17UTF-8 + 1.1.83.5.15 + + + + org.springframework.aispring-ai-bom1.1.8pomimport + org.springframework.bootspring-boot-dependencies${spring-boot.version}pomimport + + + + org.springframework.aispring-ai-model + org.springframework.aispring-ai-ollama + org.springframework.aispring-ai-client-chat + org.springframework.aispring-ai-client-chat + ai.traceaitraceai-java-core1.0.0-candidatesystem/candidate-jars/traceai-java-core-1.0.0.jar + ai.traceaitraceai-spring-ai1.0.0-candidatesystem/candidate-jars/traceai-spring-ai-1.0.0.jar + io.opentelemetryopentelemetry-sdk-testing1.49.0 + com.google.code.gsongson2.10.1 + + org.apache.maven.pluginsmaven-compiler-plugin3.12.1 + diff --git a/java/compatibility/spring-ai/primary/src/main/java/local/th7477/CommonConsumer.java b/java/compatibility/spring-ai/primary/src/main/java/local/th7477/CommonConsumer.java new file mode 100644 index 00000000..22058daf --- /dev/null +++ b/java/compatibility/spring-ai/primary/src/main/java/local/th7477/CommonConsumer.java @@ -0,0 +1,59 @@ +package local.th7477; + +import ai.traceai.FITracer; +import ai.traceai.SemanticConventions; +import ai.traceai.spring.TracedChatModel; +import ai.traceai.spring.TracedEmbeddingModel; +import io.opentelemetry.api.common.AttributeKey; +import io.opentelemetry.sdk.testing.exporter.InMemorySpanExporter; +import io.opentelemetry.sdk.trace.SdkTracerProvider; +import io.opentelemetry.sdk.trace.data.SpanData; +import io.opentelemetry.sdk.trace.export.SimpleSpanProcessor; +import org.springframework.ai.chat.messages.UserMessage; +import org.springframework.ai.chat.model.ChatResponse; +import org.springframework.ai.chat.prompt.Prompt; +import org.springframework.ai.embedding.EmbeddingRequest; + +import java.security.CodeSource; +import java.security.MessageDigest; +import java.util.HexFormat; +import java.util.List; + +/** Same compiled classes, same candidate JARs, run on both Spring AI cells. Asserts common API only. */ +public final class CommonConsumer { + public static void main(String[] args) throws Exception { + String cell = args.length > 0 ? args[0] : "unknown"; + InMemorySpanExporter exporter = InMemorySpanExporter.create(); + SdkTracerProvider provider = SdkTracerProvider.builder().addSpanProcessor(SimpleSpanProcessor.create(exporter)).build(); + FITracer tracer = new FITracer(provider.get("th7477-matrix")); + + ChatResponse expected = SpanTestSupport.response("matrix-answer", SpanTestSupport.usage(4, 2, 6), "matrix-model"); + TracedChatModel traced = new TracedChatModel(SpanTestSupport.syncModel(p -> expected), tracer, "synthetic"); + ChatResponse actual = traced.call(new Prompt(List.of(new UserMessage("matrix input " + cell)))); + if (actual != expected) throw new AssertionError("call identity lost on " + cell); + + SpanData span = exporter.getFinishedSpanItems().get(exporter.getFinishedSpanItems().size() - 1); + eq(span.getAttributes().get(AttributeKey.stringKey(SemanticConventions.OUTPUT_VALUE)), "matrix-answer", "output"); + eq(span.getAttributes().get(AttributeKey.longKey(SemanticConventions.LLM_TOKEN_COUNT_PROMPT)), 4L, "prompt tokens"); + eq(span.getAttributes().get(AttributeKey.longKey(SemanticConventions.LLM_TOKEN_COUNT_COMPLETION)), 2L, "completion tokens"); + eq(span.getAttributes().get(AttributeKey.longKey(SemanticConventions.LLM_TOKEN_COUNT_TOTAL)), 6L, "total tokens"); + + traced.stream(new Prompt("stream " + cell)).blockLast(); + long streamSpans = exporter.getFinishedSpanItems().stream().filter(s -> s.getName().contains("Stream")).count(); + if (streamSpans != 1) throw new AssertionError("expected one stream span, got " + streamSpans); + + TracedEmbeddingModel tracedEmbeddings = new TracedEmbeddingModel(new SpanTestSupport.FixedEmbeddingModel(), tracer, "synthetic"); + List vectors = tracedEmbeddings.embed(List.of("a", "b")); + if (vectors.size() != 2 || vectors.get(0).length != 3) throw new AssertionError("batch identity lost"); + tracedEmbeddings.call(new EmbeddingRequest(List.of("a"), null)); + + CodeSource source = TracedChatModel.class.getProtectionDomain().getCodeSource(); + String digest = HexFormat.of().formatHex(MessageDigest.getInstance("SHA-256").digest(source.getLocation().openStream().readAllBytes())); + System.out.println("CELL=" + cell + " SHA256=" + digest + " SPANS=" + exporter.getFinishedSpanItems().size() + " COMMON_CONSUMER_PASS"); + provider.close(); + } + + private static void eq(Object actual, Object expected, String what) { + if (actual == null || !actual.equals(expected)) throw new AssertionError(what + " expected " + expected + " but was " + actual); + } +} diff --git a/java/compatibility/spring-ai/primary/src/main/java/local/th7477/PrimaryConsumer.java b/java/compatibility/spring-ai/primary/src/main/java/local/th7477/PrimaryConsumer.java new file mode 100644 index 00000000..5536ad56 --- /dev/null +++ b/java/compatibility/spring-ai/primary/src/main/java/local/th7477/PrimaryConsumer.java @@ -0,0 +1,23 @@ +package local.th7477; + +import ai.traceai.FITracer; +import ai.traceai.spring.TracedEmbeddingModel; +import io.opentelemetry.sdk.testing.exporter.InMemorySpanExporter; +import io.opentelemetry.sdk.trace.SdkTracerProvider; +import io.opentelemetry.sdk.trace.export.SimpleSpanProcessor; +import org.springframework.ai.document.Document; + +/** P-only: getEmbeddingContent exists in Spring AI 1.1.8, not 1.1.0. Must not open a span. */ +public final class PrimaryConsumer { + public static void main(String[] args) { + InMemorySpanExporter exporter = InMemorySpanExporter.create(); + SdkTracerProvider provider = SdkTracerProvider.builder().addSpanProcessor(SimpleSpanProcessor.create(exporter)).build(); + FITracer tracer = new FITracer(provider.get("th7477-primary")); + TracedEmbeddingModel model = new TracedEmbeddingModel(new SpanTestSupport.FixedEmbeddingModel(), tracer, "synthetic"); + String content = model.getEmbeddingContent(new Document("doc-text")); + if (!"doc-text".equals(content)) throw new AssertionError("getEmbeddingContent returned " + content); + if (!exporter.getFinishedSpanItems().isEmpty()) throw new AssertionError("getEmbeddingContent opened a span"); + System.out.println("PRIMARY_ONLY_PASS getEmbeddingContent=" + content); + provider.close(); + } +} diff --git a/java/compatibility/spring-ai/primary/src/main/java/local/th7477/RealAgentCandidate.java b/java/compatibility/spring-ai/primary/src/main/java/local/th7477/RealAgentCandidate.java new file mode 100644 index 00000000..ea0a806f --- /dev/null +++ b/java/compatibility/spring-ai/primary/src/main/java/local/th7477/RealAgentCandidate.java @@ -0,0 +1,90 @@ +package local.th7477; + +import ai.traceai.FITracer; +import ai.traceai.SemanticConventions; +import ai.traceai.spring.TracedChatModel; +import io.opentelemetry.api.common.AttributeKey; +import io.opentelemetry.sdk.testing.exporter.InMemorySpanExporter; +import io.opentelemetry.sdk.trace.SdkTracerProvider; +import io.opentelemetry.sdk.trace.data.SpanData; +import io.opentelemetry.sdk.trace.export.SimpleSpanProcessor; +import org.springframework.ai.chat.client.ChatClient; +import org.springframework.ai.ollama.OllamaChatModel; +import org.springframework.ai.ollama.api.OllamaApi; +import org.springframework.ai.ollama.api.OllamaChatOptions; +import org.springframework.ai.tool.annotation.Tool; +import org.springframework.http.client.ClientHttpRequestInterceptor; +import org.springframework.web.client.RestClient; +import org.springframework.web.reactive.function.client.WebClient; + +import java.util.List; +import java.util.concurrent.atomic.AtomicInteger; + +/** + * After-fix real local agent: ChatClient + OllamaChatModel + a bounded harmless @Tool, through the repaired + * wrapper. Counts actual outbound /api/chat requests at the transport boundary and denies the fourth before + * any I/O, because the provider recurses privately for tool continuations. + */ +public final class RealAgentCandidate { + static final AtomicInteger REQUESTS = new AtomicInteger(); + static final int REQUEST_BUDGET = 3; + + public static final class LocalTools { + int calls; + @Tool(description = "Get the current status of the synthetic local service. Returns its exact status code. No network or side effects.") + public String getLocalStatus() { + if (++calls > 2) throw new IllegalStateException("Bounded local tool-call budget exceeded"); + System.out.println("ACTUAL_TOOL_EXECUTION getLocalStatus call=" + calls); + return "LOCAL_OK_7477"; + } + } + + public static void main(String[] args) { + String mode = args.length > 0 ? args[0] : "sync"; + String url = System.getenv().getOrDefault("LOCAL_OLLAMA_URL", "http://host.docker.internal:17477"); + ClientHttpRequestInterceptor budget = (request, body, execution) -> { + if (REQUESTS.incrementAndGet() > REQUEST_BUDGET) { + throw new IllegalStateException("provider request budget exceeded before I/O"); + } + return execution.execute(request, body); + }; + OllamaApi api = OllamaApi.builder().baseUrl(url) + .restClientBuilder(RestClient.builder().requestInterceptor(budget)) + .webClientBuilder(WebClient.builder().filter((request, next) -> { + if (REQUESTS.incrementAndGet() > REQUEST_BUDGET) { + return reactor.core.publisher.Mono.error(new IllegalStateException("provider request budget exceeded before I/O")); + } + return next.exchange(request); + })) + .build(); + OllamaChatModel model = OllamaChatModel.builder().ollamaApi(api) + .defaultOptions(OllamaChatOptions.builder().model("qwen3:1.7b").temperature(0.0).numCtx(4096).numPredict(128) + .disableThinking().keepAlive("2m").build()) + .build(); + + InMemorySpanExporter exporter = InMemorySpanExporter.create(); + SdkTracerProvider provider = SdkTracerProvider.builder().addSpanProcessor(SimpleSpanProcessor.create(exporter)).build(); + FITracer tracer = new FITracer(provider.get("th7477-real-agent")); + TracedChatModel traced = new TracedChatModel(model, tracer, "ollama"); + + LocalTools tools = new LocalTools(); + String prompt = "Call getLocalStatus once to check the local service. Include the exact returned status code in your short answer. Do not guess."; + ChatClient client = ChatClient.create(traced); + String answer = "sync".equals(mode) + ? client.prompt(prompt).tools(tools).call().content() + : client.prompt(prompt).tools(tools).stream().content().collectList().block().stream().reduce("", String::concat); + System.out.println("REAL_MODEL_ANSWER=" + answer); + + if (tools.calls < 1 || answer == null || !answer.contains("LOCAL_OK_7477")) { + throw new AssertionError("Real model tool journey failed; calls=" + tools.calls + ", answer=" + answer); + } + List spans = exporter.getFinishedSpanItems(); + boolean sawOutput = spans.stream().anyMatch(s -> "LOCAL_OK_7477".equals( + s.getAttributes().get(AttributeKey.stringKey(SemanticConventions.OUTPUT_VALUE))) + || String.valueOf(s.getAttributes().get(AttributeKey.stringKey(SemanticConventions.OUTPUT_VALUE))).contains("LOCAL_OK_7477")); + System.out.println("MODE=" + mode + " TOOL_CALLS=" + tools.calls + " PROVIDER_REQUESTS=" + REQUESTS.get() + + " SPANS=" + spans.size() + " OUTPUT_CAPTURED=" + sawOutput + " REAL_AGENT_PASS"); + if (spans.isEmpty()) throw new AssertionError("no exported spans"); + provider.close(); + } +} diff --git a/java/compatibility/spring-ai/primary/src/main/java/local/th7477/SpanTestSupport.java b/java/compatibility/spring-ai/primary/src/main/java/local/th7477/SpanTestSupport.java new file mode 100644 index 00000000..af34a897 --- /dev/null +++ b/java/compatibility/spring-ai/primary/src/main/java/local/th7477/SpanTestSupport.java @@ -0,0 +1,58 @@ +package local.th7477; + +import org.springframework.ai.chat.messages.AssistantMessage; +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.document.Document; +import org.springframework.ai.embedding.Embedding; +import org.springframework.ai.embedding.EmbeddingModel; +import org.springframework.ai.embedding.EmbeddingRequest; +import org.springframework.ai.embedding.EmbeddingResponse; +import org.springframework.ai.embedding.EmbeddingResponseMetadata; +import reactor.core.publisher.Flux; + +import java.util.List; +import java.util.function.Function; + +/** Spring AI objects shared by both matrix cells. Only APIs present in 1.1.0 and 1.1.8. */ +final class SpanTestSupport { + private SpanTestSupport() {} + + 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 Usage usage(int prompt, int completion, int total) { + 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 "matrix"; } + }; + } + + static ChatModel syncModel(Function call) { + return new ChatModel() { + @Override public ChatResponse call(Prompt prompt) { return call.apply(prompt); } + @Override public Flux stream(Prompt prompt) { return Flux.just(response("streamed-" + prompt.getContents(), null, null)); } + @Override public ChatOptions getDefaultOptions() { return null; } + }; + } + + static final class FixedEmbeddingModel implements EmbeddingModel { + private final EmbeddingResponse response = new EmbeddingResponse( + List.of(new Embedding(new float[]{1f, 2f, 3f}, 0), new Embedding(new float[]{4f, 5f, 6f}, 1)), + new EmbeddingResponseMetadata("m", null)); + @Override public EmbeddingResponse call(EmbeddingRequest request) { return response; } + @Override public float[] embed(Document document) { return new float[]{1f, 2f, 3f}; } + @Override public List embed(List texts) { return List.of(new float[]{1f, 2f, 3f}, new float[]{4f, 5f, 6f}); } + } +} diff --git a/java/compatibility/spring-ai/primary/src/src/main/java/local/th7477/CommonConsumer.java b/java/compatibility/spring-ai/primary/src/src/main/java/local/th7477/CommonConsumer.java new file mode 100644 index 00000000..22058daf --- /dev/null +++ b/java/compatibility/spring-ai/primary/src/src/main/java/local/th7477/CommonConsumer.java @@ -0,0 +1,59 @@ +package local.th7477; + +import ai.traceai.FITracer; +import ai.traceai.SemanticConventions; +import ai.traceai.spring.TracedChatModel; +import ai.traceai.spring.TracedEmbeddingModel; +import io.opentelemetry.api.common.AttributeKey; +import io.opentelemetry.sdk.testing.exporter.InMemorySpanExporter; +import io.opentelemetry.sdk.trace.SdkTracerProvider; +import io.opentelemetry.sdk.trace.data.SpanData; +import io.opentelemetry.sdk.trace.export.SimpleSpanProcessor; +import org.springframework.ai.chat.messages.UserMessage; +import org.springframework.ai.chat.model.ChatResponse; +import org.springframework.ai.chat.prompt.Prompt; +import org.springframework.ai.embedding.EmbeddingRequest; + +import java.security.CodeSource; +import java.security.MessageDigest; +import java.util.HexFormat; +import java.util.List; + +/** Same compiled classes, same candidate JARs, run on both Spring AI cells. Asserts common API only. */ +public final class CommonConsumer { + public static void main(String[] args) throws Exception { + String cell = args.length > 0 ? args[0] : "unknown"; + InMemorySpanExporter exporter = InMemorySpanExporter.create(); + SdkTracerProvider provider = SdkTracerProvider.builder().addSpanProcessor(SimpleSpanProcessor.create(exporter)).build(); + FITracer tracer = new FITracer(provider.get("th7477-matrix")); + + ChatResponse expected = SpanTestSupport.response("matrix-answer", SpanTestSupport.usage(4, 2, 6), "matrix-model"); + TracedChatModel traced = new TracedChatModel(SpanTestSupport.syncModel(p -> expected), tracer, "synthetic"); + ChatResponse actual = traced.call(new Prompt(List.of(new UserMessage("matrix input " + cell)))); + if (actual != expected) throw new AssertionError("call identity lost on " + cell); + + SpanData span = exporter.getFinishedSpanItems().get(exporter.getFinishedSpanItems().size() - 1); + eq(span.getAttributes().get(AttributeKey.stringKey(SemanticConventions.OUTPUT_VALUE)), "matrix-answer", "output"); + eq(span.getAttributes().get(AttributeKey.longKey(SemanticConventions.LLM_TOKEN_COUNT_PROMPT)), 4L, "prompt tokens"); + eq(span.getAttributes().get(AttributeKey.longKey(SemanticConventions.LLM_TOKEN_COUNT_COMPLETION)), 2L, "completion tokens"); + eq(span.getAttributes().get(AttributeKey.longKey(SemanticConventions.LLM_TOKEN_COUNT_TOTAL)), 6L, "total tokens"); + + traced.stream(new Prompt("stream " + cell)).blockLast(); + long streamSpans = exporter.getFinishedSpanItems().stream().filter(s -> s.getName().contains("Stream")).count(); + if (streamSpans != 1) throw new AssertionError("expected one stream span, got " + streamSpans); + + TracedEmbeddingModel tracedEmbeddings = new TracedEmbeddingModel(new SpanTestSupport.FixedEmbeddingModel(), tracer, "synthetic"); + List vectors = tracedEmbeddings.embed(List.of("a", "b")); + if (vectors.size() != 2 || vectors.get(0).length != 3) throw new AssertionError("batch identity lost"); + tracedEmbeddings.call(new EmbeddingRequest(List.of("a"), null)); + + CodeSource source = TracedChatModel.class.getProtectionDomain().getCodeSource(); + String digest = HexFormat.of().formatHex(MessageDigest.getInstance("SHA-256").digest(source.getLocation().openStream().readAllBytes())); + System.out.println("CELL=" + cell + " SHA256=" + digest + " SPANS=" + exporter.getFinishedSpanItems().size() + " COMMON_CONSUMER_PASS"); + provider.close(); + } + + private static void eq(Object actual, Object expected, String what) { + if (actual == null || !actual.equals(expected)) throw new AssertionError(what + " expected " + expected + " but was " + actual); + } +} diff --git a/java/compatibility/spring-ai/primary/src/src/main/java/local/th7477/SpanTestSupport.java b/java/compatibility/spring-ai/primary/src/src/main/java/local/th7477/SpanTestSupport.java new file mode 100644 index 00000000..af34a897 --- /dev/null +++ b/java/compatibility/spring-ai/primary/src/src/main/java/local/th7477/SpanTestSupport.java @@ -0,0 +1,58 @@ +package local.th7477; + +import org.springframework.ai.chat.messages.AssistantMessage; +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.document.Document; +import org.springframework.ai.embedding.Embedding; +import org.springframework.ai.embedding.EmbeddingModel; +import org.springframework.ai.embedding.EmbeddingRequest; +import org.springframework.ai.embedding.EmbeddingResponse; +import org.springframework.ai.embedding.EmbeddingResponseMetadata; +import reactor.core.publisher.Flux; + +import java.util.List; +import java.util.function.Function; + +/** Spring AI objects shared by both matrix cells. Only APIs present in 1.1.0 and 1.1.8. */ +final class SpanTestSupport { + private SpanTestSupport() {} + + 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 Usage usage(int prompt, int completion, int total) { + 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 "matrix"; } + }; + } + + static ChatModel syncModel(Function call) { + return new ChatModel() { + @Override public ChatResponse call(Prompt prompt) { return call.apply(prompt); } + @Override public Flux stream(Prompt prompt) { return Flux.just(response("streamed-" + prompt.getContents(), null, null)); } + @Override public ChatOptions getDefaultOptions() { return null; } + }; + } + + static final class FixedEmbeddingModel implements EmbeddingModel { + private final EmbeddingResponse response = new EmbeddingResponse( + List.of(new Embedding(new float[]{1f, 2f, 3f}, 0), new Embedding(new float[]{4f, 5f, 6f}, 1)), + new EmbeddingResponseMetadata("m", null)); + @Override public EmbeddingResponse call(EmbeddingRequest request) { return response; } + @Override public float[] embed(Document document) { return new float[]{1f, 2f, 3f}; } + @Override public List embed(List texts) { return List.of(new float[]{1f, 2f, 3f}, new float[]{4f, 5f, 6f}); } + } +} diff --git a/java/examples/spring-ai-example/pom.xml b/java/examples/spring-ai-example/pom.xml index 0a71151e..adf3a48e 100644 --- a/java/examples/spring-ai-example/pom.xml +++ b/java/examples/spring-ai-example/pom.xml @@ -7,7 +7,7 @@ org.springframework.boot spring-boot-starter-parent - 3.2.1 + 3.5.15 ai.traceai.examples @@ -17,42 +17,37 @@ 17 - 1.0.0-M4 + 1.1.8 + + + + org.springframework.ai + spring-ai-bom + ${spring-ai.version} + pom + import + + + + - org.springframework.boot spring-boot-starter-web - - ai.traceai traceai-spring-boot-starter 1.0.0 - - org.springframework.ai - spring-ai-openai-spring-boot-starter - ${spring-ai.version} + spring-ai-starter-model-openai - - - spring-milestones - Spring Milestones - https://repo.spring.io/milestone - - false - - - - diff --git a/java/examples/spring-ai-example/src/main/java/ai/traceai/examples/SpringAIExampleApplication.java b/java/examples/spring-ai-example/src/main/java/ai/traceai/examples/SpringAIExampleApplication.java index dfc1dd17..6db71a73 100644 --- a/java/examples/spring-ai-example/src/main/java/ai/traceai/examples/SpringAIExampleApplication.java +++ b/java/examples/spring-ai-example/src/main/java/ai/traceai/examples/SpringAIExampleApplication.java @@ -63,12 +63,12 @@ public CommandLineRunner demo(TracedChatModel tracedChatModel) { new Prompt("What is 2 + 2? Answer in one word.") ); - System.out.println("Response: " + response.getResult().getOutput().getContent()); + System.out.println("Response: " + response.getResult().getOutput().getText()); if (response.getMetadata() != null && response.getMetadata().getUsage() != null) { var usage = response.getMetadata().getUsage(); System.out.println("Tokens - Prompt: " + usage.getPromptTokens() - + ", Generation: " + usage.getGenerationTokens()); + + ", Completion: " + usage.getCompletionTokens()); } System.out.println("\nTraces are being exported to TraceAI."); @@ -94,13 +94,13 @@ public ChatController(TracedChatModel chatModel) { @GetMapping public String chat(@RequestParam String message) { var response = chatModel.call(new Prompt(message)); - return response.getResult().getOutput().getContent(); + return response.getResult().getOutput().getText(); } @PostMapping public String chatPost(@RequestBody ChatRequest request) { var response = chatModel.call(new Prompt(request.message())); - return response.getResult().getOutput().getContent(); + return response.getResult().getOutput().getText(); } record ChatRequest(String message) {} diff --git a/java/pom.xml b/java/pom.xml index eae1236b..cf36d9d4 100644 --- a/java/pom.xml +++ b/java/pom.xml @@ -80,8 +80,8 @@ 1.0.0-beta.16 0.1.0 0.35.0 - 1.0.0-M4 - 3.2.1 + 1.1.8 + 3.5.15 2.25.67 1.7.0 diff --git a/java/traceai-spring-ai/pom.xml b/java/traceai-spring-ai/pom.xml index 624d187b..4052ea65 100644 --- a/java/traceai-spring-ai/pom.xml +++ b/java/traceai-spring-ai/pom.xml @@ -12,7 +12,40 @@ traceai-spring-ai TraceAI Spring AI - Spring AI instrumentation for TraceAI Java SDK + Spring AI instrumentation for TraceAI Java SDK (Spring AI 1.1.x GA) + + + + 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 + + + @@ -21,19 +54,17 @@ traceai-java-core - + org.springframework.ai - spring-ai-core - ${spring-ai.version} + spring-ai-model provided - + org.springframework spring-context - 6.1.3 provided @@ -49,19 +80,30 @@ test - org.mockito - mockito-core - test - - - org.mockito - mockito-junit-jupiter + io.opentelemetry + opentelemetry-sdk-testing test - io.opentelemetry - opentelemetry-sdk-testing + io.projectreactor + reactor-test test + + + + + org.apache.maven.plugins + maven-surefire-plugin + + + + --add-opens java.base/java.lang=ALL-UNNAMED + --add-opens java.base/java.util=ALL-UNNAMED + + + + + diff --git a/java/traceai-spring-ai/src/main/java/ai/traceai/spring/SpringAITracingContext.java b/java/traceai-spring-ai/src/main/java/ai/traceai/spring/SpringAITracingContext.java new file mode 100644 index 00000000..d7aac4fa --- /dev/null +++ b/java/traceai-spring-ai/src/main/java/ai/traceai/spring/SpringAITracingContext.java @@ -0,0 +1,181 @@ +package ai.traceai.spring; + +import ai.traceai.ContextAttributes; +import ai.traceai.FISpanKind; +import ai.traceai.FITracer; +import ai.traceai.SemanticConventions; +import io.opentelemetry.api.trace.Span; +import io.opentelemetry.context.Context; +import reactor.util.context.ContextView; + +import java.util.Collections; +import java.util.LinkedHashMap; +import java.util.Map; +import java.util.Objects; +import java.util.Set; +import java.util.function.UnaryOperator; + +/** + * Explicit operation-context supply for the Spring AI reactive wrappers. + * + *

TraceAI session/user/metadata/tag attributes live in {@link ContextAttributes} thread-locals and the + * OpenTelemetry parent span lives in the current {@link Context}. Both are thread-bound, which is unreliable + * for a {@code Flux} that is assembled on one thread, subscribed on another and signalled from a provider + * I/O thread. Use this helper to bind the parent context and the TraceAI attributes to the Reactor + * subscriber context instead:

+ * + * 
+ * tracedModel.stream(prompt)
+ *     .contextWrite(SpringAITracingContext.withContext(Context.current(), ContextAttributes.getAttributesFromContext()))
+ *
+ * + *

Selection order for every span the wrappers create:

+ *
    + *
  1. an explicit snapshot supplied with {@link #withContext(Context, Map)} (parent and attributes override + * everything else as one unit; {@code Context.root()} plus an empty map is an explicit empty context, + * not absence);
    2. + *
  2. otherwise the subscription-time current OpenTelemetry span (when valid) together with the current + * {@link ContextAttributes} snapshot;
    3. + *
  3. otherwise, when only subscription-time attributes are present, those attributes with the current + * (possibly root) parent;
    4. + *
  4. otherwise the snapshot captured when {@code stream(prompt)} was assembled.
    5. + *
+ * + *

Reusing one publisher across tenants therefore requires an explicit snapshot per subscription.

+ */ +public final class SpringAITracingContext { + + /** Private namespaced Reactor context key. Callers use the helper, never this key. */ + static final String CONTEXT_KEY = "ai.traceai.spring.operation-context.v1"; + + private static final Set ALLOWED_KEYS = Set.of( + SemanticConventions.SESSION_ID, + SemanticConventions.GEN_AI_CONVERSATION_ID, + SemanticConventions.USER_ID, + SemanticConventions.METADATA, + SemanticConventions.TAG_TAGS); + + private SpringAITracingContext() { + throw new UnsupportedOperationException("Utility class"); + } + + /** + * Builds a Reactor context operator that binds the given OpenTelemetry parent and TraceAI attributes to + * every span the wrapped model creates for that subscription. + * + * @param parent the OpenTelemetry parent context; use {@link Context#root()} for no parent + * @param attributes TraceAI attribute strings as produced by {@link ContextAttributes#getAttributesFromContext()} + * (keys {@code session.id}, {@code gen_ai.conversation.id}, {@code user.id}, {@code metadata}, + * {@code tag.tags}); the map is copied, metadata/tags stay verbatim serialized strings + * @return an operator for {@code Flux.contextWrite(...)} + * @throws NullPointerException if parent or attributes is null + * @throws IllegalArgumentException for unknown keys or null values + */ + public static UnaryOperator withContext(Context parent, Map attributes) { + Objects.requireNonNull(parent, "parent context must not be null"); + Objects.requireNonNull(attributes, "attributes must not be null"); + Map copy = new LinkedHashMap<>(); + for (Map.Entry entry : attributes.entrySet()) { + String key = entry.getKey(); + if (key == null || !ALLOWED_KEYS.contains(key)) { + throw new IllegalArgumentException("Unsupported TraceAI context attribute key: " + key); + } + if (entry.getValue() == null) { + throw new IllegalArgumentException("TraceAI context attribute value must not be null: " + key); + } + copy.put(key, entry.getValue()); + } + Snapshot snapshot = new Snapshot(parent, Collections.unmodifiableMap(copy)); + return reactorContext -> reactorContext.put(CONTEXT_KEY, snapshot); + } + + /** Captures the current thread's parent context and TraceAI attributes. */ + static Snapshot capture() { + return new Snapshot(Context.current(), ContextAttributes.getAttributesFromContext()); + } + + /** + * Selects the operation context for one subscription. + * + * @param subscriberContext the downstream Reactor context at subscription time + * @param assembly the snapshot captured at assembly time (fallback) + */ + static Snapshot select(ContextView subscriberContext, Snapshot assembly) { + Snapshot explicit = subscriberContext.getOrDefault(CONTEXT_KEY, null); + if (explicit != null) { + return explicit; + } + Context current = Context.current(); + Map attributes = ContextAttributes.getAttributesFromContext(); + Span currentSpan = Span.fromContextOrNull(current); + if (currentSpan != null && currentSpan.getSpanContext().isValid()) { + return new Snapshot(current, attributes); + } + if (!attributes.isEmpty()) { + return new Snapshot(current, attributes); + } + return assembly; + } + + /** + * Starts a span from the selected snapshot: thread-local attributes are cleared while the span is created + * so stale worker-thread values never leak in, then the snapshot attributes are applied verbatim. + */ + static Span startSpan(FITracer tracer, String name, FISpanKind kind, Snapshot snapshot) { + AutoCloseable session = ContextAttributes.usingSession(null); + try { + AutoCloseable user = ContextAttributes.usingUser(null); + try { + AutoCloseable metadata = ContextAttributes.usingMetadata(null); + try { + AutoCloseable tags = ContextAttributes.usingTags(null); + try { + Span span = tracer.startSpan(name, kind, snapshot.parent()); + for (Map.Entry entry : snapshot.attributes().entrySet()) { + span.setAttribute(entry.getKey(), entry.getValue()); + } + return span; + } finally { + closeQuietly(tags); + } + } finally { + closeQuietly(metadata); + } + } finally { + closeQuietly(user); + } + } finally { + closeQuietly(session); + } + } + + private static void closeQuietly(AutoCloseable scope) { + try { + scope.close(); + } catch (RuntimeException e) { + throw e; + } catch (Exception e) { + // ContextAttributes scopes only restore thread-locals; a checked exception cannot occur here. + throw new IllegalStateException("Failed to restore TraceAI context attributes", e); + } + } + + /** Immutable selected operation context. */ + static final class Snapshot { + private final Context parent; + private final Map attributes; + + Snapshot(Context parent, Map attributes) { + this.parent = Objects.requireNonNull(parent, "parent"); + this.attributes = Collections.unmodifiableMap(new LinkedHashMap<>(Objects.requireNonNull(attributes, "attributes"))); + } + + Context parent() { + return parent; + } + + Map 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..bcac94af 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,127 @@ 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"); + } + preview.append(texts.get(i)); + } + 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..b62ed411 --- /dev/null +++ b/java/traceai-spring-ai/src/main/java/ai/traceai/spring/TracingSubscriber.java @@ -0,0 +1,178 @@ +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 (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); + } 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..faeabcc6 --- /dev/null +++ b/java/traceai-spring-ai/src/test/java/ai/traceai/spring/TracedChatModelTest.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.SystemMessage; +import org.springframework.ai.chat.messages.UserMessage; +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.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 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 + + From c4c544da2a71448c6cd0c247f135f9067961df1f Mon Sep 17 00:00:00 2001 From: Nikhil Pareek Date: Sat, 3 Oct 2026 20:45:22 -0700 Subject: [PATCH 2/7] chore(spring-ai): drop duplicated compatibility sources --- .../java/local/th7477/CommonConsumer.java | 59 ------------------- .../java/local/th7477/SpanTestSupport.java | 58 ------------------ .../java/local/th7477/CommonConsumer.java | 59 ------------------- .../java/local/th7477/SpanTestSupport.java | 58 ------------------ 4 files changed, 234 deletions(-) delete mode 100644 java/compatibility/spring-ai/baseline-cell/src/src/main/java/local/th7477/CommonConsumer.java delete mode 100644 java/compatibility/spring-ai/baseline-cell/src/src/main/java/local/th7477/SpanTestSupport.java delete mode 100644 java/compatibility/spring-ai/primary/src/src/main/java/local/th7477/CommonConsumer.java delete mode 100644 java/compatibility/spring-ai/primary/src/src/main/java/local/th7477/SpanTestSupport.java diff --git a/java/compatibility/spring-ai/baseline-cell/src/src/main/java/local/th7477/CommonConsumer.java b/java/compatibility/spring-ai/baseline-cell/src/src/main/java/local/th7477/CommonConsumer.java deleted file mode 100644 index 22058daf..00000000 --- a/java/compatibility/spring-ai/baseline-cell/src/src/main/java/local/th7477/CommonConsumer.java +++ /dev/null @@ -1,59 +0,0 @@ -package local.th7477; - -import ai.traceai.FITracer; -import ai.traceai.SemanticConventions; -import ai.traceai.spring.TracedChatModel; -import ai.traceai.spring.TracedEmbeddingModel; -import io.opentelemetry.api.common.AttributeKey; -import io.opentelemetry.sdk.testing.exporter.InMemorySpanExporter; -import io.opentelemetry.sdk.trace.SdkTracerProvider; -import io.opentelemetry.sdk.trace.data.SpanData; -import io.opentelemetry.sdk.trace.export.SimpleSpanProcessor; -import org.springframework.ai.chat.messages.UserMessage; -import org.springframework.ai.chat.model.ChatResponse; -import org.springframework.ai.chat.prompt.Prompt; -import org.springframework.ai.embedding.EmbeddingRequest; - -import java.security.CodeSource; -import java.security.MessageDigest; -import java.util.HexFormat; -import java.util.List; - -/** Same compiled classes, same candidate JARs, run on both Spring AI cells. Asserts common API only. */ -public final class CommonConsumer { - public static void main(String[] args) throws Exception { - String cell = args.length > 0 ? args[0] : "unknown"; - InMemorySpanExporter exporter = InMemorySpanExporter.create(); - SdkTracerProvider provider = SdkTracerProvider.builder().addSpanProcessor(SimpleSpanProcessor.create(exporter)).build(); - FITracer tracer = new FITracer(provider.get("th7477-matrix")); - - ChatResponse expected = SpanTestSupport.response("matrix-answer", SpanTestSupport.usage(4, 2, 6), "matrix-model"); - TracedChatModel traced = new TracedChatModel(SpanTestSupport.syncModel(p -> expected), tracer, "synthetic"); - ChatResponse actual = traced.call(new Prompt(List.of(new UserMessage("matrix input " + cell)))); - if (actual != expected) throw new AssertionError("call identity lost on " + cell); - - SpanData span = exporter.getFinishedSpanItems().get(exporter.getFinishedSpanItems().size() - 1); - eq(span.getAttributes().get(AttributeKey.stringKey(SemanticConventions.OUTPUT_VALUE)), "matrix-answer", "output"); - eq(span.getAttributes().get(AttributeKey.longKey(SemanticConventions.LLM_TOKEN_COUNT_PROMPT)), 4L, "prompt tokens"); - eq(span.getAttributes().get(AttributeKey.longKey(SemanticConventions.LLM_TOKEN_COUNT_COMPLETION)), 2L, "completion tokens"); - eq(span.getAttributes().get(AttributeKey.longKey(SemanticConventions.LLM_TOKEN_COUNT_TOTAL)), 6L, "total tokens"); - - traced.stream(new Prompt("stream " + cell)).blockLast(); - long streamSpans = exporter.getFinishedSpanItems().stream().filter(s -> s.getName().contains("Stream")).count(); - if (streamSpans != 1) throw new AssertionError("expected one stream span, got " + streamSpans); - - TracedEmbeddingModel tracedEmbeddings = new TracedEmbeddingModel(new SpanTestSupport.FixedEmbeddingModel(), tracer, "synthetic"); - List vectors = tracedEmbeddings.embed(List.of("a", "b")); - if (vectors.size() != 2 || vectors.get(0).length != 3) throw new AssertionError("batch identity lost"); - tracedEmbeddings.call(new EmbeddingRequest(List.of("a"), null)); - - CodeSource source = TracedChatModel.class.getProtectionDomain().getCodeSource(); - String digest = HexFormat.of().formatHex(MessageDigest.getInstance("SHA-256").digest(source.getLocation().openStream().readAllBytes())); - System.out.println("CELL=" + cell + " SHA256=" + digest + " SPANS=" + exporter.getFinishedSpanItems().size() + " COMMON_CONSUMER_PASS"); - provider.close(); - } - - private static void eq(Object actual, Object expected, String what) { - if (actual == null || !actual.equals(expected)) throw new AssertionError(what + " expected " + expected + " but was " + actual); - } -} diff --git a/java/compatibility/spring-ai/baseline-cell/src/src/main/java/local/th7477/SpanTestSupport.java b/java/compatibility/spring-ai/baseline-cell/src/src/main/java/local/th7477/SpanTestSupport.java deleted file mode 100644 index af34a897..00000000 --- a/java/compatibility/spring-ai/baseline-cell/src/src/main/java/local/th7477/SpanTestSupport.java +++ /dev/null @@ -1,58 +0,0 @@ -package local.th7477; - -import org.springframework.ai.chat.messages.AssistantMessage; -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.document.Document; -import org.springframework.ai.embedding.Embedding; -import org.springframework.ai.embedding.EmbeddingModel; -import org.springframework.ai.embedding.EmbeddingRequest; -import org.springframework.ai.embedding.EmbeddingResponse; -import org.springframework.ai.embedding.EmbeddingResponseMetadata; -import reactor.core.publisher.Flux; - -import java.util.List; -import java.util.function.Function; - -/** Spring AI objects shared by both matrix cells. Only APIs present in 1.1.0 and 1.1.8. */ -final class SpanTestSupport { - private SpanTestSupport() {} - - 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 Usage usage(int prompt, int completion, int total) { - 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 "matrix"; } - }; - } - - static ChatModel syncModel(Function call) { - return new ChatModel() { - @Override public ChatResponse call(Prompt prompt) { return call.apply(prompt); } - @Override public Flux stream(Prompt prompt) { return Flux.just(response("streamed-" + prompt.getContents(), null, null)); } - @Override public ChatOptions getDefaultOptions() { return null; } - }; - } - - static final class FixedEmbeddingModel implements EmbeddingModel { - private final EmbeddingResponse response = new EmbeddingResponse( - List.of(new Embedding(new float[]{1f, 2f, 3f}, 0), new Embedding(new float[]{4f, 5f, 6f}, 1)), - new EmbeddingResponseMetadata("m", null)); - @Override public EmbeddingResponse call(EmbeddingRequest request) { return response; } - @Override public float[] embed(Document document) { return new float[]{1f, 2f, 3f}; } - @Override public List embed(List texts) { return List.of(new float[]{1f, 2f, 3f}, new float[]{4f, 5f, 6f}); } - } -} diff --git a/java/compatibility/spring-ai/primary/src/src/main/java/local/th7477/CommonConsumer.java b/java/compatibility/spring-ai/primary/src/src/main/java/local/th7477/CommonConsumer.java deleted file mode 100644 index 22058daf..00000000 --- a/java/compatibility/spring-ai/primary/src/src/main/java/local/th7477/CommonConsumer.java +++ /dev/null @@ -1,59 +0,0 @@ -package local.th7477; - -import ai.traceai.FITracer; -import ai.traceai.SemanticConventions; -import ai.traceai.spring.TracedChatModel; -import ai.traceai.spring.TracedEmbeddingModel; -import io.opentelemetry.api.common.AttributeKey; -import io.opentelemetry.sdk.testing.exporter.InMemorySpanExporter; -import io.opentelemetry.sdk.trace.SdkTracerProvider; -import io.opentelemetry.sdk.trace.data.SpanData; -import io.opentelemetry.sdk.trace.export.SimpleSpanProcessor; -import org.springframework.ai.chat.messages.UserMessage; -import org.springframework.ai.chat.model.ChatResponse; -import org.springframework.ai.chat.prompt.Prompt; -import org.springframework.ai.embedding.EmbeddingRequest; - -import java.security.CodeSource; -import java.security.MessageDigest; -import java.util.HexFormat; -import java.util.List; - -/** Same compiled classes, same candidate JARs, run on both Spring AI cells. Asserts common API only. */ -public final class CommonConsumer { - public static void main(String[] args) throws Exception { - String cell = args.length > 0 ? args[0] : "unknown"; - InMemorySpanExporter exporter = InMemorySpanExporter.create(); - SdkTracerProvider provider = SdkTracerProvider.builder().addSpanProcessor(SimpleSpanProcessor.create(exporter)).build(); - FITracer tracer = new FITracer(provider.get("th7477-matrix")); - - ChatResponse expected = SpanTestSupport.response("matrix-answer", SpanTestSupport.usage(4, 2, 6), "matrix-model"); - TracedChatModel traced = new TracedChatModel(SpanTestSupport.syncModel(p -> expected), tracer, "synthetic"); - ChatResponse actual = traced.call(new Prompt(List.of(new UserMessage("matrix input " + cell)))); - if (actual != expected) throw new AssertionError("call identity lost on " + cell); - - SpanData span = exporter.getFinishedSpanItems().get(exporter.getFinishedSpanItems().size() - 1); - eq(span.getAttributes().get(AttributeKey.stringKey(SemanticConventions.OUTPUT_VALUE)), "matrix-answer", "output"); - eq(span.getAttributes().get(AttributeKey.longKey(SemanticConventions.LLM_TOKEN_COUNT_PROMPT)), 4L, "prompt tokens"); - eq(span.getAttributes().get(AttributeKey.longKey(SemanticConventions.LLM_TOKEN_COUNT_COMPLETION)), 2L, "completion tokens"); - eq(span.getAttributes().get(AttributeKey.longKey(SemanticConventions.LLM_TOKEN_COUNT_TOTAL)), 6L, "total tokens"); - - traced.stream(new Prompt("stream " + cell)).blockLast(); - long streamSpans = exporter.getFinishedSpanItems().stream().filter(s -> s.getName().contains("Stream")).count(); - if (streamSpans != 1) throw new AssertionError("expected one stream span, got " + streamSpans); - - TracedEmbeddingModel tracedEmbeddings = new TracedEmbeddingModel(new SpanTestSupport.FixedEmbeddingModel(), tracer, "synthetic"); - List vectors = tracedEmbeddings.embed(List.of("a", "b")); - if (vectors.size() != 2 || vectors.get(0).length != 3) throw new AssertionError("batch identity lost"); - tracedEmbeddings.call(new EmbeddingRequest(List.of("a"), null)); - - CodeSource source = TracedChatModel.class.getProtectionDomain().getCodeSource(); - String digest = HexFormat.of().formatHex(MessageDigest.getInstance("SHA-256").digest(source.getLocation().openStream().readAllBytes())); - System.out.println("CELL=" + cell + " SHA256=" + digest + " SPANS=" + exporter.getFinishedSpanItems().size() + " COMMON_CONSUMER_PASS"); - provider.close(); - } - - private static void eq(Object actual, Object expected, String what) { - if (actual == null || !actual.equals(expected)) throw new AssertionError(what + " expected " + expected + " but was " + actual); - } -} diff --git a/java/compatibility/spring-ai/primary/src/src/main/java/local/th7477/SpanTestSupport.java b/java/compatibility/spring-ai/primary/src/src/main/java/local/th7477/SpanTestSupport.java deleted file mode 100644 index af34a897..00000000 --- a/java/compatibility/spring-ai/primary/src/src/main/java/local/th7477/SpanTestSupport.java +++ /dev/null @@ -1,58 +0,0 @@ -package local.th7477; - -import org.springframework.ai.chat.messages.AssistantMessage; -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.document.Document; -import org.springframework.ai.embedding.Embedding; -import org.springframework.ai.embedding.EmbeddingModel; -import org.springframework.ai.embedding.EmbeddingRequest; -import org.springframework.ai.embedding.EmbeddingResponse; -import org.springframework.ai.embedding.EmbeddingResponseMetadata; -import reactor.core.publisher.Flux; - -import java.util.List; -import java.util.function.Function; - -/** Spring AI objects shared by both matrix cells. Only APIs present in 1.1.0 and 1.1.8. */ -final class SpanTestSupport { - private SpanTestSupport() {} - - 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 Usage usage(int prompt, int completion, int total) { - 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 "matrix"; } - }; - } - - static ChatModel syncModel(Function call) { - return new ChatModel() { - @Override public ChatResponse call(Prompt prompt) { return call.apply(prompt); } - @Override public Flux stream(Prompt prompt) { return Flux.just(response("streamed-" + prompt.getContents(), null, null)); } - @Override public ChatOptions getDefaultOptions() { return null; } - }; - } - - static final class FixedEmbeddingModel implements EmbeddingModel { - private final EmbeddingResponse response = new EmbeddingResponse( - List.of(new Embedding(new float[]{1f, 2f, 3f}, 0), new Embedding(new float[]{4f, 5f, 6f}, 1)), - new EmbeddingResponseMetadata("m", null)); - @Override public EmbeddingResponse call(EmbeddingRequest request) { return response; } - @Override public float[] embed(Document document) { return new float[]{1f, 2f, 3f}; } - @Override public List embed(List texts) { return List.of(new float[]{1f, 2f, 3f}, new float[]{4f, 5f, 6f}); } - } -} From 46bfcb2ea5fd3f83d98b6cedae36f8a04163a04f Mon Sep 17 00:00:00 2001 From: Nikhil Pareek Date: Sat, 3 Oct 2026 20:47:47 -0700 Subject: [PATCH 3/7] build(java): drop hardcoded ByteBuddy agent from Surefire The parent Surefire argLine pinned -javaagent to a ByteBuddy agent JAR at a fixed path in the local repository. On a clean CI cache that path does not exist, so the forked test JVM dies during initialization before running anything ("Error occurred during initialization of VM"), failing traceai-java-core with zero tests executed. Java 17 runs the suite without a startup agent. Verified traceai-java-core (13 tests) green with an empty local repository. TH-7477 --- java/pom.xml | 6 ++++-- 1 file changed, 4 insertions(+), 2 deletions(-) diff --git a/java/pom.xml b/java/pom.xml index cf36d9d4..a7fdd09c 100644 --- a/java/pom.xml +++ b/java/pom.xml @@ -183,9 +183,11 @@ maven-surefire-plugin 3.2.3 + - -XX:+EnableDynamicAgentLoading - -javaagent:${user.home}/.m2/repository/net/bytebuddy/byte-buddy-agent/1.14.10/byte-buddy-agent-1.14.10.jar --add-opens java.base/java.lang=ALL-UNNAMED --add-opens java.base/java.lang.reflect=ALL-UNNAMED --add-opens java.base/java.util=ALL-UNNAMED From b658b5963f75abb8fcf833133e152e1415d7c207 Mon Sep 17 00:00:00 2001 From: Nikhil Pareek Date: Sat, 3 Oct 2026 20:50:11 -0700 Subject: [PATCH 4/7] build(spring-ai): resolve candidate JARs relative to the repo root The compatibility consumer POMs pointed at /candidate-jars, which only exists inside the local evidence runner's mount, so CI failed dependency resolution after the suite itself passed. Resolve from /candidate-jars via the module directory instead, and recompile the consumers in CI. TH-7477 --- .github/workflows/java-spring-ai-compatibility.yml | 8 ++++---- java/compatibility/spring-ai/baseline-cell/pom.xml | 4 ++-- java/compatibility/spring-ai/common/pom.xml | 4 ++-- java/compatibility/spring-ai/primary/pom.xml | 4 ++-- 4 files changed, 10 insertions(+), 10 deletions(-) diff --git a/.github/workflows/java-spring-ai-compatibility.yml b/.github/workflows/java-spring-ai-compatibility.yml index a4b256b9..bdd3fb4e 100644 --- a/.github/workflows/java-spring-ai-compatibility.yml +++ b/.github/workflows/java-spring-ai-compatibility.yml @@ -31,7 +31,7 @@ jobs: - name: Package candidate JARs working-directory: java run: mvn -B -ntp -pl traceai-spring-ai,traceai-spring-boot-starter -am -DskipTests package - - name: Stage candidate artifacts + - name: Stage candidate artifacts at the repo root run: | mkdir -p candidate-jars cp java/traceai-java-core/target/traceai-java-core-1.0.0.jar candidate-jars/ @@ -39,13 +39,13 @@ jobs: sha256sum candidate-jars/*.jar - name: Same-binary consumer on Spring AI 1.1.8 working-directory: java/compatibility/spring-ai/common - run: mvn -B -ntp -q exec:java -Dexec.mainClass=local.th7477.CommonConsumer -Dexec.args=P -Dexec.classpathScope=compile + run: mvn -B -ntp -q clean compile exec:java -Dexec.mainClass=local.th7477.CommonConsumer -Dexec.args=P -Dexec.classpathScope=compile - name: Same-binary consumer on Spring AI 1.1.0 working-directory: java/compatibility/spring-ai/baseline-cell - run: mvn -B -ntp -o -q exec:java -Dexec.mainClass=local.th7477.CommonConsumer -Dexec.args=B -Dexec.classpathScope=compile + run: mvn -B -ntp -q clean compile exec:java -Dexec.mainClass=local.th7477.CommonConsumer -Dexec.args=B -Dexec.classpathScope=compile - name: 1.1.8-only embedding accessor working-directory: java/compatibility/spring-ai/primary - run: mvn -B -ntp -o -q exec:java -Dexec.mainClass=local.th7477.PrimaryConsumer -Dexec.classpathScope=compile + run: mvn -B -ntp -q clean compile exec:java -Dexec.mainClass=local.th7477.PrimaryConsumer -Dexec.classpathScope=compile - name: Example compiles against 1.1.8 working-directory: java/examples/spring-ai-example run: mvn -B -ntp -DskipTests compile diff --git a/java/compatibility/spring-ai/baseline-cell/pom.xml b/java/compatibility/spring-ai/baseline-cell/pom.xml index 4326e269..4450751d 100644 --- a/java/compatibility/spring-ai/baseline-cell/pom.xml +++ b/java/compatibility/spring-ai/baseline-cell/pom.xml @@ -15,8 +15,8 @@ org.springframework.aispring-ai-model org.springframework.aispring-ai-client-chat - ai.traceaitraceai-java-core1.0.0-candidatesystem/candidate-jars/traceai-java-core-1.0.0.jar - ai.traceaitraceai-spring-ai1.0.0-candidatesystem/candidate-jars/traceai-spring-ai-1.0.0.jar + ai.traceaitraceai-java-core1.0.0-candidatesystem${project.basedir}/../../../../candidate-jars/traceai-java-core-1.0.0.jar + ai.traceaitraceai-spring-ai1.0.0-candidatesystem${project.basedir}/../../../../candidate-jars/traceai-spring-ai-1.0.0.jar io.opentelemetryopentelemetry-sdk-testing1.49.0 com.google.code.gsongson2.10.1 diff --git a/java/compatibility/spring-ai/common/pom.xml b/java/compatibility/spring-ai/common/pom.xml index 8cdcc03e..c368bf6d 100644 --- a/java/compatibility/spring-ai/common/pom.xml +++ b/java/compatibility/spring-ai/common/pom.xml @@ -15,8 +15,8 @@ org.springframework.aispring-ai-model org.springframework.aispring-ai-client-chat - ai.traceaitraceai-java-core1.0.0-candidatesystem/candidate-jars/traceai-java-core-1.0.0.jar - ai.traceaitraceai-spring-ai1.0.0-candidatesystem/candidate-jars/traceai-spring-ai-1.0.0.jar + ai.traceaitraceai-java-core1.0.0-candidatesystem${project.basedir}/../../../../candidate-jars/traceai-java-core-1.0.0.jar + ai.traceaitraceai-spring-ai1.0.0-candidatesystem${project.basedir}/../../../../candidate-jars/traceai-spring-ai-1.0.0.jar io.opentelemetryopentelemetry-sdk-testing1.49.0 com.google.code.gsongson2.10.1 diff --git a/java/compatibility/spring-ai/primary/pom.xml b/java/compatibility/spring-ai/primary/pom.xml index f4eff421..85870b34 100644 --- a/java/compatibility/spring-ai/primary/pom.xml +++ b/java/compatibility/spring-ai/primary/pom.xml @@ -17,8 +17,8 @@ org.springframework.aispring-ai-ollama org.springframework.aispring-ai-client-chat org.springframework.aispring-ai-client-chat - ai.traceaitraceai-java-core1.0.0-candidatesystem/candidate-jars/traceai-java-core-1.0.0.jar - ai.traceaitraceai-spring-ai1.0.0-candidatesystem/candidate-jars/traceai-spring-ai-1.0.0.jar + ai.traceaitraceai-java-core1.0.0-candidatesystem${project.basedir}/../../../../candidate-jars/traceai-java-core-1.0.0.jar + ai.traceaitraceai-spring-ai1.0.0-candidatesystem${project.basedir}/../../../../candidate-jars/traceai-spring-ai-1.0.0.jar io.opentelemetryopentelemetry-sdk-testing1.49.0 com.google.code.gsongson2.10.1 From b085ce06054205a9fe071ef68e60a2219394ff7d Mon Sep 17 00:00:00 2001 From: Nikhil Pareek Date: Sat, 3 Oct 2026 22:17:16 -0700 Subject: [PATCH 5/7] ci(spring-ai): install the candidate before compiling the example The example lives outside the reactor and depends on traceai-spring-boot-starter:1.0.0, which is not published. CI resolved everything else and then failed only on that dependency. Install the just-built modules into the local repository first. TH-7477 --- .github/workflows/java-spring-ai-compatibility.yml | 3 +++ 1 file changed, 3 insertions(+) diff --git a/.github/workflows/java-spring-ai-compatibility.yml b/.github/workflows/java-spring-ai-compatibility.yml index bdd3fb4e..bf84a6bb 100644 --- a/.github/workflows/java-spring-ai-compatibility.yml +++ b/.github/workflows/java-spring-ai-compatibility.yml @@ -46,6 +46,9 @@ jobs: - name: 1.1.8-only embedding accessor working-directory: java/compatibility/spring-ai/primary run: mvn -B -ntp -q clean compile exec:java -Dexec.mainClass=local.th7477.PrimaryConsumer -Dexec.classpathScope=compile + - name: Install the candidate into the local repository + working-directory: java + run: mvn -B -ntp -pl traceai-spring-boot-starter -am -DskipTests install - name: Example compiles against 1.1.8 working-directory: java/examples/spring-ai-example run: mvn -B -ntp -DskipTests compile From 31a84b87dc611fdd31896147f8c768891e3bc9d3 Mon Sep 17 00:00:00 2001 From: Nikhil Pareek Date: Sat, 3 Oct 2026 23:33:52 -0700 Subject: [PATCH 6/7] test(spring-ai): cover the starter auto-configuration in isolation Review of b085ce0 found the approved AC-12 starter checks were never written, so a disabled context, a custom tracer bean and the uninitialized wrapper were unverified. Add them, each in its own class because TraceAI's global registration cannot be reset and the starter forks a fresh JVM per class. Also close the three non-blocking notes: record a real EmptyUsage as absent while keeping a measured zero, reject a null upstream subscription and end the span even when recording the error throws, and skip null texts in the embedding preview. TH-7477 --- .../traceai/spring/TracedEmbeddingModel.java | 5 +- .../ai/traceai/spring/TracingSubscriber.java | 5 ++ .../traceai/spring/TracedChatModelTest.java | 23 ++++++ .../TraceAIAutoConfigurationTest.java | 82 +++++++++++++++++++ .../TraceAICustomTracerTest.java | 37 +++++++++ .../autoconfigure/TraceAIDisabledTest.java | 25 ++++++ .../TraceAIUninitializedTest.java | 28 +++++++ 7 files changed, 204 insertions(+), 1 deletion(-) create mode 100644 java/traceai-spring-boot-starter/src/test/java/ai/traceai/spring/autoconfigure/TraceAIAutoConfigurationTest.java create mode 100644 java/traceai-spring-boot-starter/src/test/java/ai/traceai/spring/autoconfigure/TraceAICustomTracerTest.java create mode 100644 java/traceai-spring-boot-starter/src/test/java/ai/traceai/spring/autoconfigure/TraceAIDisabledTest.java create mode 100644 java/traceai-spring-boot-starter/src/test/java/ai/traceai/spring/autoconfigure/TraceAIUninitializedTest.java 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 bcac94af..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 @@ -202,7 +202,10 @@ private void captureInputTexts(Span span, List texts) { if (i > 0) { preview.append("\n---\n"); } - preview.append(texts.get(i)); + 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"); 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 index b62ed411..1b2e3f19 100644 --- 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 @@ -59,6 +59,9 @@ public reactor.util.context.Context 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; @@ -135,6 +138,8 @@ public void onError(Throwable t) { } 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(); } 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 index faeabcc6..05d9cce8 100644 --- 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 @@ -8,6 +8,8 @@ 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; @@ -138,6 +140,27 @@ void tokenCountsAreIndependentlyNullableAndZeroIsPreserved() { } } + @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()) { 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"); + } +} From c1e3c6b8318c325c3b887c124e66a305ec739115 Mon Sep 17 00:00:00 2001 From: Nikhil Pareek Date: Wed, 7 Oct 2026 02:21:49 -0700 Subject: [PATCH 7/7] ci(spring-ai): run the starter tests in the compatibility job The deterministic step only ran traceai-spring-ai (and its dependencies), so the four starter auto-configuration test classes added in 31a84b8 never ran in CI; the starter was only packaged and installed with -DskipTests. Co-Authored-By: Claude Opus 5.5 --- .github/workflows/java-spring-ai-compatibility.yml | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/.github/workflows/java-spring-ai-compatibility.yml b/.github/workflows/java-spring-ai-compatibility.yml index bf84a6bb..ec378a16 100644 --- a/.github/workflows/java-spring-ai-compatibility.yml +++ b/.github/workflows/java-spring-ai-compatibility.yml @@ -25,9 +25,9 @@ jobs: image: maven@sha256:e4a7ace3dc0d645ed97f8d9ad0b0d3f0b14fa8d150138f27f116d7105a639b82 steps: - uses: actions/checkout@v4 - - name: Deterministic Spring module suite (Spring AI 1.1.8) + - name: Deterministic Spring module and starter suites (Spring AI 1.1.8) working-directory: java - run: mvn -B -ntp -pl traceai-spring-ai -am -Dsurefire.failIfNoSpecifiedTests=false test + run: mvn -B -ntp -pl traceai-spring-ai,traceai-spring-boot-starter -am -Dsurefire.failIfNoSpecifiedTests=false test - name: Package candidate JARs working-directory: java run: mvn -B -ntp -pl traceai-spring-ai,traceai-spring-boot-starter -am -DskipTests package