Skip to content

Commit 7705523

Browse files
Fix UUID non-determinism, null metadata NPE, and unbounded tool loop
T5: Replace UUID.randomUUID() with Workflow.randomUUID() in LocalActivityToolCallbackWrapper to ensure deterministic replay. T7: Convert recursive tool call loop in ActivityChatModel.call() to iterative loop with MAX_TOOL_CALL_ITERATIONS (10) limit to prevent infinite recursion from misbehaving models. T14: Fix NPE when ChatResponse metadata is null by only calling .metadata() on the builder when metadata is non-null. Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
1 parent dc4a39c commit 7705523

File tree

2 files changed

+30
-21
lines changed

2 files changed

+30
-21
lines changed

temporal-spring-ai/src/main/java/io/temporal/springai/model/ActivityChatModel.java

Lines changed: 29 additions & 19 deletions
Original file line numberDiff line numberDiff line change
@@ -83,6 +83,9 @@ public class ActivityChatModel implements ChatModel {
8383
/** Default maximum retry attempts for chat model activity calls. */
8484
public static final int DEFAULT_MAX_ATTEMPTS = 3;
8585

86+
/** Maximum number of tool call iterations before aborting to prevent infinite loops. */
87+
public static final int MAX_TOOL_CALL_ITERATIONS = 10;
88+
8689
private final ChatModelActivity chatModelActivity;
8790
private final String modelName;
8891
private final ToolCallingManager toolCallingManager;
@@ -190,33 +193,41 @@ public ChatOptions getDefaultOptions() {
190193

191194
@Override
192195
public ChatResponse call(Prompt prompt) {
193-
// Convert prompt to activity input and call the activity
194-
ChatModelTypes.ChatModelActivityInput input = createActivityInput(prompt);
195-
ChatModelTypes.ChatModelActivityOutput output = chatModelActivity.callChatModel(input);
196+
Prompt currentPrompt = prompt;
197+
198+
for (int iteration = 0; iteration < MAX_TOOL_CALL_ITERATIONS; iteration++) {
199+
// Convert prompt to activity input and call the activity
200+
ChatModelTypes.ChatModelActivityInput input = createActivityInput(currentPrompt);
201+
ChatModelTypes.ChatModelActivityOutput output = chatModelActivity.callChatModel(input);
196202

197-
// Convert activity output to ChatResponse
198-
ChatResponse response = toResponse(output);
203+
// Convert activity output to ChatResponse
204+
ChatResponse response = toResponse(output);
199205

200-
// Handle tool calls if the model requested them
201-
if (prompt.getOptions() != null
202-
&& toolExecutionEligibilityPredicate.isToolExecutionRequired(
203-
prompt.getOptions(), response)) {
206+
// If no tool calls requested, return the response
207+
if (currentPrompt.getOptions() == null
208+
|| !toolExecutionEligibilityPredicate.isToolExecutionRequired(
209+
currentPrompt.getOptions(), response)) {
210+
return response;
211+
}
204212

205-
var toolExecutionResult = toolCallingManager.executeToolCalls(prompt, response);
213+
var toolExecutionResult = toolCallingManager.executeToolCalls(currentPrompt, response);
206214

207215
if (toolExecutionResult.returnDirect()) {
208-
// Return tool execution result directly
209216
return ChatResponse.builder()
210217
.from(response)
211218
.generations(ToolExecutionResult.buildGenerations(toolExecutionResult))
212219
.build();
213-
} else {
214-
// Send tool results back to the model (recursive call)
215-
return call(new Prompt(toolExecutionResult.conversationHistory(), prompt.getOptions()));
216220
}
221+
222+
// Continue loop with tool results sent back to the model
223+
currentPrompt =
224+
new Prompt(toolExecutionResult.conversationHistory(), currentPrompt.getOptions());
217225
}
218226

219-
return response;
227+
throw new IllegalStateException(
228+
"Chat model exceeded maximum tool call iterations ("
229+
+ MAX_TOOL_CALL_ITERATIONS
230+
+ "). This may indicate the model is stuck in a tool-calling loop.");
220231
}
221232

222233
private ChatModelTypes.ChatModelActivityInput createActivityInput(Prompt prompt) {
@@ -341,12 +352,11 @@ private ChatResponse toResponse(ChatModelTypes.ChatModelActivityOutput output) {
341352
.map(gen -> new Generation(toAssistantMessage(gen.message())))
342353
.collect(Collectors.toList());
343354

344-
ChatResponseMetadata metadata = null;
355+
var builder = ChatResponse.builder().generations(generations);
345356
if (output.metadata() != null) {
346-
metadata = ChatResponseMetadata.builder().model(output.metadata().model()).build();
357+
builder.metadata(ChatResponseMetadata.builder().model(output.metadata().model()).build());
347358
}
348-
349-
return ChatResponse.builder().generations(generations).metadata(metadata).build();
359+
return builder.build();
350360
}
351361

352362
private AssistantMessage toAssistantMessage(ChatModelTypes.Message message) {

temporal-spring-ai/src/main/java/io/temporal/springai/tool/LocalActivityToolCallbackWrapper.java

Lines changed: 1 addition & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -4,7 +4,6 @@
44
import io.temporal.workflow.Workflow;
55
import java.time.Duration;
66
import java.util.Map;
7-
import java.util.UUID;
87
import java.util.concurrent.ConcurrentHashMap;
98
import org.springframework.ai.chat.model.ToolContext;
109
import org.springframework.ai.tool.ToolCallback;
@@ -87,7 +86,7 @@ public ToolMetadata getToolMetadata() {
8786

8887
@Override
8988
public String call(String toolInput) {
90-
String callbackId = UUID.randomUUID().toString();
89+
String callbackId = Workflow.randomUUID().toString();
9190
try {
9291
CALLBACK_REGISTRY.put(callbackId, delegate);
9392
return stub.call(callbackId, toolInput);

0 commit comments

Comments
 (0)