fix(tool-calls): broken trace of tool call (#1954)

* fix(observation): broken trace of tool call

* add license

Author: Cirilla-zmh

spring-ai-alibaba-core/src/main/java/com/alibaba/cloud/ai/dashscope/chat/DashScopeChatModel.java

@@ -30,6 +30,7 @@
 import com.alibaba.cloud.ai.dashscope.api.DashScopeApi.FunctionTool;
 import com.alibaba.cloud.ai.dashscope.chat.observation.DashScopeChatModelObservationConvention;
 import com.alibaba.cloud.ai.dashscope.common.DashScopeApiConstants;
+import com.alibaba.cloud.ai.tool.observation.inner.ToolCallReactiveContextHolder;
 import io.micrometer.observation.Observation;
 import io.micrometer.observation.ObservationRegistry;
 import io.micrometer.observation.contextpropagation.ObservationThreadLocalAccessor;
@@ -260,22 +261,25 @@ public Flux<ChatResponse> internalStream(Prompt prompt, ChatResponse previousCha
 			// @formatter:off
 			Flux<ChatResponse> flux = chatResponse.flatMap(response -> {
 					if (toolExecutionEligibilityPredicate.isToolExecutionRequired(prompt.getOptions(), response)) {
-						return Flux.defer(
-								() -> {
-									var toolExecutionResult = this.toolCallingManager.executeToolCalls(prompt, response);
-									if (toolExecutionResult.returnDirect()) {
-										// Return tool execution result directly to the client.
-										return Flux.just(ChatResponse.builder().from(response)
-												.generations(ToolExecutionResult.buildGenerations(toolExecutionResult))
-												.build());
-									} else {
-										// Send the tool execution result back to the model.
-										return this.internalStream(new Prompt(toolExecutionResult.conversationHistory(), prompt.getOptions()),
-												response);
-									}
-								}
-						).subscribeOn(Schedulers.boundedElastic());
-
+						return Flux.deferContextual((ctx) -> {
+							ToolExecutionResult toolExecutionResult;
+							try {
+								ToolCallReactiveContextHolder.setContext(ctx);
+								toolExecutionResult = this.toolCallingManager.executeToolCalls(prompt, response);
+							} finally {
+								ToolCallReactiveContextHolder.clearContext();
+							}
+							if (toolExecutionResult.returnDirect()) {
+								// Return tool execution result directly to the client.
+								return Flux.just(ChatResponse.builder().from(response)
+										.generations(ToolExecutionResult.buildGenerations(toolExecutionResult))
+										.build());
+							} else {
+								// Send the tool execution result back to the model.
+								return this.internalStream(new Prompt(toolExecutionResult.conversationHistory(), prompt.getOptions()),
+										response);
+							}
+						}).subscribeOn(Schedulers.boundedElastic());
 					}
 					else {
 						return Flux.just(response);

spring-ai-alibaba-core/src/main/java/com/alibaba/cloud/ai/tool/ObservableToolCallingManager.java

@@ -18,7 +18,9 @@
 import com.alibaba.cloud.ai.tool.observation.ArmsToolCallingObservationContext;
 import com.alibaba.cloud.ai.tool.observation.ArmsToolCallingObservationConvention;
 import com.alibaba.cloud.ai.tool.observation.ArmsToolCallingObservationDocumentation;
+import com.alibaba.cloud.ai.tool.observation.inner.ToolCallReactiveContextHolder;
 import io.micrometer.observation.ObservationRegistry;
+import io.micrometer.observation.contextpropagation.ObservationThreadLocalAccessor;
 import java.util.ArrayList;
 import java.util.HashMap;
 import java.util.Iterator;
@@ -49,6 +51,7 @@
 import org.springframework.util.Assert;
 import org.springframework.util.CollectionUtils;
 import org.springframework.util.StringUtils;
+import reactor.util.context.ContextView;
 
 /**
  * Inspired from org.springframework.ai.model.tool.DefaultToolCallingManager.
@@ -215,6 +218,12 @@ private InternalToolExecutionResult executeToolCall(Prompt prompt, AssistantMess
 				.returnDirect(returnDirect)
 				.build();
 
+			ContextView contextView = ToolCallReactiveContextHolder.getContext();
+			if (contextView != null) {
+				observationContext
+					.setParentObservation(contextView.getOrDefault(ObservationThreadLocalAccessor.KEY, null));
+			}
+
 			String toolResult = ArmsToolCallingObservationDocumentation.EXECUTE_TOOL_OPERATION
 				.observation(this.observationConvention, DEFAULT_OBSERVATION_CONVENTION, () -> observationContext,
 						this.observationRegistry)

spring-ai-alibaba-core/src/main/java/com/alibaba/cloud/ai/tool/observation/inner/ToolCallReactiveContextHolder.java

@@ -0,0 +1,48 @@
+/*
+ * Copyright 2024-2025 the original author or authors.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      https://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package com.alibaba.cloud.ai.tool.observation.inner;
+
+import reactor.util.context.Context;
+import reactor.util.context.ContextView;
+
+/**
+ * Copied from <a
+ * href=https://github.com/spring-projects/spring-ai/blob/main/spring-ai-model/src/main/java/org/springframework/ai/model/tool/internal/ToolCallReactiveContextHolder.java>upstream
+ * (1.0.0)</a>. <br/>
+ * This class bridges blocking Tools call and the reactive context. When calling tools, it
+ * captures the context in a thread local, making it available to re-inject in a nested
+ * reactive call.
+ *
+ * @author Daniel Garnier-Moiroux
+ * @since 1.1.0
+ */
+public class ToolCallReactiveContextHolder {
+
+	private static final ThreadLocal<ContextView> context = ThreadLocal.withInitial(Context::empty);
+
+	public static void setContext(ContextView contextView) {
+		context.set(contextView);
+	}
+
+	public static ContextView getContext() {
+		return context.get();
+	}
+
+	public static void clearContext() {
+		context.remove();
+	}
+
+}