Skip to content

Commit b72dab5

Browse files
committed
feat(core): Support return_direct tool call
1 parent 9c0ff69 commit b72dab5

10 files changed

Lines changed: 289 additions & 0 deletions

File tree

agentscope-core/src/main/java/io/agentscope/core/ReActAgent.java

Lines changed: 35 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -724,6 +724,19 @@ private Mono<Msg> acting(int iter) {
724724
buildSuspendedMsg(pendingPairs));
725725
}
726726

727+
// Return directly: if any tool result has
728+
// return_directly flag, breaking iterations
729+
boolean hasReturnDirect =
730+
successPairs.stream()
731+
.anyMatch(
732+
entry ->
733+
entry.getValue()
734+
.isReturnDirect());
735+
if (hasReturnDirect) {
736+
return Mono.just(
737+
buildReturnDirectMsg(successPairs));
738+
}
739+
727740
// Continue next iteration
728741
return executeIteration(iter + 1);
729742
});
@@ -753,6 +766,28 @@ private Msg buildSuspendedMsg(List<Map.Entry<ToolUseBlock, ToolResultBlock>> pen
753766
.build();
754767
}
755768

769+
/**
770+
* Build a return_direct message from tool results.
771+
*
772+
* @param successPairs List of (ToolUseBlock, ToolResultBlock) pairs for successful results
773+
* @return Msg with {@link GenerateReason#TOOL_RETURN_DIRECT}
774+
*/
775+
private Msg buildReturnDirectMsg(List<Map.Entry<ToolUseBlock, ToolResultBlock>> successPairs) {
776+
List<ContentBlock> content = new ArrayList<>();
777+
for (Map.Entry<ToolUseBlock, ToolResultBlock> pair : successPairs) {
778+
content.addAll(pair.getValue().getOutput());
779+
}
780+
Msg returnDirectMsg =
781+
Msg.builder()
782+
.name(getName())
783+
.role(MsgRole.ASSISTANT)
784+
.content(content)
785+
.generateReason(GenerateReason.TOOL_RETURN_DIRECT)
786+
.build();
787+
memory.addMessage(returnDirectMsg);
788+
return returnDirectMsg;
789+
}
790+
756791
/**
757792
* Execute tool calls and return paired results.
758793
*

agentscope-core/src/main/java/io/agentscope/core/message/GenerateReason.java

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -47,6 +47,9 @@ public enum GenerateReason {
4747
/** Tool execution was suspended, waiting for user to provide results. */
4848
TOOL_SUSPENDED,
4949

50+
/** Tool result marked with return_directly, breaking ReAct iterations. */
51+
TOOL_RETURN_DIRECT,
52+
5053
/** Reasoning phase was stopped by a Hook (PostReasoningEvent.stopAgent()). */
5154
REASONING_STOP_REQUESTED,
5255

agentscope-core/src/main/java/io/agentscope/core/message/ToolResultBlock.java

Lines changed: 29 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -20,6 +20,7 @@
2020
import com.fasterxml.jackson.annotation.JsonProperty;
2121
import io.agentscope.core.tool.ToolSuspendException;
2222
import java.beans.Transient;
23+
import java.util.HashMap;
2324
import java.util.List;
2425
import java.util.Map;
2526

