feat(context-enhance): add RuntimeContext onStateLoaded callback for anchor-based AG-UI message auto-merge

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

@@ -146,6 +146,7 @@
 import java.util.concurrent.atomic.AtomicBoolean;
 import java.util.concurrent.atomic.AtomicLong;
 import java.util.concurrent.atomic.AtomicReference;
+import java.util.function.BiConsumer;
 import java.util.function.Function;
 import java.util.stream.Collectors;
 import java.util.stream.IntStream;
@@ -717,6 +718,11 @@ protected Object beforeAgentExecution(List<Msg> msgs, RuntimeContext rc) {
         // the active session's state via rc.getAgentState() (call-scoped, concurrency-safe)
         // rather than agent.getAgentState() (not call-scoped under concurrency).
         ctx.setAgentState(scope.state);
+        // state 加载完成，触发 onStateLoaded 回调
+        BiConsumer<RuntimeContext, List<Msg>> onStateLoaded = ctx.getOnStateLoaded();
+        if (onStateLoaded != null) {
+            onStateLoaded.accept(ctx, msgs);
+        }
         this.activeRc = ctx;
         bindRuntimeContextToHooks(ctx);
         // Seed per-call state onto the active execution scope. The system message is initialised

agentscope-core/src/main/java/io/agentscope/core/agent/RuntimeContext.java

@@ -15,13 +15,16 @@
  */
 package io.agentscope.core.agent;
 
+import io.agentscope.core.message.Msg;
 import io.agentscope.core.state.AgentState;
 import io.agentscope.core.tool.ContextStore;
 import io.agentscope.core.tool.ToolExecutionContext;
 import java.util.HashMap;
+import java.util.List;
 import java.util.Map;
 import java.util.concurrent.ConcurrentHashMap;
 import java.util.concurrent.ConcurrentMap;
+import java.util.function.BiConsumer;
 
 /**
  * Per-call metadata for an agent run: session-scoped fields plus a thread-safe attribute bag and
@@ -45,6 +48,14 @@ public class RuntimeContext {
      */
     private volatile AgentState agentState;
 
+    /**
+     * Callback fired after {@link #agentState} is loaded and set on this context (inside the
+     * agent's {@code beforeAgentExecution}, right after {@link #setAgentState(AgentState)}). The
+     * callback receives this RuntimeContext and the mutable incoming message list — it may modify
+     * either in place. {@code null} when no callback is registered.
+     */
+    private volatile BiConsumer<RuntimeContext, List<Msg>> onStateLoaded;
+
     /** String-keyed extras (legacy and generic extension). */
     private final ConcurrentMap<String, Object> stringAttributes;
 
@@ -63,6 +74,7 @@ private RuntimeContext(Builder builder) {
         this.typedAttributes = new ConcurrentHashMap<>();
         this.toolExecutionContext = builder.toolExecutionContext;
         this.agentState = builder.agentState;
+        this.onStateLoaded = builder.onStateLoaded;
         if (builder.stringExtras != null) {
             this.stringAttributes.putAll(builder.stringExtras);
         }
@@ -115,6 +127,25 @@ public void setAgentState(AgentState agentState) {
         this.agentState = agentState;
     }
 
+    /**
+     * Returns the callback fired after AgentState is loaded and set on this context, or {@code null}
+     * if none is registered.
+     */
+    public BiConsumer<RuntimeContext, List<Msg>> getOnStateLoaded() {
+        return onStateLoaded;
+    }
+
+    /**
+     * Registers a callback fired after AgentState is loaded (inside {@code beforeAgentExecution},
+     * right after {@link #setAgentState(AgentState)}). The callback receives this RuntimeContext and
+     * the mutable incoming message list — it may modify either in place.
+     *
+     * @param onStateLoaded the callback, or {@code null} to clear
+     */
+    public void setOnStateLoaded(BiConsumer<RuntimeContext, List<Msg>> onStateLoaded) {
+        this.onStateLoaded = onStateLoaded;
+    }
+
     /**
      * Resolves the live {@link AgentState} for the current call, preferring the call-scoped state
      * carried on {@code ctx} (concurrency-safe) and falling back to {@code fallbackAgent}'s state
@@ -329,6 +360,7 @@ public static class Builder {
         private final Map<Class<?>, Map<String, Object>> typedValues = new HashMap<>();
         private ToolExecutionContext toolExecutionContext;
         private AgentState agentState;
+        private BiConsumer<RuntimeContext, List<Msg>> onStateLoaded;
 
         public Builder sessionId(String sessionId) {
             this.sessionId = sessionId;
@@ -345,6 +377,17 @@ public Builder agentState(AgentState agentState) {
             return this;
         }
 
+        /**
+         * Registers the {@link RuntimeContext#getOnStateLoaded()} callback on the built context.
+         *
+         * @param onStateLoaded the callback, or {@code null} to leave unset
+         * @return this builder
+         */
+        public Builder onStateLoaded(BiConsumer<RuntimeContext, List<Msg>> onStateLoaded) {
+            this.onStateLoaded = onStateLoaded;
+            return this;
+        }
+
         public Builder put(String key, Object value) {
             if (this.stringExtras == null) {
                 this.stringExtras = new ConcurrentHashMap<>();
@@ -383,6 +426,7 @@ public Builder from(RuntimeContext source) {
             this.sessionId = source.sessionId;
             this.userId = source.userId;
             this.agentState = source.agentState;
+            this.onStateLoaded = source.onStateLoaded;
             this.toolExecutionContext = source.toolExecutionContext;
             if (!source.stringAttributes.isEmpty()) {
                 this.stringExtras = new ConcurrentHashMap<>(source.stringAttributes);

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

@@ -19,10 +19,18 @@
 import static org.junit.jupiter.api.Assertions.assertInstanceOf;
 import static org.junit.jupiter.api.Assertions.assertNull;
 import static org.junit.jupiter.api.Assertions.assertSame;
+import static org.junit.jupiter.api.Assertions.assertTrue;
 
+import io.agentscope.core.message.Msg;
+import io.agentscope.core.message.UserMessage;
+import io.agentscope.core.state.AgentState;
 import io.agentscope.core.tool.ToolExecutionContext;
+import java.util.ArrayList;
+import java.util.List;
 import java.util.concurrent.CyclicBarrier;
 import java.util.concurrent.atomic.AtomicInteger;
+import java.util.concurrent.atomic.AtomicReference;
+import java.util.function.BiConsumer;
 import org.junit.jupiter.api.DisplayName;
 import org.junit.jupiter.api.Test;
 
@@ -157,6 +165,91 @@ void builderCopyHandlesNullSource() {
         assertNull(empty.get("missing", Marker.class));
     }
 
+    @Test
+    @DisplayName("onStateLoaded defaults to null")
+    void onStateLoaded_defaultsToNull() {
+        RuntimeContext ctx = RuntimeContext.empty();
+        assertNull(ctx.getOnStateLoaded());
+    }
+
+    @Test
+    @DisplayName("builder.onStateLoaded propagates the callback to the built context")
+    void onStateLoaded_builderPropagatesCallback() {
+        BiConsumer<RuntimeContext, List<Msg>> callback = (c, m) -> {};
+        RuntimeContext ctx = RuntimeContext.builder().onStateLoaded(callback).build();
+        assertSame(callback, ctx.getOnStateLoaded());
+    }
+
+    @Test
+    @DisplayName("setOnStateLoaded replaces and clears the callback")
+    void onStateLoaded_setterReplacesAndClears() {
+        BiConsumer<RuntimeContext, List<Msg>> first = (c, m) -> {};
+        BiConsumer<RuntimeContext, List<Msg>> second = (c, m) -> {};
+        RuntimeContext ctx = RuntimeContext.builder().onStateLoaded(first).build();
+        assertSame(first, ctx.getOnStateLoaded());
+
+        ctx.setOnStateLoaded(second);
+        assertSame(second, ctx.getOnStateLoaded());
+
+        ctx.setOnStateLoaded(null);
+        assertNull(ctx.getOnStateLoaded());
+    }
+
+    @Test
+    @DisplayName("builder(source) preserves the onStateLoaded callback")
+    void onStateLoaded_builderCopyPreservesCallback() {
+        BiConsumer<RuntimeContext, List<Msg>> callback = (c, m) -> {};
+        RuntimeContext source = RuntimeContext.builder().onStateLoaded(callback).build();
+        RuntimeContext copy = RuntimeContext.builder(source).build();
+        assertSame(callback, copy.getOnStateLoaded());
+    }
+
+    @Test
+    @DisplayName("onStateLoaded callback receives the context and mutable message list")
+    void onStateLoaded_callbackReceivesContextAndMsgs() {
+        RuntimeContext ctx = RuntimeContext.empty();
+        AtomicInteger fired = new AtomicInteger();
+        ctx.setOnStateLoaded(
+                (c, m) -> {
+                    fired.incrementAndGet();
+                    assertSame(c, ctx);
+                    m.clear();
+                });
+        AgentState state = AgentState.builder().build();
+        List<Msg> msgs = new ArrayList<>(List.of(new UserMessage("hi")));
+
+        ctx.setAgentState(state);
+        ctx.getOnStateLoaded().accept(ctx, msgs);
+
+        assertEquals(1, fired.get());
+        assertTrue(msgs.isEmpty());
+    }
+
+    @Test
+    @DisplayName("onStateLoaded callback can read the agent state just set on the context")
+    void onStateLoaded_callbackReadsFreshlySetState() {
+        AgentState state = AgentState.builder().summary("loaded").build();
+        AtomicReference<AgentState> seen = new AtomicReference<>();
+        BiConsumer<RuntimeContext, List<Msg>> callback = (c, m) -> seen.set(c.getAgentState());
+        RuntimeContext ctx = RuntimeContext.builder().onStateLoaded(callback).build();
+        ctx.setAgentState(state);
+        ctx.getOnStateLoaded().accept(ctx, List.of());
+
+        assertSame(state, seen.get());
+    }
+
+    @Test
+    @DisplayName("onStateLoaded can be registered on a copied context without affecting the source")
+    void onStateLoaded_copyIsIndependentAfterRegistration() {
+        BiConsumer<RuntimeContext, List<Msg>> original = (c, m) -> {};
+        RuntimeContext source = RuntimeContext.builder().onStateLoaded(original).build();
+        BiConsumer<RuntimeContext, List<Msg>> override = (c, m) -> {};
+        RuntimeContext copy = RuntimeContext.builder(source).onStateLoaded(override).build();
+
+        assertSame(override, copy.getOnStateLoaded());
+        assertSame(original, source.getOnStateLoaded());
+    }
+
     @Test
     @DisplayName("concurrent puts on distinct keys from multiple threads")
     void threadSafety() throws Exception {

agentscope-extensions/agentscope-extensions-protocol/agentscope-extensions-agui/src/main/java/io/agentscope/core/agui/adapter/AguiAgentAdapter.java

@@ -36,6 +36,7 @@
 import io.agentscope.core.message.ToolResultBlock;
 import io.agentscope.core.message.ToolUseBlock;
 import io.agentscope.core.model.ToolSchema;
+import io.agentscope.core.state.AgentState;
 import io.agentscope.core.tool.AgentTool;
 import io.agentscope.core.tool.SchemaOnlyTool;
 import io.agentscope.core.tool.Toolkit;
@@ -50,6 +51,7 @@
 import java.util.Set;
 import java.util.UUID;
 import java.util.concurrent.atomic.AtomicBoolean;
+import java.util.function.BiConsumer;
 import java.util.function.Supplier;
 import reactor.core.publisher.Flux;
 
@@ -308,9 +310,39 @@ protected RuntimeContext buildRuntimeContext(
                 .put(RUNTIME_CONTEXT_STATE_KEY, input.getState())
                 .put(RUNTIME_CONTEXT_FORWARDED_PROPS_KEY, input.getForwardedProps())
                 .put(RUNTIME_CONTEXT_RESUME_KEY, input.getResume())
+                .onStateLoaded(createMessageMergeHandler())
                 .build();
     }
 
+    /**
+     * Creates the {@code onStateLoaded} callback that deduplicates incoming messages against the
+     * already-persisted AgentState context.
+     *
+     * <p>The last message id in {@code state.getContext()} is used as the anchor: if it is found in
+     * the incoming {@code msgs} list, the anchor and everything before it is removed in place (the
+     * input was a full transcript); otherwise the incoming list is treated as purely incremental and
+     * left untouched. When the context is empty or the state is null the callback is a no-op.
+     */
+    private BiConsumer<RuntimeContext, List<Msg>> createMessageMergeHandler() {
+        return (ctx, msgs) -> {
+            AgentState state = ctx.getAgentState();
+            if (state == null) {
+                return;
+            }
+            List<Msg> context = state.getContext();
+            if (context.isEmpty()) {
+                return;
+            }
+            String anchorId = context.get(context.size() - 1).getId();
+            for (int i = msgs.size() - 1; i >= 0; i--) {
+                if (anchorId.equals(msgs.get(i).getId())) {
+                    msgs.subList(0, i + 1).clear();
+                    return;
+                }
+            }
+        };
+    }
+
     @SuppressWarnings("unchecked")
     private Map<String, AguiEvent.Interrupt> resumeInterrupts(RuntimeContext runtimeContext) {
         if (runtimeContext == null) {

agentscope-extensions/agentscope-extensions-protocol/agentscope-extensions-agui/src/main/java/io/agentscope/core/agui/processor/AguiRequestProcessor.java

@@ -21,7 +21,6 @@
 import io.agentscope.core.agui.adapter.AguiAgentAdapter;
 import io.agentscope.core.agui.adapter.AguiAgentAdapterFactory;
 import io.agentscope.core.agui.event.AguiEvent;
-import io.agentscope.core.agui.model.AguiMessage;
 import io.agentscope.core.agui.model.RunAgentInput;
 import java.util.ArrayList;
 import java.util.List;
@@ -40,7 +39,6 @@
  * <p><b>Responsibilities:</b>
  * <ul>
  *   <li>Agent ID resolution from multiple sources</li>
- *   <li>Message extraction for server-side memory scenarios</li>
  *   <li>Agent resolution via {@link AgentResolver}</li>
  *   <li>Event stream generation via {@link AguiAgentAdapter}</li>
  * </ul>
@@ -139,16 +137,10 @@ public ProcessResult process(
                             }
 
                             try {
-                                // Determine effective input based on server-side memory
+                                // Full input is forwarded; message dedup against persisted
+                                // AgentState context is handled by the onStateLoaded callback
+                                // registered in AguiAgentAdapter.buildRuntimeContext().
                                 RunAgentInput effectiveInput = input;
-                                if (agentResolver.hasMemory(threadId)) {
-                                    logger.debug(
-                                            "Using server-side memory for thread {}, extracting"
-                                                    + " latest user message",
-                                            threadId);
-                                    effectiveInput = extractLatestUserMessage(input);
-                                }
-
                                 RuntimeContext effectiveRuntimeContext =
                                         resumeCoordinator.addResumeInterrupts(
                                                 input, runtimeContext);
@@ -264,48 +256,6 @@ public String resolveAgentId(RunAgentInput input, String headerAgentId, String p
         return "default";
     }
 
-    /**
-     * Extract only the latest user message from the input.
-     *
-     * <p>This is used when server-side memory is enabled and the agent already
-     * has conversation history. Only the latest user message needs to be passed.
-     *
-     * @param input The original input
-     * @return A new input with only the latest user message
-     */
-    public RunAgentInput extractLatestUserMessage(RunAgentInput input) {
-        List<AguiMessage> messages = input.getMessages();
-        if (messages == null || messages.isEmpty()) {
-            return input;
-        }
-
-        // Find the last user message
-        AguiMessage lastUserMessage = null;
-        for (int i = messages.size() - 1; i >= 0; i--) {
-            AguiMessage msg = messages.get(i);
-            if ("user".equalsIgnoreCase(msg.getRole())) {
-                lastUserMessage = msg;
-                break;
-            }
-        }
-
-        if (lastUserMessage == null) {
-            return input;
-        }
-
-        // Create new input with only the last user message
-        return RunAgentInput.builder()
-                .threadId(input.getThreadId())
-                .runId(input.getRunId())
-                .messages(List.of(lastUserMessage))
-                .tools(input.getTools())
-                .context(input.getContext())
-                .state(input.getState())
-                .forwardedProps(input.getForwardedProps())
-                .resume(input.getResume())
-                .build();
-    }
-
     /**
      * Creates a new builder for AguiRequestProcessor.
      *