Skip to content

fix(agui): reconcile replayed tool results with session state #3046

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

Closed
guslegend0510 wants to merge 18 commits into from
Closed

fix(agui): reconcile replayed tool results with session state #3046

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
18 commits
Select commit Hold shift + click to select a range
8514975
fix(agui): reconcile replayed tool results with session state
guslegend Sep 8, 2026
ddbae63
Merge remote-tracking branch 'origin/main' into HEAD
guslegend Sep 8, 2026
693dd82
Merge remote-tracking branch 'origin/main' into HEAD
guslegend Sep 9, 2026
ce138aa
Merge remote-tracking branch 'origin/main' into HEAD
guslegend Sep 9, 2026
4d6b9b4
Merge remote-tracking branch 'origin/main' into HEAD
guslegend Sep 9, 2026
1ffa26b
Merge remote-tracking branch 'origin/main' into HEAD
guslegend Sep 10, 2026
2155611
Merge remote-tracking branch 'origin/main' into HEAD
guslegend Sep 10, 2026
c2b5272
Merge remote-tracking branch 'origin/main' into HEAD
guslegend Sep 10, 2026
d3f9724
Merge remote-tracking branch 'origin/main' into HEAD
guslegend Sep 10, 2026
5829957
Merge remote-tracking branch 'origin/main' into HEAD
guslegend Sep 11, 2026
2c0d409
Merge remote-tracking branch 'origin/main' into HEAD
guslegend Sep 12, 2026
f20857d
Merge remote-tracking branch 'origin/main' into HEAD
guslegend Sep 12, 2026
5f862d0
Merge remote-tracking branch 'origin/main' into HEAD
guslegend Sep 12, 2026
cff2db9
Merge remote-tracking branch 'origin/main' into HEAD
guslegend Sep 13, 2026
d8652fd
fix(core): preserve replay message context
guslegend Sep 14, 2026
e16a4da
Merge remote-tracking branch 'origin/main' into HEAD
guslegend Sep 14, 2026
d26d43f
fix(core): tighten replay result filtering
guslegend Sep 14, 2026
347c304
Merge remote-tracking branch 'origin/main' into HEAD
guslegend Sep 14, 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
156 changes: 149 additions & 7 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 @@ -757,6 +757,12 @@ protected Object beforeAgentExecution(List<Msg> msgs, RuntimeContext rc) {
return scope;
}

@Override
protected List<Msg> prepareInputMessages(List<Msg> msgs, Object callScope) {
CallExecution scope = (CallExecution) callScope;

@oss-maintainer oss-maintainer Sep 11, 2026

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

prepareInputMessages does an unchecked cast (CallExecution) callScope and then calls selectReplayedMessages. Two things worth hardening: (1) if callScope is ever null or a different type (a subclass of ReActAgent overriding beforeAgentExecution, or a future hook that passes a different scope) this becomes a raw ClassCastException/NPE inside the call path instead of a clear error; (2) the base-class contract for prepareInputMessages is not documented here. An instanceof check that returns msgs unchanged on mismatch (or making the hook signature generic/typed) would keep the replay optimisation strictly best-effort and non-breaking for subclasses.

@oss-maintainer oss-maintainer Sep 13, 2026

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

prepareInputMessages casts callScope to CallExecution without a null check. If beforeAgentExecution ever returns null (e.g., an early-exit subclass), this will NPE before the existing null-safe path in runLifecycleBody can handle it. Add a null guard and return msgs unchanged.

@oss-maintainer oss-maintainer Sep 19, 2026

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

The callScope parameter arrives as Object from AgentBase.runLifecycleBody, and the previous review round already asked for this cast to be made safe. It is still unguarded: if beforeAgentExecution ever returns null (or a non-CallExecution scope, e.g. for a subclass/mixin that overrides beforeAgentExecution), this throws ClassCastException/NullPointerException from inside prepareInputMessages — before tracing and pre-call hooks — so the failure surfaces as an opaque run error rather than "replay handling not applicable". A cheap guard keeps the hook total:

if (!(callScope instanceof CallExecution scope)) {
    return msgs;
}
return scope.selectReplayedMessages(msgs);

(plus a msgs == null short-circuit if the base class can pass null).

return scope.selectReplayedMessages(msgs);
}

