fix: stream partial chunks, then one aggregated response

contrib/spring-ai/src/main/java/com/google/adk/models/springai/SpringAI.java

@@ -18,20 +18,26 @@
 import com.fasterxml.jackson.databind.ObjectMapper;
 import com.google.adk.models.BaseLlm;
 import com.google.adk.models.BaseLlmConnection;
+import com.google.adk.models.FunctionCallIds;
 import com.google.adk.models.LlmRequest;
 import com.google.adk.models.LlmResponse;
 import com.google.adk.models.springai.error.SpringAIErrorMapper;
 import com.google.adk.models.springai.observability.SpringAIObservabilityHandler;
 import com.google.adk.models.springai.properties.SpringAIProperties;
-import io.reactivex.rxjava3.core.BackpressureStrategy;
+import com.google.genai.types.Content;
+import com.google.genai.types.FunctionCall;
+import com.google.genai.types.GenerateContentResponseUsageMetadata;
+import com.google.genai.types.Part;
 import io.reactivex.rxjava3.core.Flowable;
+import java.util.ArrayList;
+import java.util.List;
 import java.util.Objects;
+import org.jspecify.annotations.Nullable;
 import org.springframework.ai.chat.model.ChatModel;
 import org.springframework.ai.chat.model.ChatResponse;
 import org.springframework.ai.chat.model.StreamingChatModel;
 import org.springframework.ai.chat.prompt.ChatOptions;
 import org.springframework.ai.chat.prompt.Prompt;
