Skip to content

fix(agui): use unique message IDs for discontinuous text segments #3010

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
ZHCHJ888 wants to merge 3 commits into
base: main
Choose a base branch
from
Open

fix(agui): use unique message IDs for discontinuous text segments #3010

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
3 commits
Select commit Hold shift + click to select a range
c629078
fix(agui): use unique ids for text segments
ZHCHJ888 Sep 6, 2026
31bbf1e
fix(agui): use unique ids for reasoning segments
ZHCHJ888 Sep 6, 2026
656495d
Merge branch 'main' into fix/agui-text-segment-ids
jujn Sep 7, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
47 changes: 37 additions & 10 deletions agentscope-core/src/main/java/io/agentscope/core/ReActAgent.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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);
Expand All @@ -2624,13 +2629,17 @@ private void emitBlockEvents(
*
* <p>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<String> currentTextBlockId = new AtomicReference<>();
private final AtomicBoolean thinkingStarted = new AtomicBoolean(false);
private final AtomicLong thinkingSegmentSequence = new AtomicLong(0);
private final AtomicReference<String> currentThinkingBlockId = new AtomicReference<>();
private final Map<String, String> startedToolCalls = new ConcurrentHashMap<>();

private ModelCallBlockLifecycle(String replyId) {
Expand All @@ -2640,16 +2649,30 @@ private ModelCallBlockLifecycle(String replyId) {
private void startText(List<AgentEvent> 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<AgentEvent> 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<AgentEvent> events) {
if (toolId == null || startedToolCalls.containsKey(toolId)) {
return;
Expand All @@ -2665,13 +2688,15 @@ private void startToolCall(String toolId, String toolName, List<AgentEvent> even

private void flushText(List<AgentEvent> 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<AgentEvent> events) {
if (thinkingStarted.compareAndSet(true, false)) {
events.add(new ThinkingBlockEndEvent(replyId, "thinking"));
String blockId = currentThinkingBlockId.getAndSet(null);
events.add(new ThinkingBlockEndEvent(replyId, blockId));
}
}

Expand Down Expand Up @@ -3654,7 +3679,8 @@ private Flux<AgentEvent> summaryModelCallStream(
new TextBlockDeltaEvent(
blockLifecycle
.replyId,
"text",
blockLifecycle
.currentTextBlockId(),
tb.getText()));
}
} else if (block
Expand All @@ -3668,7 +3694,8 @@ private Flux<AgentEvent> summaryModelCallStream(
new ThinkingBlockDeltaEvent(
blockLifecycle
.replyId,
"thinking",
blockLifecycle
.currentThinkingBlockId(),
tb
.getThinking()));
}
Expand Down
118 changes: 118 additions & 0 deletions agentscope-core/src/test/java/io/agentscope/core/agent/ReActAgentNewLoopReplyTest.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand Down Expand Up @@ -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<AgentEvent> events = agent.streamEvents(List.of()).collectList().block();
assertNotNull(events);

int firstModelEnd = indexOf(events, ModelCallEndEvent.class);
List<TextBlockStartEvent> starts =
events.subList(0, firstModelEnd).stream()
.filter(TextBlockStartEvent.class::isInstance)
.map(TextBlockStartEvent.class::cast)
.toList();
List<TextBlockEndEvent> ends =
events.subList(0, firstModelEnd).stream()
.filter(TextBlockEndEvent.class::isInstance)
.map(TextBlockEndEvent.class::cast)
.toList();
List<TextBlockDeltaEvent> 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<AgentEvent> events = agent.streamEvents(List.of()).collectList().block();
assertNotNull(events);

int firstModelEnd = indexOf(events, ModelCallEndEvent.class);
List<ThinkingBlockStartEvent> starts =
events.subList(0, firstModelEnd).stream()
.filter(ThinkingBlockStartEvent.class::isInstance)
.map(ThinkingBlockStartEvent.class::cast)
.toList();
List<ThinkingBlockEndEvent> ends =
events.subList(0, firstModelEnd).stream()
.filter(ThinkingBlockEndEvent.class::isInstance)
.map(ThinkingBlockEndEvent.class::cast)
.toList();
List<ThinkingBlockDeltaEvent> 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 =
Expand Down
65 changes: 61 additions & 4 deletions ...nsions-agui/src/main/java/io/agentscope/core/agui/adapter/strategy/AguiStreamContext.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -53,8 +53,12 @@ public class AguiStreamContext {

private final Set<String> startedTextMessages = new LinkedHashSet<>();
private final Set<String> endedTextMessages = new LinkedHashSet<>();
private final Map<TextBlockKey, String> textMessageIds = new LinkedHashMap<>();
private final Map<String, String> firstTextBlockByReply = new LinkedHashMap<>();
private final Set<String> startedReasoningMessages = new LinkedHashSet<>();
private final Set<String> endedReasoningMessages = new LinkedHashSet<>();
private final Map<ThinkingBlockKey, String> reasoningMessageIds = new LinkedHashMap<>();
private final Map<String, String> firstThinkingBlockByReply = new LinkedHashMap<>();
private final Set<String> startedToolCalls = new LinkedHashSet<>();
private final Set<String> endedToolCalls = new LinkedHashSet<>();
private String currentTextMessageId;
Expand Down Expand Up @@ -148,6 +152,26 @@ public void appendTextDelta(String messageId, String delta) {
}
}

public String textMessageId(String replyId, String blockId) {
String normalizedBlockId = isBlank(blockId) ? "text" : blockId;
TextBlockKey key = new TextBlockKey(replyId, normalizedBlockId);
return textMessageIds.computeIfAbsent(
key,
ignored -> {
String firstBlockId =
firstTextBlockByReply.putIfAbsent(replyId, normalizedBlockId);
if (firstBlockId == null || Objects.equals(firstBlockId, normalizedBlockId)) {
return replyId;
}
return replyId + "-" + normalizedBlockId;
});
}

public String existingTextMessageId(String replyId, String blockId) {
String normalizedBlockId = isBlank(blockId) ? "text" : blockId;
return textMessageIds.get(new TextBlockKey(replyId, normalizedBlockId));
}

public void closeActiveTextMessage() {
if (currentTextMessageId == null) {
return;
Expand All @@ -169,7 +193,7 @@ public void closeTextMessage(String messageId) {
}

public void startReasoningMessage(String messageId) {
String reasoningMessageId = reasoningMessageId(messageId);
String reasoningMessageId = withReasoningSuffix(messageId);
if (startedReasoningMessages.add(reasoningMessageId)) {
emit(
new AguiEvent.ReasoningMessageStart(
Expand All @@ -183,10 +207,36 @@ public void appendReasoningDelta(String messageId, String delta) {
startReasoningMessage(messageId);
emit(
new AguiEvent.ReasoningMessageContent(
threadId, runId, reasoningMessageId(messageId), delta));
threadId, runId, withReasoningSuffix(messageId), delta));
}
}

public String reasoningMessageId(String replyId, String blockId) {
String normalizedBlockId = isBlank(blockId) ? "thinking" : blockId;
ThinkingBlockKey key = new ThinkingBlockKey(replyId, normalizedBlockId);
String messageId =
reasoningMessageIds.computeIfAbsent(
key,
ignored -> {
String firstBlockId =
firstThinkingBlockByReply.putIfAbsent(
replyId, normalizedBlockId);
if (firstBlockId == null
|| Objects.equals(firstBlockId, normalizedBlockId)) {
return replyId;
}
return replyId + "-" + normalizedBlockId;
});
return withReasoningSuffix(messageId);
}

public String existingReasoningMessageId(String replyId, String blockId) {
String normalizedBlockId = isBlank(blockId) ? "thinking" : blockId;
String messageId =
reasoningMessageIds.get(new ThinkingBlockKey(replyId, normalizedBlockId));
return messageId == null ? null : withReasoningSuffix(messageId);
}

public void closeActiveReasoningMessage() {
if (currentReasoningMessageId == null) {
return;
Expand All @@ -195,7 +245,7 @@ public void closeActiveReasoningMessage() {
}

public void closeReasoningMessage(String messageId) {
String reasoningMessageId = reasoningMessageId(messageId);
String reasoningMessageId = withReasoningSuffix(messageId);
if (reasoningMessageId == null
|| !startedReasoningMessages.contains(reasoningMessageId)
|| endedReasoningMessages.contains(reasoningMessageId)) {
Expand Down Expand Up @@ -335,7 +385,10 @@ private static String normalizeToolCallName(String toolCallName) {
return toolCallName != null && !toolCallName.isBlank() ? toolCallName : "unknown";
}

private static String reasoningMessageId(String messageId) {
private static String withReasoningSuffix(String messageId) {
if (messageId == null) {
return null;
}
if (messageId.endsWith(REASONING_MESSAGE_ID_SUFFIX)) {
return messageId;
}
Expand All @@ -357,6 +410,10 @@ private static boolean isBlank(String value) {
return value == null || value.isBlank();
}

private record TextBlockKey(String replyId, String blockId) {}

private record ThinkingBlockKey(String replyId, String blockId) {}

private void warnMissingToolCallId(String eventName) {
if (!warnedMissingToolCallIdOperations.add(eventName)) {
return;
Expand Down
Loading
Loading