@Override
protected Mono<Msg> seedSystemMsg(Object callExectution) {
RuntimeContext rc =
Expand Down Expand Up @@ -1784,20 +1790,24 @@ private Mono<Msg> doCallInner(List<Msg> msgs) {

Set<String> pendingIds = getPendingToolUseIds();

// No pending tools -> normal processing
if (pendingIds.isEmpty()) {
addToContext(msgs);
return coreAgent();
}

// Permission HITL: if any pending tool is ASKING, the caller MUST supply
// ConfirmResults (via Msg.METADATA_CONFIRM_RESULTS) before we can proceed.
List<ToolUseBlock> asking = askingToolCalls();
List<ToolUseBlock> asking = pendingIds.isEmpty() ? List.of() : askingToolCalls();
if (!asking.isEmpty()) {
validateAndAcceptConfirmResults(msgs, asking);
return resumeAgent();
}

// Transcript-replaying clients may resend results already consumed by this session.
// Normalize against the call-scoped state before either recovery or normal input
// handling, including calls with no pending tools. Unknown IDs remain invalid.
msgs = removeConsumedToolResults(msgs, pendingIds);

@oss-maintainer oss-maintainer Sep 19, 2026

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

removeConsumedToolResults is applied to every call, not just replays (only selectReplayedMessages is gated on RuntimeContext.REPLAYED_INPUT). For sessions with pendingIds empty, a client that echoes the last tool result together with its follow-up message — legal under the old addToContext path — now hits the Invalid tool result IDs throw at line 2185. The docs update covers it, and unknownTrailingResultFailsWithoutPoisoningTheNextCorrectedRequest pins the behaviour, but this is a behavior change for existing non-replay callers: would it be worth an explicit @apiNote on the ReActAgent class / a CHANGELOG entry mentioning the stricter contract (and that REPLAYED_INPUT is the opt-out path)?


if (pendingIds.isEmpty()) {
addToContext(msgs);
return coreAgent();
}

// Pending-tool-call recovery: auto-patch orphaned pending tool calls with synthetic
// error results so the agent can continue instead of crashing. This must happen after
// the permission HITL flow so ASKING tool calls are handled by confirmation first.
Expand Down Expand Up @@ -2077,6 +2087,138 @@ private void applyToolUseBlockReplacements(Map<String, ToolUseBlock> replacement
}
}

private List<Msg> selectReplayedMessages(List<Msg> msgs) {
if (rc == null
|| !Boolean.TRUE.equals(rc.get(RuntimeContext.REPLAYED_INPUT))
|| msgs == null
|| msgs.isEmpty()) {
return msgs;
}
int lastAssistant = -1;
for (int i = msgs.size() - 1; i >= 0; i--) {
if (msgs.get(i).getRole() == MsgRole.ASSISTANT) {
lastAssistant = i;
break;
}
}
if (lastAssistant < 0) {

@oss-maintainer oss-maintainer Sep 11, 2026

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

selected.addAll(msgs.subList(lastAssistant + 1, msgs.size())) keeps everything after the last assistant message, and in the selected.isEmpty() fallback the single last USER message is re-appended. Please confirm the ordering guarantees for the case where a tool result for a pending call appears after the last assistant message in the client transcript — with the fallback path, a mid-transcript pending result would be dropped and only the old user prompt replayed, which is exactly the failure mode the PR sets out to fix. If that transcript shape is possible with the AG-UI client, a scan over the whole transcript (not only indices <= lastAssistant) would be more robust; if it is guaranteed impossible, it would help to state that as an assumption in the javadoc of RuntimeContext.REPLAYED_INPUT.

return msgs;
}

Set<String> pendingIds = getPendingToolUseIds();
List<Msg> selected = new ArrayList<>();
// Clients can split one server tool round into several local assistant turns. Do
// not let those turns hide a result for a call the server is still waiting for.
for (int i = 0; i <= lastAssistant; i++) {
Msg msg = msgs.get(i);
List<ContentBlock> replayContent = new ArrayList<>();
boolean hasPendingResult = false;
for (ContentBlock block : msg.getContent()) {
if (block instanceof ToolResultBlock result) {
if (pendingIds.contains(result.getId())) {
replayContent.add(block);
hasPendingResult = true;
}
} else {
replayContent.add(block);
}
}
if (hasPendingResult) {
selected.add(msg.withContent(replayContent));
}
}
selected.addAll(msgs.subList(lastAssistant + 1, msgs.size()));
// Preserve regenerate/continue behavior, but never replay the old user prompt in
// place of a pending tool result found in the middle of the transcript.

@oss-maintainer oss-maintainer Sep 13, 2026

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

selectReplayedMessages returns an empty list for a replay that contains only an assistant message with no preceding user message. Consider whether this should fail fast or fall back to the latest user message even when it appears after the assistant.

if (selected.isEmpty()) {

@oss-maintainer oss-maintainer Sep 19, 2026

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This fallback can re-append an already-answered user prompt. selected is empty when (a) no message up to lastAssistant carries a result for a currently pending call, and (b) msgs has nothing after lastAssistant — i.e. a replayed transcript that ends on the client's assistant turn with no pending tools. In that case the USER message at lastAssistant - 1 is added back and removeConsumedToolResults passes it through untouched (it only filters ToolResultBlocks), so the model is asked to answer the original instruction a second time. The comment above ("never replay the old user prompt") reads as if this path were excluded, and the no-pending-results case is exactly when the guarantee is not held.

For a "regenerate" request this is intended; for a stale/duplicate replay of a completed round it is a user-visible double execution (extra tool calls, extra model spend). Could you either confirm the intent and add a case for "REPLAYED_INPUT + transcript ends on assistant + pendingIds empty", or return the input unfiltered / reject when there is nothing to reconcile?

for (int i = lastAssistant - 1; i >= 0; i--) {
if (msgs.get(i).getRole() == MsgRole.USER) {
selected.add(msgs.get(i));
break;
}
}
}
return selected;
}

private List<Msg> removeConsumedToolResults(List<Msg> msgs, Set<String> pendingIds) {
if (msgs == null
|| msgs.stream().noneMatch(m -> m.hasContentBlocks(ToolResultBlock.class))) {
return msgs;
}

Set<String> consumedIds =
state.contextMutable().stream()
.flatMap(m -> m.getContentBlocks(ToolResultBlock.class).stream())
.map(ToolResultBlock::getId)
.collect(Collectors.toSet());
Set<String> historicalToolIds =
state.contextMutable().stream()
.filter(m -> m.getRole() == MsgRole.ASSISTANT)
.flatMap(m -> m.getContentBlocks(ToolUseBlock.class).stream())
.map(ToolUseBlock::getId)
.collect(Collectors.toSet());
consumedIds.retainAll(historicalToolIds);
// A stateless caller can supply complete tool-call/result history in its input.
boolean initialTranscript = state.contextMutable().isEmpty();
Set<String> inputToolIds = new HashSet<>();
Set<String> providedIds = new HashSet<>();
Set<String> invalidIds = new HashSet<>();
for (Msg msg : msgs) {
if (initialTranscript && msg.getRole() == MsgRole.ASSISTANT) {
msg.getContentBlocks(ToolUseBlock.class).stream()
.map(ToolUseBlock::getId)
.forEach(inputToolIds::add);
}
for (ToolResultBlock result : msg.getContentBlocks(ToolResultBlock.class)) {
String id = result.getId();
if (!providedIds.add(id)) {
throw new IllegalStateException("Duplicate tool result ID: " + id);
}
if (!pendingIds.contains(id)
&& !consumedIds.contains(id)
&& !inputToolIds.contains(id)) {
invalidIds.add(id);
}
}
}
if (!invalidIds.isEmpty()) {

@oss-maintainer oss-maintainer Sep 19, 2026

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Minor typing note for the three new rejection paths here (Duplicate tool result ID, Invalid tool result IDs, Input contains only previously consumed tool results): these are all caller-input problems, but they are raised as IllegalStateException, which is also what ObjectStoreTranscriptStore.appendSegment now uses for an I/O failure. Through the AG-UI path both surface as RUN_ERROR (unknownTrailingResultFailsWithoutPoisoningTheNextCorrectedRequest pins that), so nothing is broken today — but a dedicated validation exception type (or an IllegalArgumentException for input-shaped failures) would let transports map 4xx-vs-5xx and would keep operator-facing "write failed" errors distinguishable from "client sent a bad transcript" in logs.

throw new IllegalStateException(
"Invalid tool result IDs: " + invalidIds + ". Expected: " + pendingIds);
}

Set<String> replayedIds = new HashSet<>(providedIds);
replayedIds.retainAll(consumedIds);
if (!replayedIds.isEmpty()) {
log.debug("Ignoring previously consumed tool result IDs: {}", replayedIds);
}

List<Msg> cleaned = new ArrayList<>();
for (Msg msg : msgs) {
List<ContentBlock> content =
msg.getContent().stream()
.filter(
block ->
!(block instanceof ToolResultBlock result)
|| !consumedIds.contains(result.getId()))
.toList();
if (content.size() == msg.getContent().size()) {
cleaned.add(msg);
} else if (!content.isEmpty()
|| (msg.getMetadata() != null
&& msg.getMetadata().containsKey(Msg.METADATA_CONFIRM_RESULTS))) {
cleaned.add(msg.withContent(content));
}
}
if (cleaned.isEmpty()) {
throw new IllegalStateException(
"Input contains only previously consumed tool results. Provide new"
+ " messages or results for pending IDs: "
+ pendingIds);
}
return cleaned;
}

private void maybePatchPendingToolCalls(List<Msg> msgs, Set<String> pendingIds) {
if (pendingIds.isEmpty()) {
return;
Expand Down
20 changes: 17 additions & 3 deletions agentscope-core/src/main/java/io/agentscope/core/agent/AgentBase.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -307,18 +307,19 @@ private Mono<Msg> runLifecycleBody(
// shutdown interrupts / saves the exact (userId, sessionId) session rather than the agent's
// no-arg "most-recently-active" accessors.
GracefulShutdownManager.getInstance().bindRequestState(requestId, stateForCall(scope));
List<Msg> preparedMsgs = prepareInputMessages(msgs, scope);
Mono<Msg> body =

@oss-maintainer oss-maintainer Sep 13, 2026

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

runLifecycleBody now passes preparedMsgs to TracerRegistry.callAgent, notifyPreCall, and the error handler instead of the original caller input. Verify that tracers/pre-call hooks are expected to observe the filtered/replayed messages rather than the raw input.

TracerRegistry.get()
.callAgent(
this,
msgs,
preparedMsgs,
() ->
notifyPreCall(msgs, scope)
notifyPreCall(preparedMsgs, scope)
.flatMap(doCallFn)
.flatMap(this::notifyPostCall)
.onErrorResume(
createErrorHandler(
msgs.toArray(new Msg[0]))));
preparedMsgs.toArray(new Msg[0]))));
return scope == null ? body : body.contextWrite(c -> c.put(CALL_SCOPE_KEY, scope));
}