-import reactor.core.publisher.Flux;
 
 /**
  * Spring AI implementation of BaseLlm that wraps Spring AI ChatModel and StreamingChatModel.
@@ -195,60 +201,158 @@ private Flowable<LlmResponse> generateContent(LlmRequest llmRequest) {
   }
 
   private Flowable<LlmResponse> generateStreamingContent(LlmRequest llmRequest) {
-    SpringAIObservabilityHandler.RequestContext context =
-        observabilityHandler.startRequest(model(), "streaming");
-
-    return Flowable.create(
-        emitter -> {
-          try {
-            Prompt prompt = messageConverter.toLlmPrompt(llmRequest, resolveDefaultOptions());
-            observabilityHandler.logRequest(prompt.toString(), model());
-
-            Flux<ChatResponse> responseFlux = streamingChatModel.stream(prompt);
-
-            responseFlux
-                .doOnError(
-                    error -> {
-                      observabilityHandler.recordError(context, error);
-                      SpringAIErrorMapper.MappedError mappedError =
-                          SpringAIErrorMapper.mapError(error);
-                      emitter.onError(
-                          new RuntimeException(mappedError.getNormalizedMessage(), error));
-                    })
-                .subscribe(
-                    chatResponse -> {
-                      try {
-                        // Use enhanced streaming-aware conversion
-                        LlmResponse llmResponse =
-                            messageConverter.toLlmResponse(chatResponse, true);
-                        emitter.onNext(llmResponse);
-                      } catch (Exception e) {
-                        observabilityHandler.recordError(context, e);
-                        SpringAIErrorMapper.MappedError mappedError =
-                            SpringAIErrorMapper.mapError(e);
-                        emitter.onError(
-                            new RuntimeException(mappedError.getNormalizedMessage(), e));
-                      }
-                    },
-                    error -> {
-                      observabilityHandler.recordError(context, error);
-                      SpringAIErrorMapper.MappedError mappedError =
-                          SpringAIErrorMapper.mapError(error);
-                      emitter.onError(
-                          new RuntimeException(mappedError.getNormalizedMessage(), error));
-                    },
-                    () -> {
-                      // Record success for streaming completion
-                      observabilityHandler.recordSuccess(context, 0, 0, 0);
-                      emitter.onComplete();
-                    });
-          } catch (Exception e) {
-            observabilityHandler.recordError(context, e);
-            SpringAIErrorMapper.MappedError mappedError = SpringAIErrorMapper.mapError(e);
-            emitter.onError(new RuntimeException(mappedError.getNormalizedMessage(), e));
+    return Flowable.defer(
+        () -> {
+          SpringAIObservabilityHandler.RequestContext context =
+              observabilityHandler.startRequest(model(), "streaming");
+          // Deferred so a failure building the prompt also reaches onErrorResumeNext.
+          return Flowable.defer(() -> streamResponses(llmRequest))
+              .doOnComplete(() -> observabilityHandler.recordSuccess(context, 0, 0, 0))
+              .onErrorResumeNext(
+                  error -> {
+                    observabilityHandler.recordError(context, error);
+                    SpringAIErrorMapper.MappedError mappedError =
+                        SpringAIErrorMapper.mapError(error);
+                    return Flowable.error(
+                        new RuntimeException(mappedError.getNormalizedMessage(), error));
+                  });
+        });
+  }
+
+  /**
+   * Emits each chunk that has content as a partial response and, if the model streamed anything,
+   * then one aggregated response, as {@code Gemini} does. ADK saves only the aggregated response to
+   * the session and runs its function calls.
+   */
+  private Flowable<LlmResponse> streamResponses(LlmRequest llmRequest) {
+    Prompt prompt = messageConverter.toLlmPrompt(llmRequest, resolveDefaultOptions());
+    observabilityHandler.logRequest(prompt.toString(), model());
+    StreamAggregation aggregation = new StreamAggregation();
+    return Flowable.fromPublisher(streamingChatModel.stream(prompt))
+        .concatMap(aggregation::processChunk)
+        .concatWith(Flowable.defer(aggregation::aggregatedResponse));
+  }
+
+  /** The state of one stream, used to build its aggregated response. */
+  private final class StreamAggregation {
+    // Parts in stream order. Consecutive text, or consecutive thought text, merges into one part;
+    // function calls, media and other parts are kept as they arrived.
+    private final List<Part> parts = new ArrayList<>();
+    private final StringBuilder text = new StringBuilder();
+    private boolean textIsThought;
+    private @Nullable GenerateContentResponseUsageMetadata usage;
+    private @Nullable String modelVersion;
+    private @Nullable LlmResponse finishingChunk;
+    private boolean receivedChunk;
+
+    private Flowable<LlmResponse> processChunk(ChatResponse chatResponse) {
+      receivedChunk = true;
+      LlmResponse chunk = withFunctionCallIds(messageConverter.toLlmResponse(chatResponse));
+      chunk.usageMetadata().ifPresent(chunkUsage -> usage = chunkUsage);
+      chunk.modelVersion().ifPresent(version -> modelVersion = version);
+      if (chunk.finishReason().isPresent()) {
+        finishingChunk = chunk;
+      }
+      List<Part> chunkParts = chunk.content().flatMap(Content::parts).orElse(List.of());
+      if (chunkParts.isEmpty()) {
+        return Flowable.empty();
+      }
+      for (Part part : chunkParts) {
+        if (part.text().isPresent()) {
+          boolean isThought = part.thought().orElse(false);
+          if (isThought != textIsThought) {
+            flushText();
           }
-        },
-        BackpressureStrategy.BUFFER);
+          textIsThought = isThought;
+          text.append(part.text().get());
+        } else {
+          flushText();
+          parts.add(part);
+        }
+      }
+      // As in Gemini, only the aggregated response reports an error.
+      return Flowable.just(
+          chunk.toBuilder()
+              .partial(true)
+              .turnComplete(null)
+              .errorCode(null)
+              .errorMessage(null)
+              .build());
+    }
+
+    private void flushText() {
+      if (text.isEmpty()) {
+        return;
+      }
+      Part.Builder part = Part.builder().text(text.toString());
+      if (textIsThought) {
+        part.thought(true);
+      }
+      parts.add(part.build());
+      text.setLength(0);
+    }
+
+    private Flowable<LlmResponse> aggregatedResponse() {
+      if (!receivedChunk) {
+        return Flowable.empty();
+      }
+      flushText();
+      LlmResponse.Builder aggregated =
+          LlmResponse.builder().partial(false).turnComplete(true).modelVersion(modelVersion);
+      // A stream without content still ends with a response, as Gemini's does.
+      if (!parts.isEmpty()) {
+        aggregated.content(Content.builder().role("model").parts(parts).build());
+      }
+      if (usage != null) {
+        aggregated.usageMetadata(usage);
+      }
+      // Report how the stream ended, from the last chunk that said so, as ADK Python does.
+      if (finishingChunk != null) {
+        aggregated
+            .finishReason(finishingChunk.finishReason().orElse(null))
+            .errorCode(finishingChunk.errorCode().orElse(null))
+            .errorMessage(finishingChunk.errorMessage().orElse(null));
+      }
+      return Flowable.just(aggregated.build());
+    }
+  }
+
+  /**
+   * Gives each function call without an id a client-side one, so the partial and aggregated
+   * responses carrying it share the same id.
+   */
+  private static LlmResponse withFunctionCallIds(LlmResponse response) {
+    if (response.content().isEmpty()) {
+      return response;
+    }
+    Content content = response.content().get();
+    List<Part> parts = content.parts().orElse(List.of());
+    if (parts.stream()
+        .noneMatch(part -> part.functionCall().filter(SpringAI::lacksId).isPresent())) {
+      return response;
+    }
+    return response.toBuilder()
+        .content(
+            content.toBuilder()
+                .parts(parts.stream().map(SpringAI::withFunctionCallId).toList())
+                .build())
+        .build();
+  }
+
+  private static Part withFunctionCallId(Part part) {
+    return part.functionCall()
+        .filter(SpringAI::lacksId)
+        .map(
+            call ->
+                part.toBuilder()
+                    .functionCall(
+                        call.toBuilder().id(FunctionCallIds.generateClientFunctionCallId()).build())
+                    .build())
+        .orElse(part);
+  }
+
+  private static boolean lacksId(FunctionCall call) {
+    return call.id().map(String::isEmpty).orElse(true);
   }
 
   @Override