@@ -37,6 +38,9 @@ public final class ToolResultBlock extends ContentBlock {
3738
/** Metadata key indicating this result is suspended for external execution. */
3839
public static final String METADATA_SUSPENDED = "agentscope_suspended";
3940

41+
/** Metadata key indicating this result should be returned directly, breaking ReAct iterations. */
42+
public static final String METADATA_RETURN_DIRECTLY = "agentscope_return_directly";
43+
4044
private final String id;
4145
private final String name;
4246
private final List<ContentBlock> output;
@@ -126,6 +130,20 @@ public boolean isSuspended() {
126130
return Boolean.TRUE.equals(metadata.get(METADATA_SUSPENDED));
127131
}
128132

133+
/**
134+
* Checks if this result should be returned directly, breaking ReAct iterations.
135+
*
136+
* <p>When true, the agent will return this tool result as the final response immediately
137+
* without sending it back to the model for additional reasoning.
138+
*
139+
* @return true if this result should be returned directly
140+
*/
141+
@Transient
142+
@JsonInclude
143+
public boolean isReturnDirect() {
144+
return Boolean.TRUE.equals(metadata.get(METADATA_RETURN_DIRECTLY));
145+
}
146+
129147
/**
130148
* Creates a suspended tool result from a ToolSuspendException.
131149
*
@@ -289,6 +307,17 @@ public ToolResultBlock withIdAndName(String id, String name) {
289307
return new ToolResultBlock(id, name, this.output, this.metadata);
290308
}
291309

310+
/**
311+
* Creates a new ToolResultBlock with the return_directly metadata flag set.
312+
*
313+
* @return A ToolResultBlock with agentscope_return_directly=true
314+
*/
315+
public ToolResultBlock withReturnDirect() {
316+
Map<String, Object> newMetadata = new HashMap<>(this.metadata);
317+
newMetadata.put(METADATA_RETURN_DIRECTLY, true);
318+
return of(this.id, this.name, this.output, newMetadata);
319+
}
320+
292321
/**
293322
* Creates a new builder for constructing ToolResultBlock instances.
294323
*

agentscope-core/src/main/java/io/agentscope/core/tool/AgentTool.java

Lines changed: 13 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -100,6 +100,19 @@ default Map<String, Object> getOutputSchema() {
100100
return null;
101101
}
102102

103+
/**
104+
* Whether this tool's result should be returned directly, breaking ReAct iterations.
105+
*
106+
* <p>When {@code true}, after this tool executes successfully, the agent returns the tool
107+
* result as the final response without sending it back to the model for additional reasoning.
108+
* This corresponds to the {@code returnDirect} attribute on the {@link Tool @Tool} annotation.
109+
*
110+
* @return true if this tool's result should be returned directly, false by default
111+
*/
112+
default boolean isReturnDirect() {
113+
return false;
114+
}
115+
103116
/**
104117
* Execute the tool with the given parameters (asynchronous).
105118
*

agentscope-core/src/main/java/io/agentscope/core/tool/Tool.java

Lines changed: 11 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -96,6 +96,17 @@
9696
*/
9797
boolean strict() default false;
9898

99+
/**
100+
* Whether this tool's result should be returned directly to the caller, breaking ReAct loop iterations.
101+
*
102+
* <p>When {@code true}, after this tool executes successfully, the agent will return the tool
103+
* result as the final response immediately — without sending it back to the model for
104+
* additional reasoning.
105+
*
106+
* @return true to enable return directly for this tool
107+
*/
108+
boolean returnDirect() default false;
109+
99110
/**
100111
* Custom result converter for this tool.
101112
*

agentscope-core/src/main/java/io/agentscope/core/tool/ToolExecutor.java

Lines changed: 7 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -242,6 +242,13 @@ private Mono<ToolResultBlock> executeCore(ToolCallParam param) {
242242
.build();
243243

244244
return tool.callAsync(executionParam)
245+
.map(
246+
result -> {
247+
if (tool.isReturnDirect()) {
248+
return result.withReturnDirect();
249+
}
250+
return result;
251+
})
245252
.onErrorResume(
246253
ToolSuspendException.class,
247254
e -> {

agentscope-core/src/main/java/io/agentscope/core/tool/Toolkit.java

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -381,6 +381,11 @@ public Boolean getStrict() {
381381
return toolAnnotation.strict() ? Boolean.TRUE : null;
382382
}
383383

384+
@Override
385+
public boolean isReturnDirect() {
386+
return toolAnnotation.returnDirect();
387+
}
388+
384389
@Override
385390
public Mono<ToolResultBlock> callAsync(ToolCallParam param) {
386391
// Pass custom converter to method invoker

agentscope-core/src/test/java/io/agentscope/core/agent/ReActAgentTest.java

Lines changed: 164 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -313,6 +313,170 @@ void testMultipleToolExecution() {
313313
TestConstants.TEST_TOOL_NAME, history.get(1), "Test tool should be called second");
314314
}
315315

316+
@Test
317+
@DisplayName("Should return tool result directly when tool is marked as returnDirect")
318+
void testToolCallingReturnDirect() {
319+
// Register a returnDirect tool
320+
mockToolkit.withReturnDirectTool(
321+
TestConstants.RETURN_DIRECT_TOOL_NAME, "ReturnDirect tool result");
322+
323+
final int[] callCount = {0};
324+
MockModel toolModel =
325+
new MockModel(
326+
messages -> {
327+
int currentCall = callCount[0]++;
328+
if (currentCall == 0) {
329+
// Return tool call for the returnDirect tool
330+
return List.of(
331+
createToolCallResponseHelper(
332+
TestConstants.RETURN_DIRECT_TOOL_NAME,
333+
"tool_call_return_direct_1",
334+
TestUtils.createToolArguments()));
335+
}
336+
// This should never be reached because returnDirect breaks the loop
337+
return List.of(
338+
ChatResponse.builder()
339+
.content(
340+
List.of(
341+
TextBlock.builder()
342+
.text("Should not reach this")
343+
.build()))
344+
.usage(new ChatUsage(10, 20, 30))
345+
.build());
346+
});
347+
348+
agent =
349+
ReActAgent.builder()
350+
.name(TestConstants.TEST_REACT_AGENT_NAME)
351+
.sysPrompt(TestConstants.DEFAULT_SYS_PROMPT)
352+
.model(toolModel)
353+
.toolkit(mockToolkit)
354+
.memory(memory)
355+
.build();
356+
357+
Msg userMsg = TestUtils.createUserMessage("User", "Call the returnDirect tool");
358+
Msg response =
359+
agent.call(userMsg).block(Duration.ofMillis(TestConstants.DEFAULT_TEST_TIMEOUT_MS));
360+
361+
// Verify response
362+
assertNotNull(response, "Response should not be null");
363+
assertEquals(
364+
GenerateReason.TOOL_RETURN_DIRECT,
365+
response.getGenerateReason(),
366+
"GenerateReason should be TOOL_RETURN_DIRECT");
367+
368+
// Verify the response contains the tool result content
369+
String text = TestUtils.extractTextContent(response);
370+
assertTrue(
371+
text.contains("ReturnDirect tool result"),
372+
"Response should contain the return direct tool result");
373+
374+
// Verify the model was called only once (no second reasoning iteration)
375+
assertEquals(1, toolModel.getCallCount(), "Model should be called only once");
376+
377+
// Verify tool was called
378+
assertTrue(
379+
mockToolkit.wasToolCalled(TestConstants.RETURN_DIRECT_TOOL_NAME),
380+
"ReturnDirect tool should be called");
381+
}
382+
383+
@Test
384+
@DisplayName(
385+
"Should return multiple tool results directly when a tool is marked as returnDirect")
386+
void testMultipleToolCallingReturnDirect() {
387+
mockToolkit.withTool(TestConstants.TEST_TOOL_NAME, "Test tool result");
388+
mockToolkit.withReturnDirectTool(
389+
TestConstants.RETURN_DIRECT_TOOL_NAME, "ReturnDirect tool result");
390+
391+
final int[] callCount = {0};
392+
MockModel toolModel =
393+
new MockModel(
394+
messages -> {
395+
int currentCall = callCount[0]++;
396+
if (currentCall == 0) {
397+
Map<String, Object> args1 = new HashMap<>();
398+
args1.put("param1", "value1");
399+
Map<String, Object> args2 = new HashMap<>();
400+
args2.put("param2", "value2");
401+
return List.of(
402+
ChatResponse.builder()
403+
.content(
404+
List.of(
405+
ToolUseBlock.builder()
406+
.id("tool_call_1")
407+
.name(
408+
TestConstants
409+
.TEST_TOOL_NAME)
410+
.input(args1)
411+
.content(
412+
JsonUtils
413+
.getJsonCodec()
414+
.toJson(
415+
args1))
416+
.build(),
417+
ToolUseBlock.builder()
418+
.id(
419+
"tool_call_return_direct_1")
420+
.name(
421+
TestConstants
422+
.RETURN_DIRECT_TOOL_NAME)
423+
.input(args2)
424+
.content(
425+
JsonUtils
426+
.getJsonCodec()
427+
.toJson(
428+
args2))
429+
.build()))
430+
.usage(new ChatUsage(10, 20, 30))
431+
.build());
432+
}
433+
return List.of(
434+
ChatResponse.builder()
435+
.content(
436+
List.of(
437+
TextBlock.builder()
438+
.text("Should not reach this")
439+
.build()))
440+
.usage(new ChatUsage(10, 20, 30))
441+
.build());
442+
});
443+
444+
agent =
445+
ReActAgent.builder()
446+
.name(TestConstants.TEST_REACT_AGENT_NAME)
447+
.sysPrompt(TestConstants.DEFAULT_SYS_PROMPT)
448+
.model(toolModel)
449+
.toolkit(mockToolkit)
450+
.memory(memory)
451+
.build();
452+
453+
Msg userMsg = TestUtils.createUserMessage("User", "Call multiple returnDirect tools");
454+
Msg response =
455+
agent.call(userMsg).block(Duration.ofMillis(TestConstants.DEFAULT_TEST_TIMEOUT_MS));
456+
457+
assertNotNull(response, "Response should not be null");
458+
assertEquals(
459+
GenerateReason.TOOL_RETURN_DIRECT,
460+
response.getGenerateReason(),
461+
"GenerateReason should be TOOL_RETURN_DIRECT");
462+
463+
String text = TestUtils.extractTextContent(response);
464+
assertTrue(text.contains("Test tool result"), "Response should contain test tool result");
465+
assertTrue(
466+
text.contains("ReturnDirect tool result"),
467+
"Response should contain return direct tool result");
468+
469+
assertEquals(1, toolModel.getCallCount(), "Model should be called only once");
470+
471+
assertTrue(
472+
mockToolkit.wasToolCalled(TestConstants.TEST_TOOL_NAME),
473+
"Test tool should be called");
474+
assertTrue(
475+
mockToolkit.wasToolCalled(TestConstants.RETURN_DIRECT_TOOL_NAME),
476+
"ReturnDirect tool should be called");
477+
assertEquals(2, mockToolkit.getCallCount(), "Two tools should be called");
478+
}
479+
316480
@Test
317481
@DisplayName("Should handle tool execution errors in acting phase")
318482
void testToolExecutionError() {

agentscope-core/src/test/java/io/agentscope/core/agent/test/MockToolkit.java

Lines changed: 21 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -64,6 +64,15 @@ public MockToolkit withTool(String toolName, Function<Map<String, Object>, Strin
6464
return this;
6565
}
6666

67+
/**
68+
* Register a tool that returns its result directly (breaking ReAct loop).
69+
*/
70+
public MockToolkit withReturnDirectTool(String toolName, String result) {
71+
toolBehaviors.put(toolName, args -> result);
72+
registerMockTool(toolName, "Mock returnDirect tool: " + toolName, true);
73+
return this;
74+
}
75+
6776
/**
6877
* Register a tool that throws an error.
6978
*/
@@ -117,6 +126,13 @@ private void registerDefaultTools() {
117126
* Register a mock tool with the toolkit.
118127
*/
119128
private void registerMockTool(String name, String description) {
129+
registerMockTool(name, description, false);
130+
}
131+
132+
/**
133+
* Register a mock tool with the toolkit, optionally marking it as returnDirect.
134+
*/
135+
private void registerMockTool(String name, String description, boolean returnDirect) {
120136
AgentTool tool =
121137
new AgentTool() {
122138
@Override
@@ -137,6 +153,11 @@ public Map<String, Object> getParameters() {
137153
return schema;
138154
}
139155

156+
@Override
157+
public boolean isReturnDirect() {
158+
return returnDirect;
159+
}
160+
140161
@Override
141162
public Mono<ToolResultBlock> callAsync(ToolCallParam param) {
142163
return Mono.fromCallable(

0 commit comments

Comments
 (0)