Expand All @@ -336,6 +337,19 @@ protected Object callSerializationKey(RuntimeContext rc) {
return null;
}

/**
* Prepares caller input after session activation and before tracing and pre-call hooks.
* Subclasses can reconcile replayed input against the call-scoped state while holding the
* session serialization gate. The default preserves the original input.
*
* @param msgs caller input
* @param callScope the scope returned by {@code beforeAgentExecution}
* @return the messages to process
*/
protected List<Msg> prepareInputMessages(List<Msg> msgs, Object callScope) {
return msgs;
}

/**
* Serializes {@code action} against other actions sharing {@code key}: this call waits for the
* previously-enqueued call with the same key to terminate before running, then becomes the tail
Expand Down
9 changes: 9 additions & 0 deletions agentscope-core/src/main/java/io/agentscope/core/agent/RuntimeContext.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -32,6 +32,15 @@
*/
public class RuntimeContext {

/**
* Boolean attribute indicating that input contains a replayed client transcript while the
* server owns conversation history. ReActAgent selects follow-up messages after the last
* client assistant turn, plus results for currently pending calls anywhere in the transcript.
* Selection happens after the authoritative session state is loaded. Omit this attribute for
* incremental input or stateless, client-owned history.
*/
public static final String REPLAYED_INPUT = "agentscope_replayed_input";

@oss-maintainer oss-maintainer Sep 11, 2026

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Nit on the public API surface: this introduces a new cross-module contract (agentscope-extensions-agui writes the flag, agentscope-core reads it) as a raw string-keyed boolean attribute. Since it is a first-class behavioural switch rather than user data, consider a typed accessor on RuntimeContext (e.g. isReplayedInput() / builder().replayedInput(true)) so the flag cannot be misspelled or set to a non-Boolean value, and so it shows up in the typed-attribute API rather than in the free-form get(String) bag.

private static final String TYPED_DEFAULT_KEY = "";

private final String sessionId;
Expand Down
Loading
Loading