diff --git a/agentscope-core/src/main/java/io/agentscope/core/ReActAgent.java b/agentscope-core/src/main/java/io/agentscope/core/ReActAgent.java index 1138ff8042..8185ae842d 100644 --- a/agentscope-core/src/main/java/io/agentscope/core/ReActAgent.java +++ b/agentscope-core/src/main/java/io/agentscope/core/ReActAgent.java @@ -2595,14 +2595,19 @@ private void emitBlockEvents( blockLifecycle.startText(events); if (tb.getText() != null && !tb.getText().isEmpty()) { events.add( - new TextBlockDeltaEvent(blockLifecycle.replyId, "text", tb.getText())); + new TextBlockDeltaEvent( + blockLifecycle.replyId, + blockLifecycle.currentTextBlockId(), + tb.getText())); } } else if (block instanceof ThinkingBlock tb) { blockLifecycle.startThinking(events); if (tb.getThinking() != null && !tb.getThinking().isEmpty()) { events.add( new ThinkingBlockDeltaEvent( - blockLifecycle.replyId, "thinking", tb.getThinking())); + blockLifecycle.replyId, + blockLifecycle.currentThinkingBlockId(), + tb.getThinking())); } } else if (withToolEvents && block instanceof ToolUseBlock tub) { String toolId = resolveToolCallId(tub, context); @@ -2624,13 +2629,17 @@ private void emitBlockEvents( * *

The model stream is consumed through {@code concatMap}, but the state holders keep the * previous thread-safe shape because model providers may deliver chunk content - * unpredictably. This helper only changes when pending end events are flushed; it does not - * change the block identity or event payloads. + * unpredictably. Each contiguous text or thinking segment receives its own block ID so its + * start, delta, and end events can be correlated independently. */ private final class ModelCallBlockLifecycle { private final String replyId; private final AtomicBoolean textStarted = new AtomicBoolean(false); + private final AtomicLong textSegmentSequence = new AtomicLong(0); + private final AtomicReference currentTextBlockId = new AtomicReference<>(); private final AtomicBoolean thinkingStarted = new AtomicBoolean(false); + private final AtomicLong thinkingSegmentSequence = new AtomicLong(0); + private final AtomicReference currentThinkingBlockId = new AtomicReference<>(); private final Map startedToolCalls = new ConcurrentHashMap<>(); private ModelCallBlockLifecycle(String replyId) { @@ -2640,16 +2649,30 @@ private ModelCallBlockLifecycle(String replyId) { private void startText(List events) { flushThinking(events); if (textStarted.compareAndSet(false, true)) { - events.add(new TextBlockStartEvent(replyId, "text")); + long segment = textSegmentSequence.incrementAndGet(); + String blockId = segment == 1 ? "text" : "text-" + segment; + currentTextBlockId.set(blockId); + events.add(new TextBlockStartEvent(replyId, blockId)); } } + private String currentTextBlockId() { + return currentTextBlockId.get(); + } + private void startThinking(List events) { if (thinkingStarted.compareAndSet(false, true)) { - events.add(new ThinkingBlockStartEvent(replyId, "thinking")); + long segment = thinkingSegmentSequence.incrementAndGet(); + String blockId = segment == 1 ? "thinking" : "thinking-" + segment; + currentThinkingBlockId.set(blockId); + events.add(new ThinkingBlockStartEvent(replyId, blockId)); } } + private String currentThinkingBlockId() { + return currentThinkingBlockId.get(); + } + private void startToolCall(String toolId, String toolName, List events) { if (toolId == null || startedToolCalls.containsKey(toolId)) { return; @@ -2665,13 +2688,15 @@ private void startToolCall(String toolId, String toolName, List even private void flushText(List events) { if (textStarted.compareAndSet(true, false)) { - events.add(new TextBlockEndEvent(replyId, "text")); + String blockId = currentTextBlockId.getAndSet(null); + events.add(new TextBlockEndEvent(replyId, blockId)); } } private void flushThinking(List events) { if (thinkingStarted.compareAndSet(true, false)) { - events.add(new ThinkingBlockEndEvent(replyId, "thinking")); + String blockId = currentThinkingBlockId.getAndSet(null); + events.add(new ThinkingBlockEndEvent(replyId, blockId)); } } @@ -3654,7 +3679,8 @@ private Flux summaryModelCallStream( new TextBlockDeltaEvent( blockLifecycle .replyId, - "text", + blockLifecycle + .currentTextBlockId(), tb.getText())); } } else if (block @@ -3668,7 +3694,8 @@ private Flux summaryModelCallStream( new ThinkingBlockDeltaEvent( blockLifecycle .replyId, - "thinking", + blockLifecycle + .currentThinkingBlockId(), tb .getThinking())); } diff --git a/agentscope-core/src/test/java/io/agentscope/core/agent/ReActAgentNewLoopReplyTest.java b/agentscope-core/src/test/java/io/agentscope/core/agent/ReActAgentNewLoopReplyTest.java index 5179068c0f..01d69d9a57 100644 --- a/agentscope-core/src/test/java/io/agentscope/core/agent/ReActAgentNewLoopReplyTest.java +++ b/agentscope-core/src/test/java/io/agentscope/core/agent/ReActAgentNewLoopReplyTest.java @@ -29,8 +29,10 @@ import io.agentscope.core.event.ModelCallEndEvent; import io.agentscope.core.event.ModelCallStartEvent; import io.agentscope.core.event.RequireExternalExecutionEvent; +import io.agentscope.core.event.TextBlockDeltaEvent; import io.agentscope.core.event.TextBlockEndEvent; import io.agentscope.core.event.TextBlockStartEvent; +import io.agentscope.core.event.ThinkingBlockDeltaEvent; import io.agentscope.core.event.ThinkingBlockEndEvent; import io.agentscope.core.event.ThinkingBlockStartEvent; import io.agentscope.core.event.ToolCallEndEvent; @@ -517,6 +519,122 @@ void consecutiveTextChunksEmitSingleStartAndEnd() { < indexOf(events, ModelCallEndEvent.class)); } + @Test + void textSeparatedByToolCallUsesDistinctBlockIds() { + ChatModelBase model = + new ScriptedModel( + List.of( + () -> + Flux.just( + chatResponse( + TextBlock.builder().text("before").build()), + chatResponse( + ToolUseBlock.builder() + .id("tc1") + .name("echo") + .input(Map.of("query", "ping")) + .build()), + chatResponse( + TextBlock.builder().text("after").build())), + () -> Flux.just(textResponse("done")))); + ReActAgent agent = + ReActAgent.builder() + .name("asst") + .model(model) + .toolkit(toolkitWith(new EchoTool())) + .build(); + + List events = agent.streamEvents(List.of()).collectList().block(); + assertNotNull(events); + + int firstModelEnd = indexOf(events, ModelCallEndEvent.class); + List starts = + events.subList(0, firstModelEnd).stream() + .filter(TextBlockStartEvent.class::isInstance) + .map(TextBlockStartEvent.class::cast) + .toList(); + List ends = + events.subList(0, firstModelEnd).stream() + .filter(TextBlockEndEvent.class::isInstance) + .map(TextBlockEndEvent.class::cast) + .toList(); + List deltas = + events.subList(0, firstModelEnd).stream() + .filter(TextBlockDeltaEvent.class::isInstance) + .map(TextBlockDeltaEvent.class::cast) + .toList(); + + assertEquals( + List.of("text", "text-2"), + starts.stream().map(TextBlockStartEvent::getBlockId).toList()); + assertEquals( + List.of("text", "text-2"), + deltas.stream().map(TextBlockDeltaEvent::getBlockId).toList()); + assertEquals( + List.of("text", "text-2"), + ends.stream().map(TextBlockEndEvent::getBlockId).toList()); + } + + @Test + void thinkingSeparatedByToolCallUsesDistinctBlockIds() { + ChatModelBase model = + new ScriptedModel( + List.of( + () -> + Flux.just( + chatResponse( + ThinkingBlock.builder() + .thinking("before") + .build()), + chatResponse( + ToolUseBlock.builder() + .id("tc1") + .name("echo") + .input(Map.of("query", "ping")) + .build()), + chatResponse( + ThinkingBlock.builder() + .thinking("after") + .build())), + () -> Flux.just(textResponse("done")))); + ReActAgent agent = + ReActAgent.builder() + .name("asst") + .model(model) + .toolkit(toolkitWith(new EchoTool())) + .build(); + + List events = agent.streamEvents(List.of()).collectList().block(); + assertNotNull(events); + + int firstModelEnd = indexOf(events, ModelCallEndEvent.class); + List starts = + events.subList(0, firstModelEnd).stream() + .filter(ThinkingBlockStartEvent.class::isInstance) + .map(ThinkingBlockStartEvent.class::cast) + .toList(); + List ends = + events.subList(0, firstModelEnd).stream() + .filter(ThinkingBlockEndEvent.class::isInstance) + .map(ThinkingBlockEndEvent.class::cast) + .toList(); + List deltas = + events.subList(0, firstModelEnd).stream() + .filter(ThinkingBlockDeltaEvent.class::isInstance) + .map(ThinkingBlockDeltaEvent.class::cast) + .toList(); + + assertEquals( + List.of("thinking", "thinking-2"), + starts.stream().map(ThinkingBlockStartEvent::getBlockId).toList()); + assertEquals( + List.of("thinking", "thinking-2"), + deltas.stream().map(ThinkingBlockDeltaEvent::getBlockId).toList()); + assertEquals( + List.of("thinking", "thinking-2"), + ends.stream().map(ThinkingBlockEndEvent::getBlockId).toList()); + } + @Test void summaryModelCallClosesThinkingBeforeTextAndFlushesTextBeforeModelEnd() { ChatModelBase model = diff --git a/agentscope-extensions/agentscope-extensions-protocol/agentscope-extensions-agui/src/main/java/io/agentscope/core/agui/adapter/strategy/AguiStreamContext.java b/agentscope-extensions/agentscope-extensions-protocol/agentscope-extensions-agui/src/main/java/io/agentscope/core/agui/adapter/strategy/AguiStreamContext.java index 11c041b279..086db9b6b0 100644 --- a/agentscope-extensions/agentscope-extensions-protocol/agentscope-extensions-agui/src/main/java/io/agentscope/core/agui/adapter/strategy/AguiStreamContext.java +++ b/agentscope-extensions/agentscope-extensions-protocol/agentscope-extensions-agui/src/main/java/io/agentscope/core/agui/adapter/strategy/AguiStreamContext.java @@ -38,10 +38,6 @@ public class AguiStreamContext { - // CopilotKit will merge reasoning and text with the same messageId, adding suffixes to the - // reasoning to distinguish them - public static final String REASONING_MESSAGE_ID_SUFFIX = "-reasoning"; - private static final Logger logger = LoggerFactory.getLogger(AguiStreamContext.class); private final String threadId; @@ -169,21 +165,16 @@ public void closeTextMessage(String messageId) { } public void startReasoningMessage(String messageId) { - String reasoningMessageId = reasoningMessageId(messageId); - if (startedReasoningMessages.add(reasoningMessageId)) { - emit( - new AguiEvent.ReasoningMessageStart( - threadId, runId, reasoningMessageId, "reasoning")); + if (startedReasoningMessages.add(messageId)) { + emit(new AguiEvent.ReasoningMessageStart(threadId, runId, messageId, "reasoning")); } - currentReasoningMessageId = reasoningMessageId; + currentReasoningMessageId = messageId; } public void appendReasoningDelta(String messageId, String delta) { if (delta != null && !delta.isEmpty()) { startReasoningMessage(messageId); - emit( - new AguiEvent.ReasoningMessageContent( - threadId, runId, reasoningMessageId(messageId), delta)); + emit(new AguiEvent.ReasoningMessageContent(threadId, runId, messageId, delta)); } } @@ -195,17 +186,16 @@ public void closeActiveReasoningMessage() { } public void closeReasoningMessage(String messageId) { - String reasoningMessageId = reasoningMessageId(messageId); - if (reasoningMessageId == null - || !startedReasoningMessages.contains(reasoningMessageId) - || endedReasoningMessages.contains(reasoningMessageId)) { + if (messageId == null + || !startedReasoningMessages.contains(messageId) + || endedReasoningMessages.contains(messageId)) { return; } - endedReasoningMessages.add(reasoningMessageId); - if (Objects.equals(reasoningMessageId, currentReasoningMessageId)) { + endedReasoningMessages.add(messageId); + if (Objects.equals(messageId, currentReasoningMessageId)) { currentReasoningMessageId = null; } - emit(new AguiEvent.ReasoningMessageEnd(threadId, runId, reasoningMessageId)); + emit(new AguiEvent.ReasoningMessageEnd(threadId, runId, messageId)); } public void startToolCall(String toolCallId, String toolCallName) { @@ -279,7 +269,6 @@ public void endToolResult(String replyId, String toolCallId) { if (endedToolCalls.add(toolCallId)) { emit(new AguiEvent.ToolCallEnd(threadId, runId, toolCallId)); } - StringBuilder content = toolResultContent.remove(toolCallId); emit( new AguiEvent.ToolCallResult( @@ -335,13 +324,6 @@ private static String normalizeToolCallName(String toolCallName) { return toolCallName != null && !toolCallName.isBlank() ? toolCallName : "unknown"; } - private static String reasoningMessageId(String messageId) { - if (messageId.endsWith(REASONING_MESSAGE_ID_SUFFIX)) { - return messageId; - } - return messageId + REASONING_MESSAGE_ID_SUFFIX; - } - private static String serialize(ContentBlock data) { if (data instanceof TextBlock textBlock) { return textBlock.getText(); @@ -398,7 +380,6 @@ private static Set frontendToolNames(RunAgentInput runInput) { } static final class TokenUsageAccumulator { - private long cumulativeInputTokens; private long cumulativeOutputTokens; private long cumulativeCachedTokens; @@ -426,7 +407,6 @@ TokenUsageSnapshot add(ChatUsage usage) { record TokenUsageSnapshot(TokenUsage delta, TokenUsage cumulative) {} record TokenUsage(long inputTokens, long outputTokens, long cachedTokens, double time) { - long totalTokens() { return inputTokens + outputTokens; } diff --git a/agentscope-extensions/agentscope-extensions-protocol/agentscope-extensions-agui/src/main/java/io/agentscope/core/agui/adapter/strategy/TextBlockEventConverter.java b/agentscope-extensions/agentscope-extensions-protocol/agentscope-extensions-agui/src/main/java/io/agentscope/core/agui/adapter/strategy/TextBlockEventConverter.java index 438c758c1f..49e1d5c8d9 100644 --- a/agentscope-extensions/agentscope-extensions-protocol/agentscope-extensions-agui/src/main/java/io/agentscope/core/agui/adapter/strategy/TextBlockEventConverter.java +++ b/agentscope-extensions/agentscope-extensions-protocol/agentscope-extensions-agui/src/main/java/io/agentscope/core/agui/adapter/strategy/TextBlockEventConverter.java @@ -33,9 +33,14 @@ public Set> eventTypes() { public void convert(AgentEvent event, AguiStreamContext context) { if (event instanceof TextBlockDeltaEvent delta) { // AguiEvent.TextMessageStart delays sending when content arrives - context.appendTextDelta(delta.getReplyId(), delta.getDelta()); + context.appendTextDelta( + messageId(delta.getReplyId(), delta.getBlockId()), delta.getDelta()); } else if (event instanceof TextBlockEndEvent end) { - context.closeTextMessage(end.getReplyId()); + context.closeTextMessage(messageId(end.getReplyId(), end.getBlockId())); } } + + private String messageId(String replyId, String blockId) { + return replyId + "-" + blockId; + } } diff --git a/agentscope-extensions/agentscope-extensions-protocol/agentscope-extensions-agui/src/main/java/io/agentscope/core/agui/adapter/strategy/ThinkingBlockEventConverter.java b/agentscope-extensions/agentscope-extensions-protocol/agentscope-extensions-agui/src/main/java/io/agentscope/core/agui/adapter/strategy/ThinkingBlockEventConverter.java index bdf7ef4132..0994d2e041 100644 --- a/agentscope-extensions/agentscope-extensions-protocol/agentscope-extensions-agui/src/main/java/io/agentscope/core/agui/adapter/strategy/ThinkingBlockEventConverter.java +++ b/agentscope-extensions/agentscope-extensions-protocol/agentscope-extensions-agui/src/main/java/io/agentscope/core/agui/adapter/strategy/ThinkingBlockEventConverter.java @@ -40,9 +40,14 @@ public void convert(AgentEvent event, AguiStreamContext context) { if (event instanceof ThinkingBlockDeltaEvent delta) { // AguiEvent.ReasoningMessageStart delays sending when content arrives - context.appendReasoningDelta(delta.getReplyId(), delta.getDelta()); + context.appendReasoningDelta( + messageId(delta.getReplyId(), delta.getBlockId()), delta.getDelta()); } else if (event instanceof ThinkingBlockEndEvent end) { - context.closeReasoningMessage(end.getReplyId()); + context.closeReasoningMessage(messageId(end.getReplyId(), end.getBlockId())); } } + + private String messageId(String replyId, String blockId) { + return replyId + "-" + blockId; + } } diff --git a/agentscope-extensions/agentscope-extensions-protocol/agentscope-extensions-agui/src/test/java/io/agentscope/core/agui/adapter/AguiAgentAdapterV2Test.java b/agentscope-extensions/agentscope-extensions-protocol/agentscope-extensions-agui/src/test/java/io/agentscope/core/agui/adapter/AguiAgentAdapterV2Test.java index 62d5a1dfe2..421a390b9f 100644 --- a/agentscope-extensions/agentscope-extensions-protocol/agentscope-extensions-agui/src/test/java/io/agentscope/core/agui/adapter/AguiAgentAdapterV2Test.java +++ b/agentscope-extensions/agentscope-extensions-protocol/agentscope-extensions-agui/src/test/java/io/agentscope/core/agui/adapter/AguiAgentAdapterV2Test.java @@ -329,6 +329,122 @@ void testTextBlockEventsConvertToAguiTextMessageEvents() { assertEquals("hel", firstDelta.delta()); } + @Test + void testTextSegmentsSeparatedByToolCallUseDistinctMessageIds() { + List events = + runReActEvents( + new TextBlockStartEvent("reply-mixed", "text"), + new TextBlockDeltaEvent("reply-mixed", "text", "before"), + new TextBlockEndEvent("reply-mixed", "text"), + new ToolCallStartEvent("reply-mixed", "tool-1", "lookup"), + new ToolCallEndEvent("reply-mixed", "tool-1", "lookup"), + new TextBlockStartEvent("reply-mixed", "text-2"), + new TextBlockDeltaEvent("reply-mixed", "text-2", "after"), + new TextBlockEndEvent("reply-mixed", "text-2")); + + assertEquals( + List.of( + AguiEventType.TEXT_MESSAGE_START, + AguiEventType.TEXT_MESSAGE_CONTENT, + AguiEventType.TEXT_MESSAGE_END, + AguiEventType.TOOL_CALL_START, + AguiEventType.TOOL_CALL_END, + AguiEventType.TEXT_MESSAGE_START, + AguiEventType.TEXT_MESSAGE_CONTENT, + AguiEventType.TEXT_MESSAGE_END), + types(events)); + + List messageIds = + events.stream() + .filter( + event -> + event instanceof AguiEvent.TextMessageStart + || event instanceof AguiEvent.TextMessageContent + || event instanceof AguiEvent.TextMessageEnd) + .map( + event -> { + if (event instanceof AguiEvent.TextMessageStart start) { + return start.messageId(); + } + if (event instanceof AguiEvent.TextMessageContent content) { + return content.messageId(); + } + return ((AguiEvent.TextMessageEnd) event).messageId(); + }) + .toList(); + assertEquals( + List.of( + "reply-mixed-text", + "reply-mixed-text", + "reply-mixed-text", + "reply-mixed-text-2", + "reply-mixed-text-2", + "reply-mixed-text-2"), + messageIds); + } + + @Test + void testReasoningSegmentsSeparatedByToolCallUseDistinctMessageIds() { + List events = + runReActEvents( + AguiAdapterConfig.builder().enableReasoning(true).build(), + new ThinkingBlockStartEvent("reply-mixed", "thinking"), + new ThinkingBlockDeltaEvent("reply-mixed", "thinking", "before"), + new ThinkingBlockEndEvent("reply-mixed", "thinking"), + new ToolCallStartEvent("reply-mixed", "tool-1", "lookup"), + new ToolCallEndEvent("reply-mixed", "tool-1", "lookup"), + new ThinkingBlockStartEvent("reply-mixed", "thinking-2"), + new ThinkingBlockDeltaEvent("reply-mixed", "thinking-2", "after"), + new ThinkingBlockEndEvent("reply-mixed", "thinking-2")); + + assertEquals( + List.of( + AguiEventType.REASONING_MESSAGE_START, + AguiEventType.REASONING_MESSAGE_CONTENT, + AguiEventType.REASONING_MESSAGE_END, + AguiEventType.TOOL_CALL_START, + AguiEventType.TOOL_CALL_END, + AguiEventType.REASONING_MESSAGE_START, + AguiEventType.REASONING_MESSAGE_CONTENT, + AguiEventType.REASONING_MESSAGE_END), + types(events)); + + List messageIds = + events.stream() + .filter( + event -> + event instanceof AguiEvent.ReasoningMessageStart + || event + instanceof + AguiEvent.ReasoningMessageContent + || event + instanceof + AguiEvent.ReasoningMessageEnd) + .map( + event -> { + if (event + instanceof AguiEvent.ReasoningMessageStart start) { + return start.messageId(); + } + if (event + instanceof + AguiEvent.ReasoningMessageContent content) { + return content.messageId(); + } + return ((AguiEvent.ReasoningMessageEnd) event).messageId(); + }) + .toList(); + assertEquals( + List.of( + "reply-mixed-thinking", + "reply-mixed-thinking", + "reply-mixed-thinking", + "reply-mixed-thinking-2", + "reply-mixed-thinking-2", + "reply-mixed-thinking-2"), + messageIds); + } + @Test void testThinkingEventsAreIgnoredWhenReasoningDisabled() { List events = @@ -347,9 +463,9 @@ void testThinkingEventsConvertWhenReasoningEnabled() { List events = runReActEvents( AguiAdapterConfig.builder().enableReasoning(true).build(), - new ThinkingBlockStartEvent("reply-thinking", "block-1"), - new ThinkingBlockDeltaEvent("reply-thinking", "block-1", "visible"), - new ThinkingBlockEndEvent("reply-thinking", "block-1")); + new ThinkingBlockStartEvent("reply-thinking", "thinking"), + new ThinkingBlockDeltaEvent("reply-thinking", "thinking", "visible"), + new ThinkingBlockEndEvent("reply-thinking", "thinking")); assertEquals( List.of( @@ -363,8 +479,7 @@ void testThinkingEventsConvertWhenReasoningEnabled() { assertInstanceOf(AguiEvent.ReasoningMessageContent.class, events.get(1)); AguiEvent.ReasoningMessageEnd end = assertInstanceOf(AguiEvent.ReasoningMessageEnd.class, events.get(2)); - String expectedMessageId = - "reply-thinking" + AguiStreamContext.REASONING_MESSAGE_ID_SUFFIX; + String expectedMessageId = "reply-thinking-thinking"; assertEquals(expectedMessageId, start.messageId()); assertEquals(expectedMessageId, content.messageId()); assertEquals(expectedMessageId, end.messageId()); @@ -375,10 +490,10 @@ void testTextAndReasoningUseDifferentMessageIdsForSameReply() { List events = runReActEvents( AguiAdapterConfig.builder().enableReasoning(true).build(), - new ThinkingBlockDeltaEvent("reply-shared", "thinking-1", "think"), - new ThinkingBlockEndEvent("reply-shared", "thinking-1"), - new TextBlockDeltaEvent("reply-shared", "text-1", "answer"), - new TextBlockEndEvent("reply-shared", "text-1")); + new ThinkingBlockDeltaEvent("reply-shared", "thinking", "think"), + new ThinkingBlockEndEvent("reply-shared", "thinking"), + new TextBlockDeltaEvent("reply-shared", "text", "answer"), + new TextBlockEndEvent("reply-shared", "text")); AguiEvent.ReasoningMessageContent reasoningContent = events.stream() @@ -393,10 +508,8 @@ void testTextAndReasoningUseDifferentMessageIdsForSameReply() { .findFirst() .orElseThrow(); - assertEquals("reply-shared", textContent.messageId()); - assertEquals( - "reply-shared" + AguiStreamContext.REASONING_MESSAGE_ID_SUFFIX, - reasoningContent.messageId()); + assertEquals("reply-shared-text", textContent.messageId()); + assertEquals("reply-shared-thinking", reasoningContent.messageId()); } @Test