From 1c3c5420d5e19f9f02b4f3a7d5d2ea2b82a8f626 Mon Sep 17 00:00:00 2001 From: guslegend <1670547022@qq.com> Date: Wed, 10 Jun 2026 07:09:36 +0800 Subject: [PATCH 1/4] fix: port auto context memory to v2 --- .../pom.xml | 41 + .../memory/autocontext/AutoContextConfig.java | 144 ++++ .../memory/autocontext/AutoContextHook.java | 193 +++++ .../memory/autocontext/AutoContextMemory.java | 723 ++++++++++++++++++ .../memory/autocontext/AutoContextState.java | 68 ++ .../memory/autocontext/CompressionEvent.java | 140 ++++ .../memory/autocontext/ContextOffLoader.java | 28 + .../autocontext/ContextOffloadTool.java | 103 +++ .../core/memory/autocontext/MsgUtils.java | 378 +++++++++ .../core/memory/autocontext/PromptConfig.java | 77 ++ .../memory/autocontext/PromptProvider.java | 57 ++ .../core/memory/autocontext/Prompts.java | 38 + .../memory/autocontext/TokenCounterUtil.java | 29 + .../autocontext/AutoContextHookTest.java | 117 +++ .../autocontext/AutoContextMemoryTest.java | 112 +++ .../autocontext/AutoContextTestSupport.java | 194 +++++ .../autocontext/ContextOffloadToolTest.java | 88 +++ .../agentscope-extensions-mem/pom.xml | 1 + 18 files changed, 2531 insertions(+) create mode 100644 agentscope-extensions/agentscope-extensions-mem/agentscope-extensions-autocontext-memory/pom.xml create mode 100644 agentscope-extensions/agentscope-extensions-mem/agentscope-extensions-autocontext-memory/src/main/java/io/agentscope/core/memory/autocontext/AutoContextConfig.java create mode 100644 agentscope-extensions/agentscope-extensions-mem/agentscope-extensions-autocontext-memory/src/main/java/io/agentscope/core/memory/autocontext/AutoContextHook.java create mode 100644 agentscope-extensions/agentscope-extensions-mem/agentscope-extensions-autocontext-memory/src/main/java/io/agentscope/core/memory/autocontext/AutoContextMemory.java create mode 100644 agentscope-extensions/agentscope-extensions-mem/agentscope-extensions-autocontext-memory/src/main/java/io/agentscope/core/memory/autocontext/AutoContextState.java create mode 100644 agentscope-extensions/agentscope-extensions-mem/agentscope-extensions-autocontext-memory/src/main/java/io/agentscope/core/memory/autocontext/CompressionEvent.java create mode 100644 agentscope-extensions/agentscope-extensions-mem/agentscope-extensions-autocontext-memory/src/main/java/io/agentscope/core/memory/autocontext/ContextOffLoader.java create mode 100644 agentscope-extensions/agentscope-extensions-mem/agentscope-extensions-autocontext-memory/src/main/java/io/agentscope/core/memory/autocontext/ContextOffloadTool.java create mode 100644 agentscope-extensions/agentscope-extensions-mem/agentscope-extensions-autocontext-memory/src/main/java/io/agentscope/core/memory/autocontext/MsgUtils.java create mode 100644 agentscope-extensions/agentscope-extensions-mem/agentscope-extensions-autocontext-memory/src/main/java/io/agentscope/core/memory/autocontext/PromptConfig.java create mode 100644 agentscope-extensions/agentscope-extensions-mem/agentscope-extensions-autocontext-memory/src/main/java/io/agentscope/core/memory/autocontext/PromptProvider.java create mode 100644 agentscope-extensions/agentscope-extensions-mem/agentscope-extensions-autocontext-memory/src/main/java/io/agentscope/core/memory/autocontext/Prompts.java create mode 100644 agentscope-extensions/agentscope-extensions-mem/agentscope-extensions-autocontext-memory/src/main/java/io/agentscope/core/memory/autocontext/TokenCounterUtil.java create mode 100644 agentscope-extensions/agentscope-extensions-mem/agentscope-extensions-autocontext-memory/src/test/java/io/agentscope/core/memory/autocontext/AutoContextHookTest.java create mode 100644 agentscope-extensions/agentscope-extensions-mem/agentscope-extensions-autocontext-memory/src/test/java/io/agentscope/core/memory/autocontext/AutoContextMemoryTest.java create mode 100644 agentscope-extensions/agentscope-extensions-mem/agentscope-extensions-autocontext-memory/src/test/java/io/agentscope/core/memory/autocontext/AutoContextTestSupport.java create mode 100644 agentscope-extensions/agentscope-extensions-mem/agentscope-extensions-autocontext-memory/src/test/java/io/agentscope/core/memory/autocontext/ContextOffloadToolTest.java diff --git a/agentscope-extensions/agentscope-extensions-mem/agentscope-extensions-autocontext-memory/pom.xml b/agentscope-extensions/agentscope-extensions-mem/agentscope-extensions-autocontext-memory/pom.xml new file mode 100644 index 000000000..e78724461 --- /dev/null +++ b/agentscope-extensions/agentscope-extensions-mem/agentscope-extensions-autocontext-memory/pom.xml @@ -0,0 +1,41 @@ + + + + + 4.0.0 + + io.agentscope + agentscope-extensions-mem + ${revision} + ../pom.xml + + + agentscope-extensions-autocontext-memory + AgentScope Java - Extensions - AutoContext Memory + Auto context compression and offload for AgentScope Java + + + + io.agentscope + agentscope-core + provided + true + + + diff --git a/agentscope-extensions/agentscope-extensions-mem/agentscope-extensions-autocontext-memory/src/main/java/io/agentscope/core/memory/autocontext/AutoContextConfig.java b/agentscope-extensions/agentscope-extensions-mem/agentscope-extensions-autocontext-memory/src/main/java/io/agentscope/core/memory/autocontext/AutoContextConfig.java new file mode 100644 index 000000000..81b6c59de --- /dev/null +++ b/agentscope-extensions/agentscope-extensions-mem/agentscope-extensions-autocontext-memory/src/main/java/io/agentscope/core/memory/autocontext/AutoContextConfig.java @@ -0,0 +1,144 @@ +/* + * Copyright 2024-2026 the original author or authors. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package io.agentscope.core.memory.autocontext; + +/** Configuration for auto context compression. */ +public class AutoContextConfig { + + long largePayloadThreshold = 5 * 1024; + long maxToken = 128 * 1024; + double tokenRatio = 0.75; + int offloadSinglePreview = 200; + int msgThreshold = 100; + int lastKeep = 50; + int minConsecutiveToolMessages = 6; + double currentRoundCompressionRatio = 0.3; + int minCompressionTokenThreshold = 5000; + private PromptConfig customPrompt; + + public long getLargePayloadThreshold() { + return largePayloadThreshold; + } + + public long getMaxToken() { + return maxToken; + } + + public double getTokenRatio() { + return tokenRatio; + } + + public int getOffloadSinglePreview() { + return offloadSinglePreview; + } + + public int getMsgThreshold() { + return msgThreshold; + } + + public int getLastKeep() { + return lastKeep; + } + + public int getMinConsecutiveToolMessages() { + return minConsecutiveToolMessages; + } + + public double getCurrentRoundCompressionRatio() { + return currentRoundCompressionRatio; + } + + public int getMinCompressionTokenThreshold() { + return minCompressionTokenThreshold; + } + + public PromptConfig getCustomPrompt() { + return customPrompt; + } + + public static Builder builder() { + return new Builder(); + } + + public static class Builder { + private final AutoContextConfig config = new AutoContextConfig(); + + public Builder largePayloadThreshold(long largePayloadThreshold) { + config.largePayloadThreshold = largePayloadThreshold; + return this; + } + + public Builder maxToken(long maxToken) { + config.maxToken = maxToken; + return this; + } + + public Builder tokenRatio(double tokenRatio) { + config.tokenRatio = tokenRatio; + return this; + } + + public Builder offloadSinglePreview(int offloadSinglePreview) { + config.offloadSinglePreview = offloadSinglePreview; + return this; + } + + public Builder msgThreshold(int msgThreshold) { + config.msgThreshold = msgThreshold; + return this; + } + + public Builder lastKeep(int lastKeep) { + config.lastKeep = lastKeep; + return this; + } + + public Builder minConsecutiveToolMessages(int minConsecutiveToolMessages) { + config.minConsecutiveToolMessages = minConsecutiveToolMessages; + return this; + } + + public Builder currentRoundCompressionRatio(double currentRoundCompressionRatio) { + config.currentRoundCompressionRatio = currentRoundCompressionRatio; + return this; + } + + public Builder minCompressionTokenThreshold(int minCompressionTokenThreshold) { + config.minCompressionTokenThreshold = minCompressionTokenThreshold; + return this; + } + + public Builder customPrompt(PromptConfig customPrompt) { + config.customPrompt = customPrompt; + return this; + } + + public AutoContextConfig build() { + AutoContextConfig result = new AutoContextConfig(); + result.largePayloadThreshold = config.largePayloadThreshold; + result.maxToken = config.maxToken; + result.tokenRatio = config.tokenRatio; + result.offloadSinglePreview = config.offloadSinglePreview; + result.msgThreshold = config.msgThreshold; + result.lastKeep = config.lastKeep; + result.minConsecutiveToolMessages = config.minConsecutiveToolMessages; + result.currentRoundCompressionRatio = config.currentRoundCompressionRatio; + result.minCompressionTokenThreshold = config.minCompressionTokenThreshold; + result.customPrompt = config.customPrompt; + return result; + } + } +} diff --git a/agentscope-extensions/agentscope-extensions-mem/agentscope-extensions-autocontext-memory/src/main/java/io/agentscope/core/memory/autocontext/AutoContextHook.java b/agentscope-extensions/agentscope-extensions-mem/agentscope-extensions-autocontext-memory/src/main/java/io/agentscope/core/memory/autocontext/AutoContextHook.java new file mode 100644 index 000000000..df7912eb2 --- /dev/null +++ b/agentscope-extensions/agentscope-extensions-mem/agentscope-extensions-autocontext-memory/src/main/java/io/agentscope/core/memory/autocontext/AutoContextHook.java @@ -0,0 +1,193 @@ +/* + * Copyright 2024-2026 the original author or authors. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package io.agentscope.core.memory.autocontext; + +import io.agentscope.core.ReActAgent; +import io.agentscope.core.agent.Agent; +import io.agentscope.core.agent.RuntimeContext; +import io.agentscope.core.hook.Hook; +import io.agentscope.core.hook.HookEvent; +import io.agentscope.core.hook.PreCallEvent; +import io.agentscope.core.hook.PreReasoningEvent; +import io.agentscope.core.message.Msg; +import io.agentscope.core.state.AgentState; +import io.agentscope.core.state.AgentStateStore; +import io.agentscope.core.tool.Toolkit; +import java.util.ArrayList; +import java.util.List; +import java.util.Objects; +import java.util.concurrent.ConcurrentHashMap; +import java.util.concurrent.atomic.AtomicBoolean; +import reactor.core.publisher.Mono; +import reactor.core.scheduler.Schedulers; + +/** Hook that wires auto context compression into ReActAgent. */ +@SuppressWarnings("deprecation") +public class AutoContextHook implements Hook { + + private final AutoContextConfig config; + private final io.agentscope.core.model.Model model; + private final String hookId; + private final ConcurrentHashMap memories = new ConcurrentHashMap<>(); + private final ConcurrentHashMap registeredAgents = + new ConcurrentHashMap<>(); + private final AtomicBoolean registrationFailed = new AtomicBoolean(false); + + public AutoContextHook() { + this(AutoContextConfig.builder().build(), null); + } + + public AutoContextHook(AutoContextConfig config, io.agentscope.core.model.Model model) { + this.config = Objects.requireNonNull(config, "config"); + this.model = model; + this.hookId = java.util.UUID.randomUUID().toString().replace("-", ""); + } + + @Override + public Mono onEvent(T event) { + if (event instanceof PreCallEvent preCallEvent) { + @SuppressWarnings("unchecked") + Mono mono = (Mono) handlePreCall(preCallEvent); + return mono; + } + if (event instanceof PreReasoningEvent preReasoningEvent) { + @SuppressWarnings("unchecked") + Mono mono = (Mono) handlePreReasoning(preReasoningEvent); + return mono; + } + return Mono.just(event); + } + + @Override + public int priority() { + return 0; + } + + @Override + public List tools() { + return List.of(); + } + + Mono handlePreCall(PreCallEvent event) { + Agent agent = event.getAgent(); + if (!(agent instanceof ReActAgent reactAgent)) { + return Mono.just(event); + } + if (registeredAgents.putIfAbsent(reactAgent, Boolean.TRUE) != null) { + return Mono.just(event); + } + try { + Toolkit toolkit = reactAgent.getToolkit(); + if (toolkit != null && !toolkit.getToolNames().contains("context_reload")) { + toolkit.registerTool(new ContextOffloadTool(this)); + } + } catch (Exception e) { + registrationFailed.set(true); + } + return Mono.just(event); + } + + Mono handlePreReasoning(PreReasoningEvent event) { + Agent agent = event.getAgent(); + if (!(agent instanceof ReActAgent reactAgent)) { + return Mono.just(event); + } + return Mono.fromCallable( + () -> { + RuntimeContext runtimeContext = reactAgent.getRuntimeContext(); + AgentState state = + RuntimeContext.resolveAgentState(runtimeContext, reactAgent); + if (state == null) { + return event; + } + + AutoContextMemory memory = memoryFor(reactAgent, runtimeContext); + memory.mergeWithContext(state.getContext()); + boolean compressed = memory.compressIfNeeded(); + if (compressed) { + state.contextMutable().clear(); + state.contextMutable().addAll(memory.getMessages()); + } + + event.setInputMessages(new ArrayList<>(memory.getMessages())); + event.appendSystemContent( + "You may see compressed messages containing .\n" + + "- Use the UUID to call context_reload if you need full" + + " details.\n" + + "- Never mention UUIDs, offload tags, or internal" + + " metadata in your response."); + + persistMemory(reactAgent, runtimeContext, memory); + return event; + }) + .subscribeOn(Schedulers.boundedElastic()); + } + + List reload(String uuid, Agent agent, RuntimeContext runtimeContext) { + AutoContextMemory memory = memoryFor(agent, runtimeContext); + return memory.reload(uuid); + } + + AutoContextMemory memoryFor(Agent agent, RuntimeContext runtimeContext) { + String key = memoryKey(agent, runtimeContext); + return memories.computeIfAbsent( + key, + ignored -> { + AutoContextMemory memory = new AutoContextMemory(config, model); + AgentState state = RuntimeContext.resolveAgentState(runtimeContext, agent); + if (agent instanceof ReActAgent reactAgent) { + AgentStateStore store = reactAgent.getStateStore(); + String userId = runtimeContext != null ? runtimeContext.getUserId() : null; + String sessionId = + runtimeContext != null ? runtimeContext.getSessionId() : null; + if (store != null && sessionId != null && !sessionId.isBlank()) { + memory.loadFrom(store, userId, sessionId, stateKey(agent)); + } + } + if (state != null) { + memory.mergeWithContext(state.getContext()); + } + return memory; + }); + } + + private void persistMemory( + ReActAgent agent, RuntimeContext runtimeContext, AutoContextMemory memory) { + AgentStateStore store = agent.getStateStore(); + if (store == null || runtimeContext == null || runtimeContext.getSessionId() == null) { + return; + } + memory.saveTo( + store, runtimeContext.getUserId(), runtimeContext.getSessionId(), stateKey(agent)); + } + + private String memoryKey(Agent agent, RuntimeContext runtimeContext) { + return stateKey(agent) + + "|" + + (runtimeContext != null ? safe(runtimeContext.getUserId()) : "__anon__") + + "|" + + (runtimeContext != null ? safe(runtimeContext.getSessionId()) : "__default__"); + } + + private String stateKey(Agent agent) { + return "autocontext_" + hookId + "_" + Integer.toHexString(System.identityHashCode(agent)); + } + + private String safe(String value) { + return value == null || value.isBlank() ? "__anon__" : value; + } +} diff --git a/agentscope-extensions/agentscope-extensions-mem/agentscope-extensions-autocontext-memory/src/main/java/io/agentscope/core/memory/autocontext/AutoContextMemory.java b/agentscope-extensions/agentscope-extensions-mem/agentscope-extensions-autocontext-memory/src/main/java/io/agentscope/core/memory/autocontext/AutoContextMemory.java new file mode 100644 index 000000000..480cfaef6 --- /dev/null +++ b/agentscope-extensions/agentscope-extensions-mem/agentscope-extensions-autocontext-memory/src/main/java/io/agentscope/core/memory/autocontext/AutoContextMemory.java @@ -0,0 +1,723 @@ +/* + * Copyright 2024-2026 the original author or authors. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package io.agentscope.core.memory.autocontext; + +import io.agentscope.core.memory.Memory; +import io.agentscope.core.message.ContentBlock; +import io.agentscope.core.message.Msg; +import io.agentscope.core.message.MsgRole; +import io.agentscope.core.message.TextBlock; +import io.agentscope.core.message.ToolResultBlock; +import io.agentscope.core.model.ChatResponse; +import io.agentscope.core.model.ChatUsage; +import io.agentscope.core.model.GenerateOptions; +import io.agentscope.core.model.Model; +import io.agentscope.core.state.AgentStateStore; +import java.util.ArrayList; +import java.util.HashMap; +import java.util.List; +import java.util.Map; +import java.util.Objects; +import java.util.UUID; +import reactor.core.publisher.Mono; +import reactor.core.scheduler.Schedulers; + +/** Auto context compressor that keeps a working buffer plus offloaded originals. */ +public class AutoContextMemory implements Memory, ContextOffLoader { + + private static final String STATE_KEY = "auto_context_state"; + + private final AutoContextConfig autoContextConfig; + private final PromptConfig customPrompt; + private Model model; + private final List workingMemoryStorage = new ArrayList<>(); + private final List originalMemoryStorage = new ArrayList<>(); + private final Map> offloadContext = new HashMap<>(); + private final List compressionEvents = new ArrayList<>(); + + public AutoContextMemory(AutoContextConfig autoContextConfig, Model model) { + this.autoContextConfig = Objects.requireNonNull(autoContextConfig, "autoContextConfig"); + this.model = model; + this.customPrompt = autoContextConfig.getCustomPrompt(); + } + + public synchronized void setModel(Model model) { + this.model = model; + } + + public synchronized AutoContextConfig getAutoContextConfig() { + return autoContextConfig; + } + + public synchronized void mergeWithContext(List context) { + if (context == null || context.isEmpty()) { + return; + } + if (workingMemoryStorage.isEmpty()) { + workingMemoryStorage.addAll(copyMessages(context)); + if (originalMemoryStorage.isEmpty()) { + originalMemoryStorage.addAll(copyMessages(context)); + } + return; + } + if (context.size() < workingMemoryStorage.size() + || !startsWith(context, workingMemoryStorage)) { + workingMemoryStorage.clear(); + workingMemoryStorage.addAll(copyMessages(context)); + if (originalMemoryStorage.isEmpty()) { + originalMemoryStorage.addAll(copyMessages(context)); + } + return; + } + if (context.size() > workingMemoryStorage.size()) { + List tail = + new ArrayList<>(context.subList(workingMemoryStorage.size(), context.size())); + workingMemoryStorage.addAll(tail); + originalMemoryStorage.addAll(tail); + } + } + + public synchronized AutoContextState snapshot() { + AutoContextState snapshot = new AutoContextState(); + snapshot.setWorkingMessages(workingMemoryStorage); + snapshot.setOriginalMessages(originalMemoryStorage); + snapshot.setOffloadContext(offloadContext); + snapshot.setCompressionEvents(compressionEvents); + return snapshot; + } + + public synchronized void restore(AutoContextState snapshot) { + if (snapshot == null) { + return; + } + workingMemoryStorage.clear(); + workingMemoryStorage.addAll(copyMessages(snapshot.getWorkingMessages())); + originalMemoryStorage.clear(); + originalMemoryStorage.addAll(copyMessages(snapshot.getOriginalMessages())); + offloadContext.clear(); + offloadContext.putAll(copyMsgMap(snapshot.getOffloadContext())); + compressionEvents.clear(); + compressionEvents.addAll( + snapshot.getCompressionEvents() != null + ? snapshot.getCompressionEvents() + : List.of()); + } + + @Override + public synchronized void saveTo(AgentStateStore stateStore, String userId, String sessionId) { + if (stateStore == null) { + return; + } + stateStore.save(userId, sessionId, STATE_KEY, snapshot()); + } + + public synchronized void saveTo( + AgentStateStore stateStore, String userId, String sessionId, String key) { + if (stateStore == null) { + return; + } + stateStore.save(userId, sessionId, key, snapshot()); + } + + @Override + public synchronized void loadFrom(AgentStateStore stateStore, String userId, String sessionId) { + loadFrom(stateStore, userId, sessionId, STATE_KEY); + } + + public synchronized void loadFrom( + AgentStateStore stateStore, String userId, String sessionId, String key) { + if (stateStore == null) { + return; + } + stateStore.get(userId, sessionId, key, AutoContextState.class).ifPresent(this::restore); + } + + @Override + public synchronized void addMessage(Msg message) { + if (message == null) { + return; + } + workingMemoryStorage.add(message); + originalMemoryStorage.add(message); + } + + @Override + public synchronized List getMessages() { + return new ArrayList<>(workingMemoryStorage); + } + + @Override + public synchronized void deleteMessage(int index) { + if (index < 0 || index >= workingMemoryStorage.size()) { + return; + } + workingMemoryStorage.remove(index); + if (index < originalMemoryStorage.size()) { + originalMemoryStorage.remove(index); + } + } + + public synchronized List getOriginalMemoryMsgs() { + return new ArrayList<>(originalMemoryStorage); + } + + public synchronized List getInteractionMsgs() { + List interactions = new ArrayList<>(); + for (Msg msg : originalMemoryStorage) { + if (msg.getRole() == MsgRole.USER || MsgUtils.isFinalAssistantResponse(msg)) { + interactions.add(msg); + } + } + return interactions; + } + + public synchronized Map> getOffloadContext() { + return offloadContext; + } + + public synchronized List getCompressionEvents() { + return compressionEvents; + } + + public synchronized boolean compressIfNeeded() { + if (workingMemoryStorage.isEmpty()) { + return false; + } + + int tokenCount = TokenCounterUtil.calculateToken(workingMemoryStorage); + boolean msgCountReached = + workingMemoryStorage.size() >= autoContextConfig.getMsgThreshold(); + boolean tokenReached = + tokenCount + >= (int) + (autoContextConfig.getMaxToken() + * autoContextConfig.getTokenRatio()); + if (!msgCountReached && !tokenReached) { + return false; + } + + if (tokenCount < autoContextConfig.getMinCompressionTokenThreshold()) { + return false; + } + + if (compressToolGroups()) { + return true; + } + if (offloadLargeMessages()) { + return true; + } + if (compressPreviousRounds()) { + return true; + } + return compressCurrentRound(); + } + + public synchronized Mono compressIfNeededAsync() { + return Mono.fromCallable(this::compressIfNeeded).subscribeOn(Schedulers.boundedElastic()); + } + + @Override + public synchronized void offload(String uuid, List messages) { + if (uuid == null || uuid.isBlank() || messages == null) { + return; + } + offloadContext.put(uuid, new ArrayList<>(messages)); + } + + @Override + public synchronized List reload(String uuid) { + List messages = offloadContext.get(uuid); + return messages == null ? new ArrayList<>() : new ArrayList<>(messages); + } + + @Override + public synchronized void clear(String uuid) { + if (uuid != null) { + offloadContext.remove(uuid); + } + } + + @Override + public synchronized void clear() { + workingMemoryStorage.clear(); + originalMemoryStorage.clear(); + offloadContext.clear(); + compressionEvents.clear(); + } + + public synchronized void setWorkingMessages(List messages) { + workingMemoryStorage.clear(); + workingMemoryStorage.addAll(copyMessages(messages)); + } + + private boolean compressToolGroups() { + int cursor = 0; + int upperLimit = Math.max(0, workingMemoryStorage.size() - autoContextConfig.getLastKeep()); + boolean changed = false; + while (cursor < upperLimit) { + IntRange range = findToolGroup(cursor, upperLimit); + if (range == null) { + break; + } + List messages = + new ArrayList<>(workingMemoryStorage.subList(range.start, range.end + 1)); + if (TokenCounterUtil.calculateToken(messages) + < autoContextConfig.getMinCompressionTokenThreshold()) { + cursor = range.end + 1; + continue; + } + if (compressRange( + range.start, + range.end, + CompressionEvent.TOOL_INVOCATION_COMPRESS, + PromptProvider.getPreviousRoundToolCompressPrompt(customPrompt), + true)) { + changed = true; + cursor = range.start + 1; + upperLimit = + Math.max(0, workingMemoryStorage.size() - autoContextConfig.getLastKeep()); + } else { + cursor = range.end + 1; + } + } + return changed; + } + + private boolean offloadLargeMessages() { + int limit = Math.max(0, workingMemoryStorage.size() - autoContextConfig.getLastKeep()); + for (int i = 0; i < limit; i++) { + Msg msg = workingMemoryStorage.get(i); + if (MsgUtils.calculateMessageCharCount(msg) + < autoContextConfig.getLargePayloadThreshold()) { + continue; + } + String uuid = UUID.randomUUID().toString(); + List original = List.of(msg); + String summary = + summarizeMessages( + original, + PromptProvider.getCurrentRoundLargeMessagePrompt(customPrompt), + false); + if (summary == null || summary.isBlank()) { + summary = fallbackSummary(original); + } + offload(uuid, original); + Msg compressed = compressSingleMessage(msg, summary, uuid); + MsgUtils.replaceMsg(workingMemoryStorage, i, i, compressed); + recordCompressionEvent( + msg.getRole() == MsgRole.TOOL + ? CompressionEvent.LARGE_MESSAGE_OFFLOAD + : CompressionEvent.LARGE_MESSAGE_OFFLOAD_WITH_PROTECTION, + i, + i, + original, + compressed, + buildCompressionMetadata(compressed)); + return true; + } + return false; + } + + private boolean compressPreviousRounds() { + boolean changed = false; + int guard = 0; + while (workingMemoryStorage.size() > autoContextConfig.getLastKeep() + 1 && guard++ < 32) { + int end = findPrefixCompressionEnd(); + if (end < 0) { + break; + } + List source = new ArrayList<>(workingMemoryStorage.subList(0, end + 1)); + if (TokenCounterUtil.calculateToken(source) + < autoContextConfig.getMinCompressionTokenThreshold()) { + break; + } + if (!compressRange( + 0, + end, + CompressionEvent.PREVIOUS_ROUND_CONVERSATION_SUMMARY, + PromptProvider.getPreviousRoundSummaryPrompt(customPrompt), + false)) { + break; + } + changed = true; + } + return changed; + } + + private boolean compressCurrentRound() { + int latestUserIndex = -1; + for (int i = workingMemoryStorage.size() - 1; i >= 0; i--) { + if (workingMemoryStorage.get(i).getRole() == MsgRole.USER) { + latestUserIndex = i; + break; + } + } + if (latestUserIndex < 0 || latestUserIndex >= workingMemoryStorage.size() - 1) { + return false; + } + + int end = workingMemoryStorage.size() - 1; + if (MsgUtils.isToolUseMessage(workingMemoryStorage.get(end))) { + end--; + } + if (end <= latestUserIndex) { + return false; + } + + List source = + new ArrayList<>(workingMemoryStorage.subList(latestUserIndex + 1, end + 1)); + if (TokenCounterUtil.calculateToken(source) + < autoContextConfig.getMinCompressionTokenThreshold()) { + return false; + } + return compressRange( + latestUserIndex + 1, + end, + CompressionEvent.CURRENT_ROUND_MESSAGE_COMPRESS, + PromptProvider.getCurrentRoundCompressPrompt(customPrompt), + true); + } + + private boolean compressRange( + int start, int end, String eventType, String prompt, boolean appendOffloadTag) { + if (start < 0 || end < start || start >= workingMemoryStorage.size()) { + return false; + } + end = Math.min(end, workingMemoryStorage.size() - 1); + List source = new ArrayList<>(workingMemoryStorage.subList(start, end + 1)); + String uuid = UUID.randomUUID().toString(); + offload(uuid, source); + String summary = + summarizeMessages( + source, + prompt, + eventType.equals(CompressionEvent.CURRENT_ROUND_MESSAGE_COMPRESS)); + if (summary == null || summary.isBlank()) { + summary = fallbackSummary(source); + } + Msg summaryMsg = buildSummaryMessage(source, summary, uuid, appendOffloadTag); + MsgUtils.replaceMsg(workingMemoryStorage, start, end, summaryMsg); + recordCompressionEvent( + eventType, start, end, source, summaryMsg, buildCompressionMetadata(summaryMsg)); + return true; + } + + private String summarizeMessages(List messages, String prompt, boolean currentRound) { + if (messages == null || messages.isEmpty()) { + return null; + } + if (model == null) { + return fallbackSummary(messages); + } + try { + List input = new ArrayList<>(); + input.add( + Msg.builder() + .role(MsgRole.USER) + .name("user") + .content(TextBlock.builder().text(prompt).build()) + .build()); + input.addAll(messages); + input.add( + Msg.builder() + .role(MsgRole.USER) + .name("user") + .content( + TextBlock.builder() + .text(Prompts.COMPRESSION_MESSAGE_LIST_END) + .build()) + .build()); + if (currentRound) { + int chars = MsgUtils.calculateMessagesCharCount(messages); + int target = + Math.max( + 1, + (int) + (chars + * autoContextConfig + .getCurrentRoundCompressionRatio())); + input.add( + Msg.builder() + .role(MsgRole.USER) + .name("user") + .content( + TextBlock.builder() + .text( + String.format( + "\n" + + "Compress to about %d" + + " characters.", + target)) + .build()) + .build()); + } + List responses = + model.stream(input, null, GenerateOptions.builder().build()) + .collectList() + .block(); + if (responses == null || responses.isEmpty()) { + return fallbackSummary(messages); + } + ChatResponse last = responses.get(responses.size() - 1); + String text = extractResponseText(last); + return text == null || text.isBlank() ? fallbackSummary(messages) : text.trim(); + } catch (Exception e) { + return fallbackSummary(messages); + } + } + + private String extractResponseText(ChatResponse response) { + if (response == null || response.getContent() == null) { + return null; + } + StringBuilder sb = new StringBuilder(); + for (ContentBlock block : response.getContent()) { + if (block instanceof TextBlock text && text.getText() != null) { + sb.append(text.getText()); + } + } + return sb.length() == 0 ? null : sb.toString(); + } + + private Msg buildSummaryMessage( + List source, String summary, String uuid, boolean appendOffloadTag) { + MsgRole role = determineSummaryRole(source); + Map metadata = new HashMap<>(); + Map compressMeta = new HashMap<>(); + compressMeta.put("offloaduuid", uuid); + if (appendOffloadTag) { + compressMeta.put("compressed_current_round", Boolean.TRUE); + } + metadata.put("_compress_meta", compressMeta); + ChatUsage usage = null; + String content = summary; + if (appendOffloadTag) { + content = content + "\n" + String.format(Prompts.CONTEXT_OFFLOAD_TAG_FORMAT, uuid); + } + return Msg.builder() + .role(role) + .name(determineName(role)) + .content(TextBlock.builder().text(content).build()) + .metadata(metadata) + .usage(usage) + .build(); + } + + private Msg compressSingleMessage(Msg message, String summary, String uuid) { + List content = new ArrayList<>(); + boolean summaryInserted = false; + for (ContentBlock block : message.getContent()) { + if (block instanceof TextBlock && !summaryInserted) { + content.add(TextBlock.builder().text(summary).build()); + summaryInserted = true; + } else if (block instanceof ToolResultBlock toolResult) { + content.add( + ToolResultBlock.builder() + .name(toolResult.getName()) + .id(toolResult.getId()) + .output(List.of(TextBlock.builder().text(summary).build())) + .build()); + } else { + content.add(block); + } + } + if (!summaryInserted) { + content.add(TextBlock.builder().text(summary).build()); + } + Map metadata = new HashMap<>(message.getMetadata()); + Map compressMeta = new HashMap<>(); + compressMeta.put("offloaduuid", uuid); + metadata.put("_compress_meta", compressMeta); + return Msg.builder() + .id(message.getId()) + .name(message.getName()) + .role(message.getRole()) + .content(content) + .metadata(metadata) + .timestamp(message.getTimestamp()) + .usage(message.getUsage()) + .build(); + } + + private Map buildCompressionMetadata(Msg msg) { + Map metadata = new HashMap<>(); + metadata.put("tokenBefore", 0); + metadata.put("tokenAfter", 0); + if (msg != null && msg.getUsage() != null) { + metadata.put("inputToken", msg.getUsage().getInputTokens()); + metadata.put("outputToken", msg.getUsage().getOutputTokens()); + metadata.put("time", msg.getUsage().getTime()); + } + return metadata; + } + + private void recordCompressionEvent( + String eventType, + int startIndex, + int endIndex, + List source, + Msg compressedMessage, + Map metadata) { + int tokenBefore = TokenCounterUtil.calculateToken(source); + int tokenAfter = TokenCounterUtil.calculateToken(List.of(compressedMessage)); + metadata.put("tokenBefore", tokenBefore); + metadata.put("tokenAfter", tokenAfter); + CompressionEvent event = + new CompressionEvent( + eventType, + System.currentTimeMillis(), + Math.max(1, endIndex - startIndex + 1), + startIndex > 0 + ? workingMemoryStorage + .get( + Math.min( + startIndex - 1, + workingMemoryStorage.size() - 1)) + .getId() + : null, + endIndex < workingMemoryStorage.size() - 1 + ? workingMemoryStorage.get(endIndex + 1).getId() + : null, + compressedMessage != null ? compressedMessage.getId() : null, + metadata); + compressionEvents.add(event); + } + + private IntRange findToolGroup(int startIndex, int upperLimit) { + int start = -1; + for (int i = startIndex; i < upperLimit; i++) { + Msg msg = workingMemoryStorage.get(i); + if (MsgUtils.isToolMessage(msg)) { + if (start < 0) { + start = i; + } + } else if (start >= 0) { + int end = i - 1; + if (end - start + 1 >= autoContextConfig.getMinConsecutiveToolMessages()) { + return new IntRange(start, end); + } + start = -1; + } + } + if (start >= 0 && upperLimit - start >= autoContextConfig.getMinConsecutiveToolMessages()) { + return new IntRange(start, upperLimit - 1); + } + return null; + } + + private int findPrefixCompressionEnd() { + int end = workingMemoryStorage.size() - autoContextConfig.getLastKeep() - 1; + if (end < 0) { + return -1; + } + while (end > 0 && MsgUtils.isToolResultMessage(workingMemoryStorage.get(end))) { + end--; + } + while (end > 0 && MsgUtils.isToolUseMessage(workingMemoryStorage.get(end))) { + end--; + } + return end; + } + + private String fallbackSummary(List messages) { + StringBuilder sb = new StringBuilder(); + sb.append("Compressed ").append(messages.size()).append(" messages."); + int previewLimit = autoContextConfig.getOffloadSinglePreview(); + String preview = buildPreview(messages); + if (!preview.isBlank()) { + sb.append(" ").append(truncate(preview, previewLimit)); + } + return sb.toString(); + } + + private String buildPreview(List messages) { + StringBuilder sb = new StringBuilder(); + for (Msg msg : messages) { + if (msg == null) { + continue; + } + String text = msg.getTextContent(); + if (text == null || text.isBlank()) { + text = msg.getContent() != null ? msg.getContent().toString() : ""; + } + if (!text.isBlank()) { + if (!sb.isEmpty()) { + sb.append(" | "); + } + sb.append(text); + } + } + return sb.toString(); + } + + private static String truncate(String text, int maxChars) { + if (text == null || maxChars <= 0 || text.length() <= maxChars) { + return text == null ? "" : text; + } + return text.substring(0, maxChars); + } + + private MsgRole determineSummaryRole(List source) { + if (source == null || source.isEmpty()) { + return MsgRole.ASSISTANT; + } + boolean allTool = true; + for (Msg msg : source) { + if (msg.getRole() != MsgRole.TOOL) { + allTool = false; + break; + } + } + return allTool ? MsgRole.TOOL : MsgRole.ASSISTANT; + } + + private String determineName(MsgRole role) { + return switch (role) { + case USER -> "user"; + case TOOL -> "tool"; + case SYSTEM, ASSISTANT -> "assistant"; + }; + } + + private boolean startsWith(List full, List prefix) { + if (full == null || prefix == null || full.size() < prefix.size()) { + return false; + } + for (int i = 0; i < prefix.size(); i++) { + if (!Objects.equals(full.get(i), prefix.get(i))) { + return false; + } + } + return true; + } + + private List copyMessages(List messages) { + return messages == null ? new ArrayList<>() : new ArrayList<>(messages); + } + + private Map> copyMsgMap(Map> input) { + Map> result = new HashMap<>(); + if (input == null) { + return result; + } + for (Map.Entry> entry : input.entrySet()) { + result.put(entry.getKey(), copyMessages(entry.getValue())); + } + return result; + } + + private record IntRange(int start, int end) {} +} diff --git a/agentscope-extensions/agentscope-extensions-mem/agentscope-extensions-autocontext-memory/src/main/java/io/agentscope/core/memory/autocontext/AutoContextState.java b/agentscope-extensions/agentscope-extensions-mem/agentscope-extensions-autocontext-memory/src/main/java/io/agentscope/core/memory/autocontext/AutoContextState.java new file mode 100644 index 000000000..521689287 --- /dev/null +++ b/agentscope-extensions/agentscope-extensions-mem/agentscope-extensions-autocontext-memory/src/main/java/io/agentscope/core/memory/autocontext/AutoContextState.java @@ -0,0 +1,68 @@ +/* + * Copyright 2024-2026 the original author or authors. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package io.agentscope.core.memory.autocontext; + +import io.agentscope.core.message.Msg; +import io.agentscope.core.state.State; +import java.util.ArrayList; +import java.util.HashMap; +import java.util.List; +import java.util.Map; + +/** Persisted snapshot for auto context state. */ +public class AutoContextState implements State { + + private List workingMessages = new ArrayList<>(); + private List originalMessages = new ArrayList<>(); + private Map> offloadContext = new HashMap<>(); + private List compressionEvents = new ArrayList<>(); + + public List getWorkingMessages() { + return workingMessages; + } + + public void setWorkingMessages(List workingMessages) { + this.workingMessages = + workingMessages == null ? new ArrayList<>() : new ArrayList<>(workingMessages); + } + + public List getOriginalMessages() { + return originalMessages; + } + + public void setOriginalMessages(List originalMessages) { + this.originalMessages = + originalMessages == null ? new ArrayList<>() : new ArrayList<>(originalMessages); + } + + public Map> getOffloadContext() { + return offloadContext; + } + + public void setOffloadContext(Map> offloadContext) { + this.offloadContext = + offloadContext == null ? new HashMap<>() : new HashMap<>(offloadContext); + } + + public List getCompressionEvents() { + return compressionEvents; + } + + public void setCompressionEvents(List compressionEvents) { + this.compressionEvents = + compressionEvents == null ? new ArrayList<>() : new ArrayList<>(compressionEvents); + } +} diff --git a/agentscope-extensions/agentscope-extensions-mem/agentscope-extensions-autocontext-memory/src/main/java/io/agentscope/core/memory/autocontext/CompressionEvent.java b/agentscope-extensions/agentscope-extensions-mem/agentscope-extensions-autocontext-memory/src/main/java/io/agentscope/core/memory/autocontext/CompressionEvent.java new file mode 100644 index 000000000..2949121f0 --- /dev/null +++ b/agentscope-extensions/agentscope-extensions-mem/agentscope-extensions-autocontext-memory/src/main/java/io/agentscope/core/memory/autocontext/CompressionEvent.java @@ -0,0 +1,140 @@ +/* + * Copyright 2024-2026 the original author or authors. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package io.agentscope.core.memory.autocontext; + +import io.agentscope.core.state.State; +import java.util.HashMap; +import java.util.Map; + +/** Compression event record persisted by auto context. */ +public class CompressionEvent implements State { + public static final String TOOL_INVOCATION_COMPRESS = "TOOL_INVOCATION_COMPRESS"; + public static final String LARGE_MESSAGE_OFFLOAD_WITH_PROTECTION = + "LARGE_MESSAGE_OFFLOAD_WITH_PROTECTION"; + public static final String LARGE_MESSAGE_OFFLOAD = "LARGE_MESSAGE_OFFLOAD"; + public static final String PREVIOUS_ROUND_CONVERSATION_SUMMARY = + "PREVIOUS_ROUND_CONVERSATION_SUMMARY"; + public static final String CURRENT_ROUND_LARGE_MESSAGE_SUMMARY = + "CURRENT_ROUND_LARGE_MESSAGE_SUMMARY"; + public static final String CURRENT_ROUND_MESSAGE_COMPRESS = "CURRENT_ROUND_MESSAGE_COMPRESS"; + + private String eventType; + private long timestamp; + private int compressedMessageCount; + private String previousMessageId; + private String nextMessageId; + private String compressedMessageId; + private Map metadata = new HashMap<>(); + + public CompressionEvent() {} + + public CompressionEvent( + String eventType, + long timestamp, + int compressedMessageCount, + String previousMessageId, + String nextMessageId, + String compressedMessageId, + Map metadata) { + this.eventType = eventType; + this.timestamp = timestamp; + this.compressedMessageCount = compressedMessageCount; + this.previousMessageId = previousMessageId; + this.nextMessageId = nextMessageId; + this.compressedMessageId = compressedMessageId; + this.metadata = metadata != null ? new HashMap<>(metadata) : new HashMap<>(); + } + + public String getEventType() { + return eventType; + } + + public void setEventType(String eventType) { + this.eventType = eventType; + } + + public long getTimestamp() { + return timestamp; + } + + public void setTimestamp(long timestamp) { + this.timestamp = timestamp; + } + + public int getCompressedMessageCount() { + return compressedMessageCount; + } + + public void setCompressedMessageCount(int compressedMessageCount) { + this.compressedMessageCount = compressedMessageCount; + } + + public String getPreviousMessageId() { + return previousMessageId; + } + + public void setPreviousMessageId(String previousMessageId) { + this.previousMessageId = previousMessageId; + } + + public String getNextMessageId() { + return nextMessageId; + } + + public void setNextMessageId(String nextMessageId) { + this.nextMessageId = nextMessageId; + } + + public String getCompressedMessageId() { + return compressedMessageId; + } + + public void setCompressedMessageId(String compressedMessageId) { + this.compressedMessageId = compressedMessageId; + } + + public Map getMetadata() { + return metadata; + } + + public void setMetadata(Map metadata) { + this.metadata = metadata != null ? new HashMap<>(metadata) : new HashMap<>(); + } + + public int getTokenBefore() { + Object value = metadata != null ? metadata.get("tokenBefore") : null; + return value instanceof Number n ? n.intValue() : 0; + } + + public int getTokenAfter() { + Object value = metadata != null ? metadata.get("tokenAfter") : null; + return value instanceof Number n ? n.intValue() : 0; + } + + public int getTokenReduction() { + return getTokenBefore() - getTokenAfter(); + } + + public int getCompressInputToken() { + Object value = metadata != null ? metadata.get("inputToken") : null; + return value instanceof Number n ? n.intValue() : 0; + } + + public int getCompressOutputToken() { + Object value = metadata != null ? metadata.get("outputToken") : null; + return value instanceof Number n ? n.intValue() : 0; + } +} diff --git a/agentscope-extensions/agentscope-extensions-mem/agentscope-extensions-autocontext-memory/src/main/java/io/agentscope/core/memory/autocontext/ContextOffLoader.java b/agentscope-extensions/agentscope-extensions-mem/agentscope-extensions-autocontext-memory/src/main/java/io/agentscope/core/memory/autocontext/ContextOffLoader.java new file mode 100644 index 000000000..179513f6d --- /dev/null +++ b/agentscope-extensions/agentscope-extensions-mem/agentscope-extensions-autocontext-memory/src/main/java/io/agentscope/core/memory/autocontext/ContextOffLoader.java @@ -0,0 +1,28 @@ +/* + * Copyright 2024-2026 the original author or authors. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package io.agentscope.core.memory.autocontext; + +import io.agentscope.core.message.Msg; +import java.util.List; + +/** Storage abstraction for offloaded context. */ +interface ContextOffLoader { + void offload(String uuid, List messages); + + List reload(String uuid); + + void clear(String uuid); +} diff --git a/agentscope-extensions/agentscope-extensions-mem/agentscope-extensions-autocontext-memory/src/main/java/io/agentscope/core/memory/autocontext/ContextOffloadTool.java b/agentscope-extensions/agentscope-extensions-mem/agentscope-extensions-autocontext-memory/src/main/java/io/agentscope/core/memory/autocontext/ContextOffloadTool.java new file mode 100644 index 000000000..248e7f255 --- /dev/null +++ b/agentscope-extensions/agentscope-extensions-mem/agentscope-extensions-autocontext-memory/src/main/java/io/agentscope/core/memory/autocontext/ContextOffloadTool.java @@ -0,0 +1,103 @@ +/* + * Copyright 2024-2026 the original author or authors. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package io.agentscope.core.memory.autocontext; + +import io.agentscope.core.agent.Agent; +import io.agentscope.core.agent.RuntimeContext; +import io.agentscope.core.message.Msg; +import io.agentscope.core.message.TextBlock; +import io.agentscope.core.tool.Tool; +import io.agentscope.core.tool.ToolParam; +import java.util.List; + +/** Tool for reloading offloaded context. */ +public class ContextOffloadTool { + + private final ContextOffLoader contextOffLoader; + private final AutoContextHook hook; + + public ContextOffloadTool(ContextOffLoader contextOffLoader) { + this.contextOffLoader = contextOffLoader; + this.hook = null; + } + + ContextOffloadTool(AutoContextHook hook) { + this.contextOffLoader = null; + this.hook = hook; + } + + public List reload(String uuid) { + return reloadInternal(uuid, null, null); + } + + @Tool( + name = "context_reload", + description = + "Reload previously offloaded context messages by UUID. Use this tool when you" + + " need to access the original tool invocation history that was compressed" + + " and stored.") + public List reload( + @ToolParam( + name = "working_context_offload_uuid", + description = "The UUID of the offloaded context to reload.") + String uuid, + Agent agent, + RuntimeContext runtimeContext) { + return reloadInternal(uuid, agent, runtimeContext); + } + + private List reloadInternal(String uuid, Agent agent, RuntimeContext runtimeContext) { + if (uuid == null || uuid.trim().isEmpty()) { + return List.of(errorMsg("Error: UUID cannot be null or empty.")); + } + + List messages; + if (hook != null) { + messages = hook.reload(uuid, agent, runtimeContext); + } else if (contextOffLoader != null) { + try { + messages = contextOffLoader.reload(uuid); + } catch (Exception e) { + return List.of( + errorMsg( + "Error reloading context with UUID " + + uuid + + ": " + + e.getMessage())); + } + } else { + return List.of( + errorMsg( + "Error: Context offloader is not available. Cannot reload context with" + + " UUID: " + + uuid)); + } + + if (messages == null || messages.isEmpty()) { + return List.of( + errorMsg( + "No messages found for UUID: " + + uuid + + ", The context may have been cleared or the UUID is" + + " invalid.")); + } + return messages; + } + + private Msg errorMsg(String text) { + return Msg.builder().content(TextBlock.builder().text(text).build()).build(); + } +} diff --git a/agentscope-extensions/agentscope-extensions-mem/agentscope-extensions-autocontext-memory/src/main/java/io/agentscope/core/memory/autocontext/MsgUtils.java b/agentscope-extensions/agentscope-extensions-mem/agentscope-extensions-autocontext-memory/src/main/java/io/agentscope/core/memory/autocontext/MsgUtils.java new file mode 100644 index 000000000..adb8c0113 --- /dev/null +++ b/agentscope-extensions/agentscope-extensions-mem/agentscope-extensions-autocontext-memory/src/main/java/io/agentscope/core/memory/autocontext/MsgUtils.java @@ -0,0 +1,378 @@ +/* + * Copyright 2024-2026 the original author or authors. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package io.agentscope.core.memory.autocontext; + +import com.fasterxml.jackson.core.type.TypeReference; +import io.agentscope.core.message.ContentBlock; +import io.agentscope.core.message.Msg; +import io.agentscope.core.message.MsgRole; +import io.agentscope.core.message.TextBlock; +import io.agentscope.core.message.ToolResultBlock; +import io.agentscope.core.message.ToolUseBlock; +import io.agentscope.core.util.JsonUtils; +import java.util.ArrayList; +import java.util.HashMap; +import java.util.HashSet; +import java.util.List; +import java.util.Map; +import java.util.Set; + +/** Misc message utilities used by auto context. */ +public final class MsgUtils { + + private static final Set PLAN_RELATED_TOOLS = + Set.of( + "create_plan", + "update_plan_info", + "revise_current_plan", + "update_subtask_state", + "finish_subtask", + "view_subtasks", + "get_subtask_count", + "finish_plan", + "view_historical_plans", + "recover_historical_plan"); + + private MsgUtils() {} + + public static Object serializeMsgList(Object messages) { + if (!(messages instanceof List list)) { + return messages; + } + List> serialized = new ArrayList<>(); + for (Object item : list) { + if (!(item instanceof Msg msg)) { + continue; + } + serialized.add( + JsonUtils.getJsonCodec() + .convertValue(msg, new TypeReference>() {})); + } + return serialized; + } + + public static Object deserializeToMsgList(Object data) { + if (!(data instanceof List list)) { + return data; + } + List restored = new ArrayList<>(); + for (Object item : list) { + if (!(item instanceof Map map)) { + continue; + } + restored.add(JsonUtils.getJsonCodec().convertValue(map, Msg.class)); + } + return restored; + } + + public static Object serializeMsgListMap(Object object) { + if (!(object instanceof Map map)) { + return object; + } + Map>> result = new HashMap<>(); + for (Map.Entry entry : map.entrySet()) { + if (entry.getKey() instanceof String key && entry.getValue() instanceof List list) { + @SuppressWarnings("unchecked") + List msgs = (List) list; + @SuppressWarnings("unchecked") + List> serialized = + (List>) serializeMsgList(msgs); + result.put(key, serialized); + } + } + return result; + } + + public static Object deserializeToMsgListMap(Object data) { + if (!(data instanceof Map map)) { + return data; + } + Map> result = new HashMap<>(); + for (Map.Entry entry : map.entrySet()) { + if (entry.getKey() instanceof String key) { + Object restored = deserializeToMsgList(entry.getValue()); + if (restored instanceof List list) { + @SuppressWarnings("unchecked") + List msgs = (List) list; + result.put(key, msgs); + } + } + } + return result; + } + + public static Object serializeCompressionEventList(Object object) { + if (!(object instanceof List list)) { + return object; + } + List> serialized = new ArrayList<>(); + for (Object item : list) { + if (!(item instanceof CompressionEvent event)) { + continue; + } + Map eventMap = new HashMap<>(); + eventMap.put("eventType", event.getEventType()); + eventMap.put("timestamp", event.getTimestamp()); + eventMap.put("compressedMessageCount", event.getCompressedMessageCount()); + eventMap.put("previousMessageId", event.getPreviousMessageId()); + eventMap.put("nextMessageId", event.getNextMessageId()); + eventMap.put("compressedMessageId", event.getCompressedMessageId()); + eventMap.put("metadata", event.getMetadata()); + serialized.add(eventMap); + } + return serialized; + } + + public static Object deserializeToCompressionEventList(Object data) { + if (!(data instanceof List list)) { + return data; + } + List result = new ArrayList<>(); + for (Object item : list) { + if (!(item instanceof Map raw)) { + continue; + } + Map map = new HashMap<>(); + for (Map.Entry entry : raw.entrySet()) { + if (entry.getKey() instanceof String key) { + map.put(key, entry.getValue()); + } + } + Map metadata = new HashMap<>(); + Object meta = map.get("metadata"); + if (meta instanceof Map metaMap) { + for (Map.Entry entry : metaMap.entrySet()) { + if (entry.getKey() instanceof String key) { + metadata.put(key, entry.getValue()); + } + } + } else { + if (map.containsKey("tokenBefore")) { + metadata.put("tokenBefore", map.get("tokenBefore")); + } + if (map.containsKey("tokenAfter")) { + metadata.put("tokenAfter", map.get("tokenAfter")); + } + if (map.containsKey("inputToken")) { + metadata.put("inputToken", map.get("inputToken")); + } + if (map.containsKey("outputToken")) { + metadata.put("outputToken", map.get("outputToken")); + } + if (map.containsKey("time")) { + metadata.put("time", map.get("time")); + } + } + result.add( + new CompressionEvent( + (String) map.get("eventType"), + ((Number) map.getOrDefault("timestamp", 0L)).longValue(), + ((Number) map.getOrDefault("compressedMessageCount", 0)).intValue(), + (String) map.get("previousMessageId"), + (String) map.get("nextMessageId"), + (String) map.get("compressedMessageId"), + metadata)); + } + return result; + } + + public static void replaceMsg(List rawMessages, int startIndex, int endIndex, Msg newMsg) { + if (rawMessages == null || newMsg == null || startIndex < 0 || endIndex < startIndex) { + return; + } + if (startIndex >= rawMessages.size()) { + return; + } + int actualEnd = Math.min(endIndex, rawMessages.size() - 1); + rawMessages.subList(startIndex, actualEnd + 1).clear(); + rawMessages.add(startIndex, newMsg); + } + + public static boolean isToolMessage(Msg msg) { + return msg != null + && (msg.getRole() == MsgRole.TOOL + || msg.hasContentBlocks(ToolUseBlock.class) + || msg.hasContentBlocks(ToolResultBlock.class)); + } + + public static boolean isToolUseMessage(Msg msg) { + return msg != null + && msg.getRole() == MsgRole.ASSISTANT + && msg.hasContentBlocks(ToolUseBlock.class); + } + + public static boolean isToolResultMessage(Msg msg) { + return msg != null + && (msg.getRole() == MsgRole.TOOL || msg.hasContentBlocks(ToolResultBlock.class)); + } + + public static boolean isCompressedMessage(Msg msg) { + if (msg == null || msg.getMetadata() == null) { + return false; + } + return msg.getMetadata().get("_compress_meta") instanceof Map; + } + + public static boolean isFinalAssistantResponse(Msg msg) { + if (msg == null || msg.getRole() != MsgRole.ASSISTANT) { + return false; + } + if (msg.getMetadata() != null) { + Object compressMeta = msg.getMetadata().get("_compress_meta"); + if (compressMeta instanceof Map meta + && Boolean.TRUE.equals(meta.get("compressed_current_round"))) { + return false; + } + } + return !msg.hasContentBlocks(ToolUseBlock.class) + && !msg.hasContentBlocks(ToolResultBlock.class); + } + + public static boolean isPlanRelatedTool(String toolName) { + return toolName != null && PLAN_RELATED_TOOLS.contains(toolName); + } + + public static boolean containsPlanRelatedToolCall(Msg msg) { + if (msg == null) { + return false; + } + List toolUseBlocks = msg.getContentBlocks(ToolUseBlock.class); + if (toolUseBlocks == null) { + return false; + } + for (ToolUseBlock toolUse : toolUseBlocks) { + if (toolUse != null && isPlanRelatedTool(toolUse.getName())) { + return true; + } + } + return false; + } + + public static List filterPlanRelatedToolCalls(List messages) { + if (messages == null || messages.isEmpty()) { + return messages; + } + List filtered = new ArrayList<>(); + Set planRelatedToolCallIds = new HashSet<>(); + for (Msg msg : messages) { + if (msg.getRole() != MsgRole.ASSISTANT) { + continue; + } + List toolUseBlocks = msg.getContentBlocks(ToolUseBlock.class); + if (toolUseBlocks == null) { + continue; + } + for (ToolUseBlock toolUse : toolUseBlocks) { + if (toolUse != null && isPlanRelatedTool(toolUse.getName())) { + planRelatedToolCallIds.add(toolUse.getId()); + } + } + } + for (Msg msg : messages) { + boolean shouldInclude = true; + if (msg.getRole() == MsgRole.ASSISTANT) { + List toolUseBlocks = msg.getContentBlocks(ToolUseBlock.class); + if (toolUseBlocks != null && !toolUseBlocks.isEmpty()) { + boolean allPlanRelated = true; + for (ToolUseBlock toolUse : toolUseBlocks) { + if (toolUse != null && !isPlanRelatedTool(toolUse.getName())) { + allPlanRelated = false; + break; + } + } + if (allPlanRelated) { + shouldInclude = false; + } + } + } + if (msg.getRole() == MsgRole.TOOL) { + List toolResultBlocks = + msg.getContentBlocks(ToolResultBlock.class); + if (toolResultBlocks != null) { + for (ToolResultBlock toolResult : toolResultBlocks) { + if (toolResult != null + && planRelatedToolCallIds.contains(toolResult.getId())) { + shouldInclude = false; + break; + } + } + } + } + if (shouldInclude) { + filtered.add(msg); + } + } + return filtered; + } + + public static int calculateMessageCharCount(Msg msg) { + if (msg == null || msg.getContent() == null) { + return 0; + } + int charCount = 0; + for (ContentBlock block : msg.getContent()) { + if (block instanceof TextBlock textBlock) { + if (textBlock.getText() != null) { + charCount += textBlock.getText().length(); + } + } else if (block instanceof ToolUseBlock toolUse) { + if (toolUse.getName() != null) { + charCount += toolUse.getName().length(); + } + if (toolUse.getId() != null) { + charCount += toolUse.getId().length(); + } + if (toolUse.getInput() != null && !toolUse.getInput().isEmpty()) { + try { + charCount += JsonUtils.getJsonCodec().toJson(toolUse.getInput()).length(); + } catch (Exception e) { + charCount += toolUse.getInput().toString().length(); + } + } + if (toolUse.getContent() != null) { + charCount += toolUse.getContent().length(); + } + } else if (block instanceof ToolResultBlock toolResult) { + if (toolResult.getName() != null) { + charCount += toolResult.getName().length(); + } + if (toolResult.getId() != null) { + charCount += toolResult.getId().length(); + } + if (toolResult.getOutput() != null) { + for (ContentBlock outputBlock : toolResult.getOutput()) { + if (outputBlock instanceof TextBlock textBlock + && textBlock.getText() != null) { + charCount += textBlock.getText().length(); + } + } + } + } + } + return charCount; + } + + public static int calculateMessagesCharCount(List messages) { + if (messages == null || messages.isEmpty()) { + return 0; + } + int total = 0; + for (Msg msg : messages) { + total += calculateMessageCharCount(msg); + } + return total; + } +} diff --git a/agentscope-extensions/agentscope-extensions-mem/agentscope-extensions-autocontext-memory/src/main/java/io/agentscope/core/memory/autocontext/PromptConfig.java b/agentscope-extensions/agentscope-extensions-mem/agentscope-extensions-autocontext-memory/src/main/java/io/agentscope/core/memory/autocontext/PromptConfig.java new file mode 100644 index 000000000..1b64b4447 --- /dev/null +++ b/agentscope-extensions/agentscope-extensions-mem/agentscope-extensions-autocontext-memory/src/main/java/io/agentscope/core/memory/autocontext/PromptConfig.java @@ -0,0 +1,77 @@ +/* + * Copyright 2024-2026 the original author or authors. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package io.agentscope.core.memory.autocontext; + +/** Optional prompt overrides for auto context compression. */ +public class PromptConfig { + private String previousRoundToolCompressPrompt; + private String previousRoundSummaryPrompt; + private String currentRoundLargeMessagePrompt; + private String currentRoundCompressPrompt; + + public static Builder builder() { + return new Builder(); + } + + public String getPreviousRoundToolCompressPrompt() { + return previousRoundToolCompressPrompt; + } + + public String getPreviousRoundSummaryPrompt() { + return previousRoundSummaryPrompt; + } + + public String getCurrentRoundLargeMessagePrompt() { + return currentRoundLargeMessagePrompt; + } + + public String getCurrentRoundCompressPrompt() { + return currentRoundCompressPrompt; + } + + public static class Builder { + private final PromptConfig config = new PromptConfig(); + + public Builder previousRoundToolCompressPrompt(String prompt) { + config.previousRoundToolCompressPrompt = prompt; + return this; + } + + public Builder previousRoundSummaryPrompt(String prompt) { + config.previousRoundSummaryPrompt = prompt; + return this; + } + + public Builder currentRoundLargeMessagePrompt(String prompt) { + config.currentRoundLargeMessagePrompt = prompt; + return this; + } + + public Builder currentRoundCompressPrompt(String prompt) { + config.currentRoundCompressPrompt = prompt; + return this; + } + + public PromptConfig build() { + PromptConfig result = new PromptConfig(); + result.previousRoundToolCompressPrompt = config.previousRoundToolCompressPrompt; + result.previousRoundSummaryPrompt = config.previousRoundSummaryPrompt; + result.currentRoundLargeMessagePrompt = config.currentRoundLargeMessagePrompt; + result.currentRoundCompressPrompt = config.currentRoundCompressPrompt; + return result; + } + } +} diff --git a/agentscope-extensions/agentscope-extensions-mem/agentscope-extensions-autocontext-memory/src/main/java/io/agentscope/core/memory/autocontext/PromptProvider.java b/agentscope-extensions/agentscope-extensions-mem/agentscope-extensions-autocontext-memory/src/main/java/io/agentscope/core/memory/autocontext/PromptProvider.java new file mode 100644 index 000000000..9d9f5d407 --- /dev/null +++ b/agentscope-extensions/agentscope-extensions-mem/agentscope-extensions-autocontext-memory/src/main/java/io/agentscope/core/memory/autocontext/PromptProvider.java @@ -0,0 +1,57 @@ +/* + * Copyright 2024-2026 the original author or authors. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package io.agentscope.core.memory.autocontext; + +/** Helper for prompt fallback resolution. */ +public final class PromptProvider { + private PromptProvider() {} + + public static String getPreviousRoundToolCompressPrompt(PromptConfig customPrompt) { + if (customPrompt != null + && customPrompt.getPreviousRoundToolCompressPrompt() != null + && !customPrompt.getPreviousRoundToolCompressPrompt().isBlank()) { + return customPrompt.getPreviousRoundToolCompressPrompt(); + } + return Prompts.PREVIOUS_ROUND_TOOL_INVOCATION_COMPRESS_PROMPT; + } + + public static String getPreviousRoundSummaryPrompt(PromptConfig customPrompt) { + if (customPrompt != null + && customPrompt.getPreviousRoundSummaryPrompt() != null + && !customPrompt.getPreviousRoundSummaryPrompt().isBlank()) { + return customPrompt.getPreviousRoundSummaryPrompt(); + } + return Prompts.PREVIOUS_ROUND_CONVERSATION_SUMMARY_PROMPT; + } + + public static String getCurrentRoundLargeMessagePrompt(PromptConfig customPrompt) { + if (customPrompt != null + && customPrompt.getCurrentRoundLargeMessagePrompt() != null + && !customPrompt.getCurrentRoundLargeMessagePrompt().isBlank()) { + return customPrompt.getCurrentRoundLargeMessagePrompt(); + } + return Prompts.CURRENT_ROUND_LARGE_MESSAGE_SUMMARY_PROMPT; + } + + public static String getCurrentRoundCompressPrompt(PromptConfig customPrompt) { + if (customPrompt != null + && customPrompt.getCurrentRoundCompressPrompt() != null + && !customPrompt.getCurrentRoundCompressPrompt().isBlank()) { + return customPrompt.getCurrentRoundCompressPrompt(); + } + return Prompts.CURRENT_ROUND_MESSAGE_COMPRESS_PROMPT; + } +} diff --git a/agentscope-extensions/agentscope-extensions-mem/agentscope-extensions-autocontext-memory/src/main/java/io/agentscope/core/memory/autocontext/Prompts.java b/agentscope-extensions/agentscope-extensions-mem/agentscope-extensions-autocontext-memory/src/main/java/io/agentscope/core/memory/autocontext/Prompts.java new file mode 100644 index 000000000..145346426 --- /dev/null +++ b/agentscope-extensions/agentscope-extensions-mem/agentscope-extensions-autocontext-memory/src/main/java/io/agentscope/core/memory/autocontext/Prompts.java @@ -0,0 +1,38 @@ +/* + * Copyright 2024-2026 the original author or authors. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package io.agentscope.core.memory.autocontext; + +/** Prompt templates used by auto context compression. */ +public final class Prompts { + private Prompts() {} + + public static final String COMPRESSION_MESSAGE_LIST_END = + "Above is the message list that needs to be compressed."; + + public static final String PREVIOUS_ROUND_TOOL_INVOCATION_COMPRESS_PROMPT = + "Compress the tool call history and keep tool names, arguments, and key results."; + + public static final String PREVIOUS_ROUND_CONVERSATION_SUMMARY_PROMPT = + "Rewrite the older conversation into a concise, factual summary."; + + public static final String CURRENT_ROUND_LARGE_MESSAGE_SUMMARY_PROMPT = + "Summarize this oversized message while keeping key facts."; + + public static final String CURRENT_ROUND_MESSAGE_COMPRESS_PROMPT = + "Compress the current round into a short, factual context."; + + public static final String CONTEXT_OFFLOAD_TAG_FORMAT = ""; +} diff --git a/agentscope-extensions/agentscope-extensions-mem/agentscope-extensions-autocontext-memory/src/main/java/io/agentscope/core/memory/autocontext/TokenCounterUtil.java b/agentscope-extensions/agentscope-extensions-mem/agentscope-extensions-autocontext-memory/src/main/java/io/agentscope/core/memory/autocontext/TokenCounterUtil.java new file mode 100644 index 000000000..4c56929ae --- /dev/null +++ b/agentscope-extensions/agentscope-extensions-mem/agentscope-extensions-autocontext-memory/src/main/java/io/agentscope/core/memory/autocontext/TokenCounterUtil.java @@ -0,0 +1,29 @@ +/* + * Copyright 2024-2026 the original author or authors. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package io.agentscope.core.memory.autocontext; + +import io.agentscope.core.message.Msg; +import java.util.List; + +/** Lightweight token estimator for compression triggers. */ +public final class TokenCounterUtil { + private TokenCounterUtil() {} + + public static int calculateToken(List messages) { + int chars = MsgUtils.calculateMessagesCharCount(messages); + return Math.max(1, chars / 4); + } +} diff --git a/agentscope-extensions/agentscope-extensions-mem/agentscope-extensions-autocontext-memory/src/test/java/io/agentscope/core/memory/autocontext/AutoContextHookTest.java b/agentscope-extensions/agentscope-extensions-mem/agentscope-extensions-autocontext-memory/src/test/java/io/agentscope/core/memory/autocontext/AutoContextHookTest.java new file mode 100644 index 000000000..1cc560955 --- /dev/null +++ b/agentscope-extensions/agentscope-extensions-mem/agentscope-extensions-autocontext-memory/src/test/java/io/agentscope/core/memory/autocontext/AutoContextHookTest.java @@ -0,0 +1,117 @@ +/* + * Copyright 2024-2026 the original author or authors. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package io.agentscope.core.memory.autocontext; + +import static org.junit.jupiter.api.Assertions.assertEquals; +import static org.junit.jupiter.api.Assertions.assertNotNull; +import static org.junit.jupiter.api.Assertions.assertSame; +import static org.junit.jupiter.api.Assertions.assertTrue; + +import io.agentscope.core.ReActAgent; +import io.agentscope.core.agent.RuntimeContext; +import io.agentscope.core.hook.PreCallEvent; +import io.agentscope.core.hook.PreReasoningEvent; +import io.agentscope.core.state.AgentState; +import java.lang.reflect.Field; +import java.util.List; +import java.util.concurrent.atomic.AtomicReference; +import org.junit.jupiter.api.Test; +import reactor.core.scheduler.Schedulers; +import reactor.test.StepVerifier; + +class AutoContextHookTest { + + @Test + void handlePreCallRegistersContextReloadTool() { + AutoContextHook hook = new AutoContextHook(); + ReActAgent agent = + ReActAgent.builder() + .name("hook-test") + .sysPrompt("system") + .model(AutoContextTestSupport.noopModel()) + .build(); + PreCallEvent event = + new PreCallEvent(agent, List.of(AutoContextTestSupport.userMessage("hello"))); + + StepVerifier.create(hook.handlePreCall(event)).expectNext(event).verifyComplete(); + + assertTrue(agent.getToolkit().getToolNames().contains("context_reload")); + assertEquals(1, event.getInputMessages().size()); + } + + @Test + void handlePreReasoningCompressesWithoutBlockingAndRewritesState() { + AtomicReference threadName = new AtomicReference<>(); + AutoContextHook hook = + new AutoContextHook( + AutoContextConfig.builder() + .msgThreshold(2) + .lastKeep(0) + .minCompressionTokenThreshold(1) + .build(), + AutoContextTestSupport.recordingModel("compressed", threadName)); + ReActAgent agent = + ReActAgent.builder() + .name("reasoning-test") + .sysPrompt("system") + .model(AutoContextTestSupport.noopModel()) + .build(); + + AgentState state = + AutoContextTestSupport.runtimeState( + "session-1", + "alice", + List.of( + AutoContextTestSupport.userMessage("first"), + AutoContextTestSupport.assistantMessage("second"))); + RuntimeContext runtimeContext = AutoContextTestSupport.runtimeContext(state); + setActiveRuntimeContext(agent, runtimeContext); + + PreReasoningEvent event = + new PreReasoningEvent(agent, "recording", null, state.getContext()); + + StepVerifier.create(hook.handlePreReasoning(event).subscribeOn(Schedulers.parallel())) + .assertNext( + returned -> { + assertSame(event, returned); + assertNotNull(returned.getSystemMessage()); + assertTrue( + returned.getSystemMessage() + .getTextContent() + .contains("context_reload")); + assertEquals(1, returned.getInputMessages().size()); + assertEquals( + "compressed", + returned.getInputMessages().get(0).getTextContent()); + assertEquals(1, state.getContext().size()); + assertEquals("compressed", state.getContext().get(0).getTextContent()); + }) + .verifyComplete(); + + assertNotNull(threadName.get()); + assertTrue(threadName.get().contains("boundedElastic")); + } + + private static void setActiveRuntimeContext(ReActAgent agent, RuntimeContext runtimeContext) { + try { + Field field = ReActAgent.class.getDeclaredField("activeRc"); + field.setAccessible(true); + field.set(agent, runtimeContext); + } catch (ReflectiveOperationException e) { + throw new AssertionError(e); + } + } +} diff --git a/agentscope-extensions/agentscope-extensions-mem/agentscope-extensions-autocontext-memory/src/test/java/io/agentscope/core/memory/autocontext/AutoContextMemoryTest.java b/agentscope-extensions/agentscope-extensions-mem/agentscope-extensions-autocontext-memory/src/test/java/io/agentscope/core/memory/autocontext/AutoContextMemoryTest.java new file mode 100644 index 000000000..005c0e539 --- /dev/null +++ b/agentscope-extensions/agentscope-extensions-mem/agentscope-extensions-autocontext-memory/src/test/java/io/agentscope/core/memory/autocontext/AutoContextMemoryTest.java @@ -0,0 +1,112 @@ +/* + * Copyright 2024-2026 the original author or authors. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package io.agentscope.core.memory.autocontext; + +import static org.junit.jupiter.api.Assertions.assertEquals; +import static org.junit.jupiter.api.Assertions.assertNotNull; +import static org.junit.jupiter.api.Assertions.assertTrue; + +import io.agentscope.core.message.Msg; +import java.util.List; +import java.util.concurrent.atomic.AtomicReference; +import org.junit.jupiter.api.Test; +import reactor.core.scheduler.Schedulers; +import reactor.test.StepVerifier; + +class AutoContextMemoryTest { + + @Test + void deleteMessageRemovesTheIndexedMessageFromBothBuffers() { + AutoContextMemory memory = new AutoContextMemory(AutoContextConfig.builder().build(), null); + Msg first = AutoContextTestSupport.userMessage("first"); + Msg second = AutoContextTestSupport.assistantMessage("second"); + Msg third = AutoContextTestSupport.userMessage("third"); + + memory.addMessage(first); + memory.addMessage(second); + memory.addMessage(third); + + memory.deleteMessage(1); + + assertEquals(2, memory.getMessages().size()); + assertEquals("first", memory.getMessages().get(0).getTextContent()); + assertEquals("third", memory.getMessages().get(1).getTextContent()); + assertEquals(2, memory.getOriginalMemoryMsgs().size()); + assertEquals("first", memory.getOriginalMemoryMsgs().get(0).getTextContent()); + assertEquals("third", memory.getOriginalMemoryMsgs().get(1).getTextContent()); + } + + @Test + void saveAndLoadRestoresSnapshotState() { + AutoContextMemory memory = new AutoContextMemory(AutoContextConfig.builder().build(), null); + memory.addMessage(AutoContextTestSupport.userMessage("hello")); + memory.addMessage(AutoContextTestSupport.assistantMessage("world")); + memory.offload("offload-1", List.of(AutoContextTestSupport.userMessage("offloaded"))); + memory.getCompressionEvents() + .add( + new CompressionEvent( + CompressionEvent.CURRENT_ROUND_MESSAGE_COMPRESS, + 123L, + 1, + null, + null, + "compressed-1", + java.util.Map.of("tokenBefore", 10, "tokenAfter", 4))); + + var store = AutoContextTestSupport.inMemoryStore(); + memory.saveTo(store, "alice", "session-1"); + + AutoContextMemory loaded = new AutoContextMemory(AutoContextConfig.builder().build(), null); + loaded.loadFrom(store, "alice", "session-1"); + + assertEquals(2, loaded.getMessages().size()); + assertEquals("hello", loaded.getMessages().get(0).getTextContent()); + assertEquals("world", loaded.getMessages().get(1).getTextContent()); + assertEquals(2, loaded.getOriginalMemoryMsgs().size()); + assertNotNull(loaded.reload("offload-1")); + assertEquals(1, loaded.reload("offload-1").size()); + assertEquals("offloaded", loaded.reload("offload-1").get(0).getTextContent()); + assertEquals(1, loaded.getCompressionEvents().size()); + assertEquals( + CompressionEvent.CURRENT_ROUND_MESSAGE_COMPRESS, + loaded.getCompressionEvents().get(0).getEventType()); + } + + @Test + void compressIfNeededAsyncRunsOnBoundedElasticAndSummarizesMessages() { + AtomicReference threadName = new AtomicReference<>(); + AutoContextMemory memory = + new AutoContextMemory( + AutoContextConfig.builder() + .msgThreshold(2) + .lastKeep(0) + .minCompressionTokenThreshold(1) + .build(), + AutoContextTestSupport.recordingModel("compressed", threadName)); + memory.addMessage(AutoContextTestSupport.userMessage("first message")); + memory.addMessage(AutoContextTestSupport.assistantMessage("second message")); + + StepVerifier.create(memory.compressIfNeededAsync().subscribeOn(Schedulers.parallel())) + .expectNext(true) + .verifyComplete(); + + assertNotNull(threadName.get()); + assertTrue(threadName.get().contains("boundedElastic")); + assertEquals(1, memory.getMessages().size()); + assertEquals("compressed", memory.getMessages().get(0).getTextContent()); + assertEquals(1, memory.getCompressionEvents().size()); + } +} diff --git a/agentscope-extensions/agentscope-extensions-mem/agentscope-extensions-autocontext-memory/src/test/java/io/agentscope/core/memory/autocontext/AutoContextTestSupport.java b/agentscope-extensions/agentscope-extensions-mem/agentscope-extensions-autocontext-memory/src/test/java/io/agentscope/core/memory/autocontext/AutoContextTestSupport.java new file mode 100644 index 000000000..089bae9ab --- /dev/null +++ b/agentscope-extensions/agentscope-extensions-mem/agentscope-extensions-autocontext-memory/src/test/java/io/agentscope/core/memory/autocontext/AutoContextTestSupport.java @@ -0,0 +1,194 @@ +/* + * Copyright 2024-2026 the original author or authors. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package io.agentscope.core.memory.autocontext; + +import io.agentscope.core.agent.RuntimeContext; +import io.agentscope.core.message.Msg; +import io.agentscope.core.message.MsgRole; +import io.agentscope.core.message.TextBlock; +import io.agentscope.core.model.ChatResponse; +import io.agentscope.core.model.GenerateOptions; +import io.agentscope.core.model.Model; +import io.agentscope.core.model.ToolSchema; +import io.agentscope.core.state.AgentState; +import io.agentscope.core.state.AgentStateStore; +import io.agentscope.core.state.State; +import java.util.ArrayList; +import java.util.List; +import java.util.Map; +import java.util.Optional; +import java.util.Set; +import java.util.concurrent.ConcurrentHashMap; +import java.util.concurrent.atomic.AtomicReference; +import reactor.core.publisher.Flux; + +final class AutoContextTestSupport { + + private AutoContextTestSupport() {} + + static Msg userMessage(String text) { + return Msg.builder() + .role(MsgRole.USER) + .name("user") + .content(TextBlock.builder().text(text).build()) + .build(); + } + + static Msg assistantMessage(String text) { + return Msg.builder() + .role(MsgRole.ASSISTANT) + .name("assistant") + .content(TextBlock.builder().text(text).build()) + .build(); + } + + static Model noopModel() { + return new Model() { + @Override + public Flux stream( + List messages, List tools, GenerateOptions options) { + return Flux.empty(); + } + + @Override + public String getModelName() { + return "noop"; + } + }; + } + + static Model recordingModel(String responseText, AtomicReference threadNameRef) { + return new Model() { + @Override + public Flux stream( + List messages, List tools, GenerateOptions options) { + threadNameRef.set(Thread.currentThread().getName()); + return Flux.just( + ChatResponse.builder() + .content(List.of(TextBlock.builder().text(responseText).build())) + .build()); + } + + @Override + public String getModelName() { + return "recording"; + } + }; + } + + static AgentState runtimeState(String sessionId, String userId, List messages) { + return AgentState.builder().sessionId(sessionId).userId(userId).context(messages).build(); + } + + static RuntimeContext runtimeContext(AgentState state) { + return RuntimeContext.builder() + .sessionId(state.getSessionId()) + .userId(state.getUserId()) + .agentState(state) + .build(); + } + + static AgentStateStore inMemoryStore() { + return new InMemoryAgentStateStore(); + } + + private static final class InMemoryAgentStateStore implements AgentStateStore { + + private final Map> data = new ConcurrentHashMap<>(); + + @Override + public void save(String userId, String sessionId, String key, State value) { + slot(userId, sessionId).put(key, value); + } + + @Override + public void save( + String userId, String sessionId, String key, List values) { + slot(userId, sessionId).put(key, new ArrayList<>(values)); + } + + @Override + public Optional get( + String userId, String sessionId, String key, Class type) { + Object value = slot(userId, sessionId).get(key); + if (type.isInstance(value)) { + return Optional.of(type.cast(value)); + } + return Optional.empty(); + } + + @Override + @SuppressWarnings("unchecked") + public List getList( + String userId, String sessionId, String key, Class itemType) { + Object value = slot(userId, sessionId).get(key); + if (!(value instanceof List list)) { + return List.of(); + } + List result = new ArrayList<>(); + for (Object item : list) { + if (itemType.isInstance(item)) { + result.add((T) item); + } + } + return result; + } + + @Override + public boolean exists(String userId, String sessionId) { + return data.containsKey(slotKey(userId, sessionId)); + } + + @Override + public void delete(String userId, String sessionId) { + data.remove(slotKey(userId, sessionId)); + } + + @Override + public void delete(String userId, String sessionId, String key) { + Map slot = data.get(slotKey(userId, sessionId)); + if (slot == null) { + return; + } + slot.remove(key); + if (slot.isEmpty()) { + data.remove(slotKey(userId, sessionId)); + } + } + + @Override + public Set listSessionIds(String userId) { + String prefix = slotKey(userId, ""); + Set sessionIds = new java.util.HashSet<>(); + for (String key : data.keySet()) { + if (key.startsWith(prefix)) { + sessionIds.add(key.substring(prefix.length())); + } + } + return sessionIds; + } + + private Map slot(String userId, String sessionId) { + return data.computeIfAbsent( + slotKey(userId, sessionId), ignored -> new ConcurrentHashMap<>()); + } + + private String slotKey(String userId, String sessionId) { + String uid = userId == null || userId.isBlank() ? "__anon__" : userId; + return uid + "|" + sessionId; + } + } +} diff --git a/agentscope-extensions/agentscope-extensions-mem/agentscope-extensions-autocontext-memory/src/test/java/io/agentscope/core/memory/autocontext/ContextOffloadToolTest.java b/agentscope-extensions/agentscope-extensions-mem/agentscope-extensions-autocontext-memory/src/test/java/io/agentscope/core/memory/autocontext/ContextOffloadToolTest.java new file mode 100644 index 000000000..97fd97158 --- /dev/null +++ b/agentscope-extensions/agentscope-extensions-mem/agentscope-extensions-autocontext-memory/src/test/java/io/agentscope/core/memory/autocontext/ContextOffloadToolTest.java @@ -0,0 +1,88 @@ +/* + * Copyright 2024-2026 the original author or authors. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package io.agentscope.core.memory.autocontext; + +import static org.junit.jupiter.api.Assertions.assertEquals; +import static org.junit.jupiter.api.Assertions.assertTrue; + +import io.agentscope.core.ReActAgent; +import io.agentscope.core.agent.RuntimeContext; +import io.agentscope.core.message.Msg; +import io.agentscope.core.state.AgentState; +import java.util.List; +import org.junit.jupiter.api.Test; + +class ContextOffloadToolTest { + + @Test + void reloadUsesDirectContextOffLoaderWhenProvided() { + ContextOffloadTool tool = + new ContextOffloadTool( + new ContextOffLoader() { + @Override + public void offload(String uuid, List messages) {} + + @Override + public List reload(String uuid) { + return List.of( + AutoContextTestSupport.userMessage("direct-one"), + AutoContextTestSupport.assistantMessage("direct-two")); + } + + @Override + public void clear(String uuid) {} + }); + + List messages = tool.reload("direct-uuid"); + + assertEquals(2, messages.size()); + assertEquals("direct-one", messages.get(0).getTextContent()); + assertEquals("direct-two", messages.get(1).getTextContent()); + } + + @Test + void reloadUsesHookRuntimeContextWhenAgentAndStateAreAvailable() { + AutoContextHook hook = new AutoContextHook(); + ContextOffloadTool tool = new ContextOffloadTool(hook); + ReActAgent agent = + ReActAgent.builder() + .name("offload-test") + .sysPrompt("system") + .model(AutoContextTestSupport.noopModel()) + .build(); + + AgentState state = + AutoContextTestSupport.runtimeState( + "session-1", + "alice", + List.of(AutoContextTestSupport.userMessage("context"))); + RuntimeContext runtimeContext = AutoContextTestSupport.runtimeContext(state); + String uuid = "offload-uuid"; + + hook.memoryFor(agent, runtimeContext) + .offload( + uuid, + List.of( + AutoContextTestSupport.userMessage("offloaded-one"), + AutoContextTestSupport.assistantMessage("offloaded-two"))); + + List messages = tool.reload(uuid, agent, runtimeContext); + + assertEquals(2, messages.size()); + assertTrue(messages.get(0).getTextContent().contains("offloaded-one")); + assertTrue(messages.get(1).getTextContent().contains("offloaded-two")); + } +} diff --git a/agentscope-extensions/agentscope-extensions-mem/pom.xml b/agentscope-extensions/agentscope-extensions-mem/pom.xml index c3aa7400e..c757855d7 100644 --- a/agentscope-extensions/agentscope-extensions-mem/pom.xml +++ b/agentscope-extensions/agentscope-extensions-mem/pom.xml @@ -33,6 +33,7 @@ agentscope-extensions-mem0 + agentscope-extensions-autocontext-memory agentscope-extensions-memory-bailian agentscope-extensions-reme From 0323696a0137ec1422818f18f1ad2a5b9264eefb Mon Sep 17 00:00:00 2001 From: guslegend <1670547022@qq.com> Date: Wed, 10 Jun 2026 07:28:58 +0800 Subject: [PATCH 2/4] fix: sync autocontext module into distribution poms --- agentscope-distribution/agentscope-all/pom.xml | 7 +++++++ agentscope-distribution/agentscope-bom/pom.xml | 7 +++++++ 2 files changed, 14 insertions(+) diff --git a/agentscope-distribution/agentscope-all/pom.xml b/agentscope-distribution/agentscope-all/pom.xml index 7367cf3d8..b8581641f 100644 --- a/agentscope-distribution/agentscope-all/pom.xml +++ b/agentscope-distribution/agentscope-all/pom.xml @@ -115,6 +115,13 @@ true + + io.agentscope + agentscope-extensions-autocontext-memory + compile + true + + io.agentscope agentscope-extensions-memory-bailian diff --git a/agentscope-distribution/agentscope-bom/pom.xml b/agentscope-distribution/agentscope-bom/pom.xml index 55eddcd89..91a27b7b7 100644 --- a/agentscope-distribution/agentscope-bom/pom.xml +++ b/agentscope-distribution/agentscope-bom/pom.xml @@ -148,6 +148,13 @@ ${project.version} + + + io.agentscope + agentscope-extensions-autocontext-memory + ${project.version} + + io.agentscope From 5197bb29c5d47202fe54f6dbbbe23f5c21a87e88 Mon Sep 17 00:00:00 2001 From: guslegend <1670547022@qq.com> Date: Wed, 10 Jun 2026 16:18:48 +0800 Subject: [PATCH 3/4] fix(autocontext): address review feedback on compression flow --- .../agentscope-bom/pom.xml | 2 +- .../memory/autocontext/AutoContextConfig.java | 37 ++ .../memory/autocontext/AutoContextHook.java | 6 + .../memory/autocontext/AutoContextMemory.java | 327 +++++++++++------- .../memory/autocontext/TokenCounterUtil.java | 8 +- .../autocontext/AutoContextMemoryTest.java | 127 ++++++- 6 files changed, 380 insertions(+), 127 deletions(-) diff --git a/agentscope-distribution/agentscope-bom/pom.xml b/agentscope-distribution/agentscope-bom/pom.xml index 91a27b7b7..d2dabb7eb 100644 --- a/agentscope-distribution/agentscope-bom/pom.xml +++ b/agentscope-distribution/agentscope-bom/pom.xml @@ -148,7 +148,7 @@ ${project.version} - + io.agentscope agentscope-extensions-autocontext-memory diff --git a/agentscope-extensions/agentscope-extensions-mem/agentscope-extensions-autocontext-memory/src/main/java/io/agentscope/core/memory/autocontext/AutoContextConfig.java b/agentscope-extensions/agentscope-extensions-mem/agentscope-extensions-autocontext-memory/src/main/java/io/agentscope/core/memory/autocontext/AutoContextConfig.java index 81b6c59de..831a3c0b8 100644 --- a/agentscope-extensions/agentscope-extensions-mem/agentscope-extensions-autocontext-memory/src/main/java/io/agentscope/core/memory/autocontext/AutoContextConfig.java +++ b/agentscope-extensions/agentscope-extensions-mem/agentscope-extensions-autocontext-memory/src/main/java/io/agentscope/core/memory/autocontext/AutoContextConfig.java @@ -127,6 +127,7 @@ public Builder customPrompt(PromptConfig customPrompt) { } public AutoContextConfig build() { + validate(); AutoContextConfig result = new AutoContextConfig(); result.largePayloadThreshold = config.largePayloadThreshold; result.maxToken = config.maxToken; @@ -140,5 +141,41 @@ public AutoContextConfig build() { result.customPrompt = config.customPrompt; return result; } + + private void validate() { + requirePositive(config.largePayloadThreshold, "largePayloadThreshold"); + requirePositive(config.maxToken, "maxToken"); + requireRatio(config.tokenRatio, "tokenRatio"); + requireNonNegative(config.offloadSinglePreview, "offloadSinglePreview"); + requirePositive(config.msgThreshold, "msgThreshold"); + requireNonNegative(config.lastKeep, "lastKeep"); + requirePositive(config.minConsecutiveToolMessages, "minConsecutiveToolMessages"); + requireRatio(config.currentRoundCompressionRatio, "currentRoundCompressionRatio"); + requirePositive(config.minCompressionTokenThreshold, "minCompressionTokenThreshold"); + } + + private void requirePositive(long value, String fieldName) { + if (value <= 0) { + throw new IllegalArgumentException(fieldName + " must be positive"); + } + } + + private void requirePositive(int value, String fieldName) { + if (value <= 0) { + throw new IllegalArgumentException(fieldName + " must be positive"); + } + } + + private void requireNonNegative(int value, String fieldName) { + if (value < 0) { + throw new IllegalArgumentException(fieldName + " must be non-negative"); + } + } + + private void requireRatio(double value, String fieldName) { + if (value <= 0.0 || value > 1.0) { + throw new IllegalArgumentException(fieldName + " must be in (0.0, 1.0]"); + } + } } } diff --git a/agentscope-extensions/agentscope-extensions-mem/agentscope-extensions-autocontext-memory/src/main/java/io/agentscope/core/memory/autocontext/AutoContextHook.java b/agentscope-extensions/agentscope-extensions-mem/agentscope-extensions-autocontext-memory/src/main/java/io/agentscope/core/memory/autocontext/AutoContextHook.java index df7912eb2..18d0e936f 100644 --- a/agentscope-extensions/agentscope-extensions-mem/agentscope-extensions-autocontext-memory/src/main/java/io/agentscope/core/memory/autocontext/AutoContextHook.java +++ b/agentscope-extensions/agentscope-extensions-mem/agentscope-extensions-autocontext-memory/src/main/java/io/agentscope/core/memory/autocontext/AutoContextHook.java @@ -31,6 +31,8 @@ import java.util.Objects; import java.util.concurrent.ConcurrentHashMap; import java.util.concurrent.atomic.AtomicBoolean; +import org.slf4j.Logger; +import org.slf4j.LoggerFactory; import reactor.core.publisher.Mono; import reactor.core.scheduler.Schedulers; @@ -38,6 +40,8 @@ @SuppressWarnings("deprecation") public class AutoContextHook implements Hook { + private static final Logger log = LoggerFactory.getLogger(AutoContextHook.class); + private final AutoContextConfig config; private final io.agentscope.core.model.Model model; private final String hookId; @@ -96,6 +100,8 @@ Mono handlePreCall(PreCallEvent event) { } } catch (Exception e) { registrationFailed.set(true); + log.warn( + "Failed to register context_reload tool for agent {}", reactAgent.getName(), e); } return Mono.just(event); } diff --git a/agentscope-extensions/agentscope-extensions-mem/agentscope-extensions-autocontext-memory/src/main/java/io/agentscope/core/memory/autocontext/AutoContextMemory.java b/agentscope-extensions/agentscope-extensions-mem/agentscope-extensions-autocontext-memory/src/main/java/io/agentscope/core/memory/autocontext/AutoContextMemory.java index 480cfaef6..ba8a1cdce 100644 --- a/agentscope-extensions/agentscope-extensions-mem/agentscope-extensions-autocontext-memory/src/main/java/io/agentscope/core/memory/autocontext/AutoContextMemory.java +++ b/agentscope-extensions/agentscope-extensions-mem/agentscope-extensions-autocontext-memory/src/main/java/io/agentscope/core/memory/autocontext/AutoContextMemory.java @@ -42,7 +42,7 @@ public class AutoContextMemory implements Memory, ContextOffLoader { private final AutoContextConfig autoContextConfig; private final PromptConfig customPrompt; - private Model model; + private volatile Model model; private final List workingMemoryStorage = new ArrayList<>(); private final List originalMemoryStorage = new ArrayList<>(); private final Map> offloadContext = new HashMap<>(); @@ -62,6 +62,15 @@ public synchronized AutoContextConfig getAutoContextConfig() { return autoContextConfig; } + /** + * Reconciles runtime context back into the working buffer. + * + *

The incoming context is expected to be either the same session prefix plus newly appended + * tail messages, or a rebuilt working context after a reset/reload. When the prefix still + * matches, only the new tail is appended into both working/original buffers; when it no longer + * matches, the working buffer is replaced and the original history is seeded only if it was + * still empty. + */ public synchronized void mergeWithContext(List context) { if (context == null || context.isEmpty()) { return; @@ -165,9 +174,6 @@ public synchronized void deleteMessage(int index) { return; } workingMemoryStorage.remove(index); - if (index < originalMemoryStorage.size()) { - originalMemoryStorage.remove(index); - } } public synchronized List getOriginalMemoryMsgs() { @@ -192,24 +198,8 @@ public synchronized List getCompressionEvents() { return compressionEvents; } - public synchronized boolean compressIfNeeded() { - if (workingMemoryStorage.isEmpty()) { - return false; - } - - int tokenCount = TokenCounterUtil.calculateToken(workingMemoryStorage); - boolean msgCountReached = - workingMemoryStorage.size() >= autoContextConfig.getMsgThreshold(); - boolean tokenReached = - tokenCount - >= (int) - (autoContextConfig.getMaxToken() - * autoContextConfig.getTokenRatio()); - if (!msgCountReached && !tokenReached) { - return false; - } - - if (tokenCount < autoContextConfig.getMinCompressionTokenThreshold()) { + public boolean compressIfNeeded() { + if (!isCompressionTriggered()) { return false; } @@ -263,93 +253,72 @@ public synchronized void setWorkingMessages(List messages) { workingMemoryStorage.addAll(copyMessages(messages)); } + private synchronized boolean isCompressionTriggered() { + if (workingMemoryStorage.isEmpty()) { + return false; + } + int tokenCount = TokenCounterUtil.calculateToken(workingMemoryStorage); + boolean msgCountReached = + workingMemoryStorage.size() >= autoContextConfig.getMsgThreshold(); + boolean tokenReached = + tokenCount + >= (int) + (autoContextConfig.getMaxToken() + * autoContextConfig.getTokenRatio()); + if (!msgCountReached && !tokenReached) { + return false; + } + return tokenCount >= autoContextConfig.getMinCompressionTokenThreshold(); + } + private boolean compressToolGroups() { - int cursor = 0; - int upperLimit = Math.max(0, workingMemoryStorage.size() - autoContextConfig.getLastKeep()); boolean changed = false; - while (cursor < upperLimit) { - IntRange range = findToolGroup(cursor, upperLimit); - if (range == null) { + int cursor = 0; + while (true) { + CompressionCandidate candidate = planToolGroupCompression(cursor); + if (candidate == null) { break; } - List messages = - new ArrayList<>(workingMemoryStorage.subList(range.start, range.end + 1)); - if (TokenCounterUtil.calculateToken(messages) + if (TokenCounterUtil.calculateToken(candidate.source()) < autoContextConfig.getMinCompressionTokenThreshold()) { - cursor = range.end + 1; + cursor = candidate.end() + 1; continue; } - if (compressRange( - range.start, - range.end, - CompressionEvent.TOOL_INVOCATION_COMPRESS, - PromptProvider.getPreviousRoundToolCompressPrompt(customPrompt), - true)) { + String summary = + summarizeMessages( + candidate.source(), candidate.prompt(), candidate.currentRound()); + if (applyCompression(candidate, summary)) { changed = true; - cursor = range.start + 1; - upperLimit = - Math.max(0, workingMemoryStorage.size() - autoContextConfig.getLastKeep()); + cursor = candidate.start() + 1; } else { - cursor = range.end + 1; + cursor = candidate.end() + 1; } } return changed; } private boolean offloadLargeMessages() { - int limit = Math.max(0, workingMemoryStorage.size() - autoContextConfig.getLastKeep()); - for (int i = 0; i < limit; i++) { - Msg msg = workingMemoryStorage.get(i); - if (MsgUtils.calculateMessageCharCount(msg) - < autoContextConfig.getLargePayloadThreshold()) { - continue; - } - String uuid = UUID.randomUUID().toString(); - List original = List.of(msg); - String summary = - summarizeMessages( - original, - PromptProvider.getCurrentRoundLargeMessagePrompt(customPrompt), - false); - if (summary == null || summary.isBlank()) { - summary = fallbackSummary(original); - } - offload(uuid, original); - Msg compressed = compressSingleMessage(msg, summary, uuid); - MsgUtils.replaceMsg(workingMemoryStorage, i, i, compressed); - recordCompressionEvent( - msg.getRole() == MsgRole.TOOL - ? CompressionEvent.LARGE_MESSAGE_OFFLOAD - : CompressionEvent.LARGE_MESSAGE_OFFLOAD_WITH_PROTECTION, - i, - i, - original, - compressed, - buildCompressionMetadata(compressed)); - return true; + CompressionCandidate candidate = planLargeMessageOffload(); + if (candidate == null) { + return false; } - return false; + String summary = + summarizeMessages(candidate.source(), candidate.prompt(), candidate.currentRound()); + return applyCompression(candidate, summary); } private boolean compressPreviousRounds() { boolean changed = false; int guard = 0; - while (workingMemoryStorage.size() > autoContextConfig.getLastKeep() + 1 && guard++ < 32) { - int end = findPrefixCompressionEnd(); - if (end < 0) { + while (guard++ < 32) { + CompressionCandidate candidate = planPreviousRoundCompression(); + if (candidate == null) { break; } - List source = new ArrayList<>(workingMemoryStorage.subList(0, end + 1)); - if (TokenCounterUtil.calculateToken(source) - < autoContextConfig.getMinCompressionTokenThreshold()) { - break; - } - if (!compressRange( - 0, - end, - CompressionEvent.PREVIOUS_ROUND_CONVERSATION_SUMMARY, - PromptProvider.getPreviousRoundSummaryPrompt(customPrompt), - false)) { + String summary = + summarizeMessages( + candidate.source(), candidate.prompt(), candidate.currentRound()); + if (!applyCompression(candidate, summary)) { break; } changed = true; @@ -358,6 +327,85 @@ private boolean compressPreviousRounds() { } private boolean compressCurrentRound() { + CompressionCandidate candidate = planCurrentRoundCompression(); + if (candidate == null) { + return false; + } + String summary = + summarizeMessages(candidate.source(), candidate.prompt(), candidate.currentRound()); + return applyCompression(candidate, summary); + } + + private synchronized CompressionCandidate planToolGroupCompression(int cursor) { + int upperLimit = Math.max(0, workingMemoryStorage.size() - autoContextConfig.getLastKeep()); + if (cursor >= upperLimit) { + return null; + } + IntRange range = findToolGroup(cursor, upperLimit); + if (range == null) { + return null; + } + List source = + new ArrayList<>(workingMemoryStorage.subList(range.start, range.end + 1)); + return new CompressionCandidate( + range.start, + range.end, + CompressionEvent.TOOL_INVOCATION_COMPRESS, + PromptProvider.getPreviousRoundToolCompressPrompt(customPrompt), + source, + false, + true, + false); + } + + private synchronized CompressionCandidate planLargeMessageOffload() { + int limit = Math.max(0, workingMemoryStorage.size() - autoContextConfig.getLastKeep()); + for (int i = 0; i < limit; i++) { + Msg msg = workingMemoryStorage.get(i); + if (MsgUtils.calculateMessageCharCount(msg) + < autoContextConfig.getLargePayloadThreshold()) { + continue; + } + return new CompressionCandidate( + i, + i, + msg.getRole() == MsgRole.TOOL + ? CompressionEvent.LARGE_MESSAGE_OFFLOAD + : CompressionEvent.LARGE_MESSAGE_OFFLOAD_WITH_PROTECTION, + PromptProvider.getCurrentRoundLargeMessagePrompt(customPrompt), + List.of(msg), + false, + false, + true); + } + return null; + } + + private synchronized CompressionCandidate planPreviousRoundCompression() { + if (workingMemoryStorage.size() <= autoContextConfig.getLastKeep() + 1) { + return null; + } + int end = findPrefixCompressionEnd(); + if (end < 0) { + return null; + } + List source = new ArrayList<>(workingMemoryStorage.subList(0, end + 1)); + if (TokenCounterUtil.calculateToken(source) + < autoContextConfig.getMinCompressionTokenThreshold()) { + return null; + } + return new CompressionCandidate( + 0, + end, + CompressionEvent.PREVIOUS_ROUND_CONVERSATION_SUMMARY, + PromptProvider.getPreviousRoundSummaryPrompt(customPrompt), + source, + false, + false, + false); + } + + private synchronized CompressionCandidate planCurrentRoundCompression() { int latestUserIndex = -1; for (int i = workingMemoryStorage.size() - 1; i >= 0; i--) { if (workingMemoryStorage.get(i).getRole() == MsgRole.USER) { @@ -366,52 +414,64 @@ private boolean compressCurrentRound() { } } if (latestUserIndex < 0 || latestUserIndex >= workingMemoryStorage.size() - 1) { - return false; + return null; } - int end = workingMemoryStorage.size() - 1; if (MsgUtils.isToolUseMessage(workingMemoryStorage.get(end))) { end--; } if (end <= latestUserIndex) { - return false; + return null; } - List source = new ArrayList<>(workingMemoryStorage.subList(latestUserIndex + 1, end + 1)); if (TokenCounterUtil.calculateToken(source) < autoContextConfig.getMinCompressionTokenThreshold()) { - return false; + return null; } - return compressRange( + return new CompressionCandidate( latestUserIndex + 1, end, CompressionEvent.CURRENT_ROUND_MESSAGE_COMPRESS, PromptProvider.getCurrentRoundCompressPrompt(customPrompt), - true); + source, + true, + true, + false); } - private boolean compressRange( - int start, int end, String eventType, String prompt, boolean appendOffloadTag) { - if (start < 0 || end < start || start >= workingMemoryStorage.size()) { + private boolean applyCompression(CompressionCandidate candidate, String summary) { + if (candidate == null) { return false; } - end = Math.min(end, workingMemoryStorage.size() - 1); - List source = new ArrayList<>(workingMemoryStorage.subList(start, end + 1)); + String resolvedSummary = + summary == null || summary.isBlank() + ? fallbackSummary(candidate.source()) + : summary.trim(); String uuid = UUID.randomUUID().toString(); - offload(uuid, source); - String summary = - summarizeMessages( - source, - prompt, - eventType.equals(CompressionEvent.CURRENT_ROUND_MESSAGE_COMPRESS)); - if (summary == null || summary.isBlank()) { - summary = fallbackSummary(source); - } - Msg summaryMsg = buildSummaryMessage(source, summary, uuid, appendOffloadTag); - MsgUtils.replaceMsg(workingMemoryStorage, start, end, summaryMsg); - recordCompressionEvent( - eventType, start, end, source, summaryMsg, buildCompressionMetadata(summaryMsg)); + Msg compressedMessage = + candidate.singleMessageOffload() + ? compressSingleMessage(candidate.source().get(0), resolvedSummary, uuid) + : buildSummaryMessage( + candidate.source(), + resolvedSummary, + uuid, + candidate.appendOffloadTag()); + synchronized (this) { + if (!matchesCandidate(candidate)) { + return false; + } + offloadContext.put(uuid, copyMessages(candidate.source())); + MsgUtils.replaceMsg( + workingMemoryStorage, candidate.start(), candidate.end(), compressedMessage); + recordCompressionEvent( + candidate.eventType(), + candidate.start(), + candidate.end(), + candidate.source(), + compressedMessage, + buildCompressionMetadata(compressedMessage)); + } return true; } @@ -419,7 +479,8 @@ private String summarizeMessages(List messages, String prompt, boolean curr if (messages == null || messages.isEmpty()) { return null; } - if (model == null) { + Model currentModel = model; + if (currentModel == null) { return fallbackSummary(messages); } try { @@ -465,7 +526,7 @@ private String summarizeMessages(List messages, String prompt, boolean curr .build()); } List responses = - model.stream(input, null, GenerateOptions.builder().build()) + currentModel.stream(input, null, GenerateOptions.builder().build()) .collectList() .block(); if (responses == null || responses.isEmpty()) { @@ -537,7 +598,10 @@ private Msg compressSingleMessage(Msg message, String summary, String uuid) { if (!summaryInserted) { content.add(TextBlock.builder().text(summary).build()); } - Map metadata = new HashMap<>(message.getMetadata()); + Map metadata = + message.getMetadata() == null + ? new HashMap<>() + : new HashMap<>(message.getMetadata()); Map compressMeta = new HashMap<>(); compressMeta.put("offloaduuid", uuid); metadata.put("_compress_meta", compressMeta); @@ -580,22 +644,29 @@ private void recordCompressionEvent( eventType, System.currentTimeMillis(), Math.max(1, endIndex - startIndex + 1), - startIndex > 0 - ? workingMemoryStorage - .get( - Math.min( - startIndex - 1, - workingMemoryStorage.size() - 1)) - .getId() + startIndex > 0 && startIndex - 1 < workingMemoryStorage.size() + ? workingMemoryStorage.get(startIndex - 1).getId() : null, - endIndex < workingMemoryStorage.size() - 1 - ? workingMemoryStorage.get(endIndex + 1).getId() + startIndex < workingMemoryStorage.size() - 1 + ? workingMemoryStorage.get(startIndex + 1).getId() : null, compressedMessage != null ? compressedMessage.getId() : null, metadata); compressionEvents.add(event); } + private boolean matchesCandidate(CompressionCandidate candidate) { + if (candidate.start() < 0 + || candidate.end() < candidate.start() + || candidate.end() >= workingMemoryStorage.size()) { + return false; + } + List currentSlice = + new ArrayList<>( + workingMemoryStorage.subList(candidate.start(), candidate.end() + 1)); + return currentSlice.equals(candidate.source()); + } + private IntRange findToolGroup(int startIndex, int upperLimit) { int start = -1; for (int i = startIndex; i < upperLimit; i++) { @@ -720,4 +791,14 @@ private Map> copyMsgMap(Map> input) { } private record IntRange(int start, int end) {} + + private record CompressionCandidate( + int start, + int end, + String eventType, + String prompt, + List source, + boolean currentRound, + boolean appendOffloadTag, + boolean singleMessageOffload) {} } diff --git a/agentscope-extensions/agentscope-extensions-mem/agentscope-extensions-autocontext-memory/src/main/java/io/agentscope/core/memory/autocontext/TokenCounterUtil.java b/agentscope-extensions/agentscope-extensions-mem/agentscope-extensions-autocontext-memory/src/main/java/io/agentscope/core/memory/autocontext/TokenCounterUtil.java index 4c56929ae..48d51e44a 100644 --- a/agentscope-extensions/agentscope-extensions-mem/agentscope-extensions-autocontext-memory/src/main/java/io/agentscope/core/memory/autocontext/TokenCounterUtil.java +++ b/agentscope-extensions/agentscope-extensions-mem/agentscope-extensions-autocontext-memory/src/main/java/io/agentscope/core/memory/autocontext/TokenCounterUtil.java @@ -18,7 +18,13 @@ import io.agentscope.core.message.Msg; import java.util.List; -/** Lightweight token estimator for compression triggers. */ +/** + * Lightweight token estimator for compression triggers. + * + *

This uses a rough {@code chars / 4} heuristic, which is usually reasonable for English but + * can undercount CJK-heavy content where token counts are often materially higher. Replace with a + * model-specific tokenizer in production-sensitive paths. + */ public final class TokenCounterUtil { private TokenCounterUtil() {} diff --git a/agentscope-extensions/agentscope-extensions-mem/agentscope-extensions-autocontext-memory/src/test/java/io/agentscope/core/memory/autocontext/AutoContextMemoryTest.java b/agentscope-extensions/agentscope-extensions-mem/agentscope-extensions-autocontext-memory/src/test/java/io/agentscope/core/memory/autocontext/AutoContextMemoryTest.java index 005c0e539..41e15ea18 100644 --- a/agentscope-extensions/agentscope-extensions-mem/agentscope-extensions-autocontext-memory/src/test/java/io/agentscope/core/memory/autocontext/AutoContextMemoryTest.java +++ b/agentscope-extensions/agentscope-extensions-mem/agentscope-extensions-autocontext-memory/src/test/java/io/agentscope/core/memory/autocontext/AutoContextMemoryTest.java @@ -17,19 +17,30 @@ import static org.junit.jupiter.api.Assertions.assertEquals; import static org.junit.jupiter.api.Assertions.assertNotNull; +import static org.junit.jupiter.api.Assertions.assertThrows; import static org.junit.jupiter.api.Assertions.assertTrue; import io.agentscope.core.message.Msg; +import io.agentscope.core.model.ChatResponse; +import io.agentscope.core.model.GenerateOptions; +import io.agentscope.core.model.Model; +import io.agentscope.core.model.ToolSchema; import java.util.List; +import java.util.concurrent.CountDownLatch; +import java.util.concurrent.ExecutorService; +import java.util.concurrent.Executors; +import java.util.concurrent.Future; +import java.util.concurrent.TimeUnit; import java.util.concurrent.atomic.AtomicReference; import org.junit.jupiter.api.Test; +import reactor.core.publisher.Flux; import reactor.core.scheduler.Schedulers; import reactor.test.StepVerifier; class AutoContextMemoryTest { @Test - void deleteMessageRemovesTheIndexedMessageFromBothBuffers() { + void deleteMessageRemovesOnlyTheWorkingBufferEntry() { AutoContextMemory memory = new AutoContextMemory(AutoContextConfig.builder().build(), null); Msg first = AutoContextTestSupport.userMessage("first"); Msg second = AutoContextTestSupport.assistantMessage("second"); @@ -44,9 +55,35 @@ void deleteMessageRemovesTheIndexedMessageFromBothBuffers() { assertEquals(2, memory.getMessages().size()); assertEquals("first", memory.getMessages().get(0).getTextContent()); assertEquals("third", memory.getMessages().get(1).getTextContent()); + assertEquals(3, memory.getOriginalMemoryMsgs().size()); + assertEquals("first", memory.getOriginalMemoryMsgs().get(0).getTextContent()); + assertEquals("second", memory.getOriginalMemoryMsgs().get(1).getTextContent()); + assertEquals("third", memory.getOriginalMemoryMsgs().get(2).getTextContent()); + } + + @Test + void deleteMessageDoesNotDeleteOriginalHistoryAfterCompression() { + AutoContextMemory memory = + new AutoContextMemory( + AutoContextConfig.builder() + .msgThreshold(2) + .lastKeep(0) + .minCompressionTokenThreshold(1) + .build(), + AutoContextTestSupport.recordingModel( + "compressed", new AtomicReference<>())); + memory.addMessage(AutoContextTestSupport.userMessage("first")); + memory.addMessage(AutoContextTestSupport.assistantMessage("second")); + + assertTrue(memory.compressIfNeeded()); + assertEquals(2, memory.getOriginalMemoryMsgs().size()); + + memory.deleteMessage(0); + + assertTrue(memory.getMessages().isEmpty()); assertEquals(2, memory.getOriginalMemoryMsgs().size()); assertEquals("first", memory.getOriginalMemoryMsgs().get(0).getTextContent()); - assertEquals("third", memory.getOriginalMemoryMsgs().get(1).getTextContent()); + assertEquals("second", memory.getOriginalMemoryMsgs().get(1).getTextContent()); } @Test @@ -109,4 +146,90 @@ void compressIfNeededAsyncRunsOnBoundedElasticAndSummarizesMessages() { assertEquals("compressed", memory.getMessages().get(0).getTextContent()); assertEquals(1, memory.getCompressionEvents().size()); } + + @Test + void compressIfNeededDoesNotHoldTheLockWhileTheModelCallIsInFlight() throws Exception { + CountDownLatch streamStarted = new CountDownLatch(1); + CountDownLatch releaseModel = new CountDownLatch(1); + Model blockingModel = + new Model() { + @Override + public Flux stream( + List messages, List tools, GenerateOptions options) { + streamStarted.countDown(); + try { + if (!releaseModel.await(5, TimeUnit.SECONDS)) { + throw new IllegalStateException("timed out waiting for release"); + } + } catch (InterruptedException e) { + Thread.currentThread().interrupt(); + throw new IllegalStateException(e); + } + return Flux.just( + ChatResponse.builder() + .content( + List.of( + io.agentscope.core.message.TextBlock + .builder() + .text("compressed") + .build())) + .build()); + } + + @Override + public String getModelName() { + return "blocking"; + } + }; + AutoContextMemory memory = + new AutoContextMemory( + AutoContextConfig.builder() + .largePayloadThreshold(1) + .msgThreshold(2) + .lastKeep(0) + .minCompressionTokenThreshold(1) + .build(), + blockingModel); + memory.addMessage( + AutoContextTestSupport.userMessage( + "first message is large enough to trigger offload")); + memory.addMessage(AutoContextTestSupport.assistantMessage("second")); + + ExecutorService executor = Executors.newFixedThreadPool(2); + try { + Future compressionFuture = executor.submit(memory::compressIfNeeded); + assertTrue(streamStarted.await(1, TimeUnit.SECONDS)); + + Future addFuture = + executor.submit( + () -> memory.addMessage(AutoContextTestSupport.userMessage("third"))); + addFuture.get(500, TimeUnit.MILLISECONDS); + + releaseModel.countDown(); + assertTrue(compressionFuture.get(5, TimeUnit.SECONDS)); + } finally { + executor.shutdownNow(); + } + + assertEquals(3, memory.getMessages().size()); + assertEquals("compressed", memory.getMessages().get(0).getTextContent()); + assertEquals("second", memory.getMessages().get(1).getTextContent()); + assertEquals("third", memory.getMessages().get(2).getTextContent()); + } + + @Test + void builderRejectsInvalidCompressionSettings() { + assertThrows( + IllegalArgumentException.class, + () -> AutoContextConfig.builder().tokenRatio(0.0).build()); + assertThrows( + IllegalArgumentException.class, + () -> AutoContextConfig.builder().tokenRatio(1.1).build()); + assertThrows( + IllegalArgumentException.class, + () -> AutoContextConfig.builder().maxToken(0).build()); + assertThrows( + IllegalArgumentException.class, + () -> AutoContextConfig.builder().lastKeep(-1).build()); + } } From f8e83943af0135d78cc69e1656812be2dddfef60 Mon Sep 17 00:00:00 2001 From: guslegend <1670547022@qq.com> Date: Wed, 10 Jun 2026 17:42:23 +0800 Subject: [PATCH 4/4] test(autocontext): improve patch coverage for review fixes --- .../autocontext/AutoContextHookTest.java | 47 +++ .../autocontext/AutoContextMemoryTest.java | 215 ++++++++++++++ .../AutoContextSupportClassesTest.java | 184 ++++++++++++ .../autocontext/ContextOffloadToolTest.java | 78 +++++ .../core/memory/autocontext/MsgUtilsTest.java | 271 ++++++++++++++++++ 5 files changed, 795 insertions(+) create mode 100644 agentscope-extensions/agentscope-extensions-mem/agentscope-extensions-autocontext-memory/src/test/java/io/agentscope/core/memory/autocontext/AutoContextSupportClassesTest.java create mode 100644 agentscope-extensions/agentscope-extensions-mem/agentscope-extensions-autocontext-memory/src/test/java/io/agentscope/core/memory/autocontext/MsgUtilsTest.java diff --git a/agentscope-extensions/agentscope-extensions-mem/agentscope-extensions-autocontext-memory/src/test/java/io/agentscope/core/memory/autocontext/AutoContextHookTest.java b/agentscope-extensions/agentscope-extensions-mem/agentscope-extensions-autocontext-memory/src/test/java/io/agentscope/core/memory/autocontext/AutoContextHookTest.java index 1cc560955..70411fc56 100644 --- a/agentscope-extensions/agentscope-extensions-mem/agentscope-extensions-autocontext-memory/src/test/java/io/agentscope/core/memory/autocontext/AutoContextHookTest.java +++ b/agentscope-extensions/agentscope-extensions-mem/agentscope-extensions-autocontext-memory/src/test/java/io/agentscope/core/memory/autocontext/AutoContextHookTest.java @@ -19,12 +19,15 @@ import static org.junit.jupiter.api.Assertions.assertNotNull; import static org.junit.jupiter.api.Assertions.assertSame; import static org.junit.jupiter.api.Assertions.assertTrue; +import static org.mockito.Mockito.mock; import io.agentscope.core.ReActAgent; +import io.agentscope.core.agent.Agent; import io.agentscope.core.agent.RuntimeContext; import io.agentscope.core.hook.PreCallEvent; import io.agentscope.core.hook.PreReasoningEvent; import io.agentscope.core.state.AgentState; +import io.agentscope.core.tool.Toolkit; import java.lang.reflect.Field; import java.util.List; import java.util.concurrent.atomic.AtomicReference; @@ -105,6 +108,50 @@ void handlePreReasoningCompressesWithoutBlockingAndRewritesState() { assertTrue(threadName.get().contains("boundedElastic")); } + @Test + void handlePreCallIgnoresNonReactAgentAndRegistersOnlyOnce() { + AutoContextHook hook = new AutoContextHook(); + Agent nonReactAgent = mock(Agent.class); + PreCallEvent plainEvent = + new PreCallEvent( + nonReactAgent, List.of(AutoContextTestSupport.userMessage("hello"))); + StepVerifier.create(hook.handlePreCall(plainEvent)).expectNext(plainEvent).verifyComplete(); + + Toolkit toolkit = new Toolkit(); + ReActAgent agent = + ReActAgent.builder() + .name("register-once") + .sysPrompt("system") + .toolkit(toolkit) + .model(AutoContextTestSupport.noopModel()) + .build(); + PreCallEvent event = + new PreCallEvent(agent, List.of(AutoContextTestSupport.userMessage("hello"))); + + StepVerifier.create(hook.handlePreCall(event)).expectNext(event).verifyComplete(); + int firstCount = agent.getToolkit().getToolNames().size(); + StepVerifier.create(hook.handlePreCall(event)).expectNext(event).verifyComplete(); + assertEquals(firstCount, agent.getToolkit().getToolNames().size()); + assertTrue(agent.getToolkit().getToolNames().contains("context_reload")); + } + + @Test + void handlePreReasoningIgnoresNonReactAgent() { + AutoContextHook hook = new AutoContextHook(); + Agent agent = mock(Agent.class); + + PreReasoningEvent event = + new PreReasoningEvent( + agent, "noop", null, List.of(AutoContextTestSupport.userMessage("x"))); + + StepVerifier.create(hook.handlePreReasoning(event)) + .assertNext(returned -> assertSame(event, returned)) + .verifyComplete(); + assertEquals(1, event.getInputMessages().size()); + assertEquals("x", event.getInputMessages().get(0).getTextContent()); + assertTrue(event.getSystemMessage() == null); + } + private static void setActiveRuntimeContext(ReActAgent agent, RuntimeContext runtimeContext) { try { Field field = ReActAgent.class.getDeclaredField("activeRc"); diff --git a/agentscope-extensions/agentscope-extensions-mem/agentscope-extensions-autocontext-memory/src/test/java/io/agentscope/core/memory/autocontext/AutoContextMemoryTest.java b/agentscope-extensions/agentscope-extensions-mem/agentscope-extensions-autocontext-memory/src/test/java/io/agentscope/core/memory/autocontext/AutoContextMemoryTest.java index 41e15ea18..dee9c27c7 100644 --- a/agentscope-extensions/agentscope-extensions-mem/agentscope-extensions-autocontext-memory/src/test/java/io/agentscope/core/memory/autocontext/AutoContextMemoryTest.java +++ b/agentscope-extensions/agentscope-extensions-mem/agentscope-extensions-autocontext-memory/src/test/java/io/agentscope/core/memory/autocontext/AutoContextMemoryTest.java @@ -17,15 +17,21 @@ import static org.junit.jupiter.api.Assertions.assertEquals; import static org.junit.jupiter.api.Assertions.assertNotNull; +import static org.junit.jupiter.api.Assertions.assertNull; import static org.junit.jupiter.api.Assertions.assertThrows; import static org.junit.jupiter.api.Assertions.assertTrue; import io.agentscope.core.message.Msg; +import io.agentscope.core.message.MsgRole; +import io.agentscope.core.message.TextBlock; +import io.agentscope.core.message.ToolResultBlock; +import io.agentscope.core.message.ToolUseBlock; import io.agentscope.core.model.ChatResponse; import io.agentscope.core.model.GenerateOptions; import io.agentscope.core.model.Model; import io.agentscope.core.model.ToolSchema; import java.util.List; +import java.util.Map; import java.util.concurrent.CountDownLatch; import java.util.concurrent.ExecutorService; import java.util.concurrent.Executors; @@ -231,5 +237,214 @@ void builderRejectsInvalidCompressionSettings() { assertThrows( IllegalArgumentException.class, () -> AutoContextConfig.builder().lastKeep(-1).build()); + assertThrows( + IllegalArgumentException.class, + () -> AutoContextConfig.builder().largePayloadThreshold(0).build()); + assertThrows( + IllegalArgumentException.class, + () -> AutoContextConfig.builder().offloadSinglePreview(-1).build()); + assertThrows( + IllegalArgumentException.class, + () -> AutoContextConfig.builder().msgThreshold(0).build()); + assertThrows( + IllegalArgumentException.class, + () -> AutoContextConfig.builder().minConsecutiveToolMessages(0).build()); + assertThrows( + IllegalArgumentException.class, + () -> AutoContextConfig.builder().currentRoundCompressionRatio(1.2).build()); + assertThrows( + IllegalArgumentException.class, + () -> AutoContextConfig.builder().minCompressionTokenThreshold(0).build()); + } + + @Test + void mergeWithContextAppendsTailOrReplacesWorkingBufferOnly() { + AutoContextMemory memory = new AutoContextMemory(AutoContextConfig.builder().build(), null); + Msg first = AutoContextTestSupport.userMessage("first"); + Msg second = AutoContextTestSupport.assistantMessage("second"); + Msg third = AutoContextTestSupport.userMessage("third"); + memory.mergeWithContext(List.of(first, second)); + memory.mergeWithContext(List.of(first, second, third)); + + assertEquals(3, memory.getMessages().size()); + assertEquals(3, memory.getOriginalMemoryMsgs().size()); + + Msg reset = AutoContextTestSupport.assistantMessage("reset"); + memory.mergeWithContext(List.of(reset)); + + assertEquals(1, memory.getMessages().size()); + assertEquals("reset", memory.getMessages().get(0).getTextContent()); + assertEquals(3, memory.getOriginalMemoryMsgs().size()); + assertEquals("first", memory.getOriginalMemoryMsgs().get(0).getTextContent()); + } + + @Test + void compressIfNeededReturnsFalseWhenThresholdsOrCandidatesDoNotMatch() { + AutoContextMemory empty = new AutoContextMemory(AutoContextConfig.builder().build(), null); + assertTrue(!empty.compressIfNeeded()); + + AutoContextMemory belowThreshold = + new AutoContextMemory( + AutoContextConfig.builder() + .msgThreshold(10) + .minCompressionTokenThreshold(1) + .build(), + null); + belowThreshold.addMessage(AutoContextTestSupport.userMessage("short")); + belowThreshold.addMessage(AutoContextTestSupport.assistantMessage("tiny")); + assertTrue(!belowThreshold.compressIfNeeded()); + } + + @Test + void compressIfNeededHandlesCurrentRoundToolUseAndLargeToolOffload() { + AutoContextMemory currentRound = + new AutoContextMemory( + AutoContextConfig.builder() + .msgThreshold(3) + .lastKeep(2) + .minCompressionTokenThreshold(1) + .build(), + AutoContextTestSupport.recordingModel( + "compressed round", new AtomicReference<>())); + currentRound.addMessage(AutoContextTestSupport.userMessage("user")); + currentRound.addMessage(AutoContextTestSupport.assistantMessage("assistant body")); + currentRound.addMessage( + Msg.builder() + .role(MsgRole.ASSISTANT) + .name("assistant") + .content( + ToolUseBlock.builder() + .id("call-1") + .name("tool") + .input(Map.of("k", "v")) + .build()) + .build()); + + assertTrue(currentRound.compressIfNeeded()); + assertEquals(3, currentRound.getMessages().size()); + assertEquals("user", currentRound.getMessages().get(0).getTextContent()); + assertTrue(currentRound.getMessages().get(1).getTextContent().contains("compressed round")); + assertTrue(currentRound.getMessages().get(1).getTextContent().contains("CONTEXT_OFFLOAD")); + assertTrue(currentRound.getMessages().get(2).hasContentBlocks(ToolUseBlock.class)); + + AutoContextMemory largeTool = + new AutoContextMemory( + AutoContextConfig.builder() + .largePayloadThreshold(1) + .msgThreshold(99) + .maxToken(1) + .tokenRatio(1.0) + .lastKeep(0) + .minCompressionTokenThreshold(1) + .build(), + AutoContextTestSupport.recordingModel( + "tool summary", new AtomicReference<>())); + Msg toolMessage = + Msg.builder() + .role(MsgRole.TOOL) + .name("tool") + .content( + ToolResultBlock.builder() + .id("call-2") + .name("search") + .output( + List.of( + TextBlock.builder() + .text("very long payload") + .build())) + .build()) + .build(); + largeTool.addMessage(toolMessage); + + assertTrue(largeTool.compressIfNeeded()); + assertEquals(MsgRole.TOOL, largeTool.getMessages().get(0).getRole()); + assertTrue(largeTool.getMessages().get(0).getMetadata().containsKey("_compress_meta")); + } + + @Test + void compressIfNeededFallsBackWhenModelUnavailableOrReturnsNoText() { + AutoContextMemory noModel = + new AutoContextMemory( + AutoContextConfig.builder() + .msgThreshold(2) + .lastKeep(0) + .minCompressionTokenThreshold(1) + .build(), + null); + noModel.addMessage(AutoContextTestSupport.userMessage("alpha")); + noModel.addMessage(AutoContextTestSupport.assistantMessage("beta")); + + assertTrue(noModel.compressIfNeeded()); + assertTrue(noModel.getMessages().get(0).getTextContent().contains("Compressed 2 messages")); + + Model emptyResponseModel = + new Model() { + @Override + public Flux stream( + List messages, List tools, GenerateOptions options) { + return Flux.just(ChatResponse.builder().content(List.of()).build()); + } + + @Override + public String getModelName() { + return "empty"; + } + }; + AutoContextMemory blankResponse = + new AutoContextMemory( + AutoContextConfig.builder() + .msgThreshold(2) + .lastKeep(0) + .minCompressionTokenThreshold(1) + .build(), + emptyResponseModel); + blankResponse.addMessage(AutoContextTestSupport.userMessage("gamma")); + blankResponse.addMessage(AutoContextTestSupport.assistantMessage("delta")); + + assertTrue(blankResponse.compressIfNeeded()); + assertTrue( + blankResponse + .getMessages() + .get(0) + .getTextContent() + .contains("Compressed 2 messages")); + } + + @Test + void clearAndSetWorkingMessagesResetVisibleBuffersOnly() { + AutoContextMemory memory = new AutoContextMemory(AutoContextConfig.builder().build(), null); + memory.addMessage(AutoContextTestSupport.userMessage("a")); + memory.addMessage(AutoContextTestSupport.assistantMessage("b")); + memory.offload("uuid", List.of(AutoContextTestSupport.userMessage("offload"))); + memory.clear("uuid"); + assertTrue(memory.reload("uuid").isEmpty()); + + memory.setWorkingMessages(List.of(AutoContextTestSupport.assistantMessage("working"))); + assertEquals(1, memory.getMessages().size()); + assertEquals("working", memory.getMessages().get(0).getTextContent()); + assertEquals(2, memory.getOriginalMemoryMsgs().size()); + + memory.clear(); + assertTrue(memory.getMessages().isEmpty()); + assertTrue(memory.getOriginalMemoryMsgs().isEmpty()); + assertTrue(memory.getCompressionEvents().isEmpty()); + assertTrue(memory.getOffloadContext().isEmpty()); + } + + @Test + void restoreHandlesNullSnapshotAndMissingCompressionEvents() { + AutoContextMemory memory = new AutoContextMemory(AutoContextConfig.builder().build(), null); + memory.restore(null); + assertTrue(memory.getMessages().isEmpty()); + + AutoContextState snapshot = new AutoContextState(); + snapshot.setWorkingMessages(List.of(AutoContextTestSupport.userMessage("saved"))); + snapshot.setOriginalMessages(List.of(AutoContextTestSupport.userMessage("saved"))); + snapshot.setCompressionEvents(null); + memory.restore(snapshot); + + assertEquals(1, memory.getMessages().size()); + assertTrue(memory.getCompressionEvents().isEmpty()); + assertNull(memory.getAutoContextConfig().getCustomPrompt()); } } diff --git a/agentscope-extensions/agentscope-extensions-mem/agentscope-extensions-autocontext-memory/src/test/java/io/agentscope/core/memory/autocontext/AutoContextSupportClassesTest.java b/agentscope-extensions/agentscope-extensions-mem/agentscope-extensions-autocontext-memory/src/test/java/io/agentscope/core/memory/autocontext/AutoContextSupportClassesTest.java new file mode 100644 index 000000000..22e1c93c3 --- /dev/null +++ b/agentscope-extensions/agentscope-extensions-mem/agentscope-extensions-autocontext-memory/src/test/java/io/agentscope/core/memory/autocontext/AutoContextSupportClassesTest.java @@ -0,0 +1,184 @@ +/* + * Copyright 2024-2026 the original author or authors. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package io.agentscope.core.memory.autocontext; + +import static org.junit.jupiter.api.Assertions.assertEquals; +import static org.junit.jupiter.api.Assertions.assertSame; +import static org.junit.jupiter.api.Assertions.assertTrue; + +import io.agentscope.core.ReActAgent; +import io.agentscope.core.hook.PostCallEvent; +import io.agentscope.core.message.Msg; +import io.agentscope.core.model.ChatUsage; +import java.util.ArrayList; +import java.util.HashMap; +import java.util.List; +import java.util.Map; +import org.junit.jupiter.api.Test; +import reactor.test.StepVerifier; + +class AutoContextSupportClassesTest { + + @Test + void promptProviderUsesCustomPromptsAndFallsBackForBlankValues() { + PromptConfig customPrompt = + PromptConfig.builder() + .previousRoundToolCompressPrompt("tool") + .previousRoundSummaryPrompt("summary") + .currentRoundLargeMessagePrompt("large") + .currentRoundCompressPrompt("current") + .build(); + + assertEquals("tool", PromptProvider.getPreviousRoundToolCompressPrompt(customPrompt)); + assertEquals("summary", PromptProvider.getPreviousRoundSummaryPrompt(customPrompt)); + assertEquals("large", PromptProvider.getCurrentRoundLargeMessagePrompt(customPrompt)); + assertEquals("current", PromptProvider.getCurrentRoundCompressPrompt(customPrompt)); + + PromptConfig blankPrompt = + PromptConfig.builder() + .previousRoundToolCompressPrompt(" ") + .previousRoundSummaryPrompt("") + .currentRoundLargeMessagePrompt(" ") + .currentRoundCompressPrompt("") + .build(); + + assertEquals( + Prompts.PREVIOUS_ROUND_TOOL_INVOCATION_COMPRESS_PROMPT, + PromptProvider.getPreviousRoundToolCompressPrompt(blankPrompt)); + assertEquals( + Prompts.PREVIOUS_ROUND_CONVERSATION_SUMMARY_PROMPT, + PromptProvider.getPreviousRoundSummaryPrompt(blankPrompt)); + assertEquals( + Prompts.CURRENT_ROUND_LARGE_MESSAGE_SUMMARY_PROMPT, + PromptProvider.getCurrentRoundLargeMessagePrompt(blankPrompt)); + assertEquals( + Prompts.CURRENT_ROUND_MESSAGE_COMPRESS_PROMPT, + PromptProvider.getCurrentRoundCompressPrompt(blankPrompt)); + } + + @Test + void compressionEventAccessorsReflectMetadataAndDefensiveCopies() { + Map metadata = new HashMap<>(); + metadata.put("tokenBefore", 12); + metadata.put("tokenAfter", 5); + metadata.put("inputToken", 9); + metadata.put("outputToken", 3); + metadata.put("time", 1.5d); + CompressionEvent event = + new CompressionEvent("type", 10L, 2, "prev", "next", "compressed", metadata); + + metadata.put("tokenBefore", 100); + + assertEquals("type", event.getEventType()); + assertEquals(10L, event.getTimestamp()); + assertEquals(2, event.getCompressedMessageCount()); + assertEquals("prev", event.getPreviousMessageId()); + assertEquals("next", event.getNextMessageId()); + assertEquals("compressed", event.getCompressedMessageId()); + assertEquals(12, event.getTokenBefore()); + assertEquals(5, event.getTokenAfter()); + assertEquals(7, event.getTokenReduction()); + assertEquals(9, event.getCompressInputToken()); + assertEquals(3, event.getCompressOutputToken()); + + CompressionEvent mutable = new CompressionEvent(); + mutable.setEventType("updated"); + mutable.setTimestamp(20L); + mutable.setCompressedMessageCount(4); + mutable.setPreviousMessageId("p"); + mutable.setNextMessageId("n"); + mutable.setCompressedMessageId("c"); + mutable.setMetadata(null); + + assertEquals("updated", mutable.getEventType()); + assertEquals(20L, mutable.getTimestamp()); + assertEquals(4, mutable.getCompressedMessageCount()); + assertEquals("p", mutable.getPreviousMessageId()); + assertEquals("n", mutable.getNextMessageId()); + assertEquals("c", mutable.getCompressedMessageId()); + assertTrue(mutable.getMetadata().isEmpty()); + } + + @Test + void autoContextStateCopiesCollectionsAndNormalizesNulls() { + List working = new ArrayList<>(List.of(AutoContextTestSupport.userMessage("work"))); + List original = + new ArrayList<>(List.of(AutoContextTestSupport.assistantMessage("original"))); + List events = + new ArrayList<>( + List.of(new CompressionEvent("type", 1L, 1, null, null, null, null))); + Map> offload = new HashMap<>(); + offload.put( + "uuid", new ArrayList<>(List.of(AutoContextTestSupport.userMessage("offload")))); + + AutoContextState state = new AutoContextState(); + state.setWorkingMessages(working); + state.setOriginalMessages(original); + state.setCompressionEvents(events); + state.setOffloadContext(offload); + + working.clear(); + original.clear(); + events.clear(); + offload.clear(); + + assertEquals(1, state.getWorkingMessages().size()); + assertEquals(1, state.getOriginalMessages().size()); + assertEquals(1, state.getCompressionEvents().size()); + assertEquals(1, state.getOffloadContext().size()); + + state.setWorkingMessages(null); + state.setOriginalMessages(null); + state.setCompressionEvents(null); + state.setOffloadContext(null); + + assertTrue(state.getWorkingMessages().isEmpty()); + assertTrue(state.getOriginalMessages().isEmpty()); + assertTrue(state.getCompressionEvents().isEmpty()); + assertTrue(state.getOffloadContext().isEmpty()); + } + + @Test + void tokenCounterUsesMinimumOfOneAndCharHeuristic() { + assertEquals(1, TokenCounterUtil.calculateToken(List.of())); + assertEquals( + 2, + TokenCounterUtil.calculateToken( + List.of(AutoContextTestSupport.userMessage("12345678")))); + } + + @Test + void onEventReturnsUnhandledEventUnchanged() { + AutoContextHook hook = new AutoContextHook(); + ReActAgent agent = + ReActAgent.builder() + .name("post-call") + .sysPrompt("system") + .model(AutoContextTestSupport.noopModel()) + .build(); + Msg finalMessage = + Msg.builder() + .role(io.agentscope.core.message.MsgRole.ASSISTANT) + .name("assistant") + .textContent("done") + .usage(ChatUsage.builder().inputTokens(1).outputTokens(2).time(0.1).build()) + .build(); + PostCallEvent event = new PostCallEvent(agent, finalMessage); + + StepVerifier.create(hook.onEvent(event)).expectNext(event).verifyComplete(); + assertSame(finalMessage, event.getFinalMessage()); + } +} diff --git a/agentscope-extensions/agentscope-extensions-mem/agentscope-extensions-autocontext-memory/src/test/java/io/agentscope/core/memory/autocontext/ContextOffloadToolTest.java b/agentscope-extensions/agentscope-extensions-mem/agentscope-extensions-autocontext-memory/src/test/java/io/agentscope/core/memory/autocontext/ContextOffloadToolTest.java index 97fd97158..abacdeffa 100644 --- a/agentscope-extensions/agentscope-extensions-mem/agentscope-extensions-autocontext-memory/src/test/java/io/agentscope/core/memory/autocontext/ContextOffloadToolTest.java +++ b/agentscope-extensions/agentscope-extensions-mem/agentscope-extensions-autocontext-memory/src/test/java/io/agentscope/core/memory/autocontext/ContextOffloadToolTest.java @@ -16,6 +16,7 @@ package io.agentscope.core.memory.autocontext; import static org.junit.jupiter.api.Assertions.assertEquals; +import static org.junit.jupiter.api.Assertions.assertSame; import static org.junit.jupiter.api.Assertions.assertTrue; import io.agentscope.core.ReActAgent; @@ -85,4 +86,81 @@ void reloadUsesHookRuntimeContextWhenAgentAndStateAreAvailable() { assertTrue(messages.get(0).getTextContent().contains("offloaded-one")); assertTrue(messages.get(1).getTextContent().contains("offloaded-two")); } + + @Test + void reloadReturnsErrorMessagesForInvalidOrUnavailableContext() { + ContextOffloadTool noLoaderTool = new ContextOffloadTool((ContextOffLoader) null); + List blankUuidMessages = noLoaderTool.reload(" "); + assertEquals(1, blankUuidMessages.size()); + assertTrue(blankUuidMessages.get(0).getTextContent().contains("UUID cannot be null")); + + List unavailableMessages = noLoaderTool.reload("uuid-1"); + assertEquals(1, unavailableMessages.size()); + assertTrue( + unavailableMessages + .get(0) + .getTextContent() + .contains("Context offloader is not available")); + } + + @Test + void reloadReturnsErrorMessageWhenLoaderThrowsOrNothingFound() { + ContextOffloadTool throwingTool = + new ContextOffloadTool( + new ContextOffLoader() { + @Override + public void offload(String uuid, List messages) {} + + @Override + public List reload(String uuid) { + throw new IllegalStateException("boom"); + } + + @Override + public void clear(String uuid) {} + }); + + List thrown = throwingTool.reload("uuid-2"); + assertEquals(1, thrown.size()); + assertTrue(thrown.get(0).getTextContent().contains("boom")); + + ContextOffloadTool emptyTool = + new ContextOffloadTool( + new ContextOffLoader() { + @Override + public void offload(String uuid, List messages) {} + + @Override + public List reload(String uuid) { + return List.of(); + } + + @Override + public void clear(String uuid) {} + }); + List missing = emptyTool.reload("uuid-3"); + assertEquals(1, missing.size()); + assertTrue(missing.get(0).getTextContent().contains("No messages found for UUID")); + } + + @Test + void reloadWithoutAgentUsesHookBackedMemoryLookup() { + AutoContextHook hook = new AutoContextHook(); + ContextOffloadTool tool = new ContextOffloadTool(hook); + ReActAgent agent = + ReActAgent.builder() + .name("hook-default") + .sysPrompt("system") + .model(AutoContextTestSupport.noopModel()) + .build(); + RuntimeContext runtimeContext = + RuntimeContext.builder().sessionId("session-2").userId("bob").build(); + List expected = List.of(AutoContextTestSupport.userMessage("hook-only")); + + hook.memoryFor(agent, runtimeContext).offload("hook-uuid", expected); + + List reloaded = tool.reload("hook-uuid", agent, runtimeContext); + assertEquals(1, reloaded.size()); + assertSame(expected.get(0), reloaded.get(0)); + } } diff --git a/agentscope-extensions/agentscope-extensions-mem/agentscope-extensions-autocontext-memory/src/test/java/io/agentscope/core/memory/autocontext/MsgUtilsTest.java b/agentscope-extensions/agentscope-extensions-mem/agentscope-extensions-autocontext-memory/src/test/java/io/agentscope/core/memory/autocontext/MsgUtilsTest.java new file mode 100644 index 000000000..b306fad97 --- /dev/null +++ b/agentscope-extensions/agentscope-extensions-mem/agentscope-extensions-autocontext-memory/src/test/java/io/agentscope/core/memory/autocontext/MsgUtilsTest.java @@ -0,0 +1,271 @@ +/* + * Copyright 2024-2026 the original author or authors. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package io.agentscope.core.memory.autocontext; + +import static org.junit.jupiter.api.Assertions.assertEquals; +import static org.junit.jupiter.api.Assertions.assertFalse; +import static org.junit.jupiter.api.Assertions.assertInstanceOf; +import static org.junit.jupiter.api.Assertions.assertSame; +import static org.junit.jupiter.api.Assertions.assertTrue; + +import io.agentscope.core.message.ContentBlock; +import io.agentscope.core.message.MessageMetadataKeys; +import io.agentscope.core.message.Msg; +import io.agentscope.core.message.MsgRole; +import io.agentscope.core.message.TextBlock; +import io.agentscope.core.message.ToolResultBlock; +import io.agentscope.core.message.ToolUseBlock; +import io.agentscope.core.util.JsonUtils; +import java.util.ArrayList; +import java.util.HashMap; +import java.util.List; +import java.util.Map; +import org.junit.jupiter.api.Test; + +class MsgUtilsTest { + + @Test + void serializeAndDeserializeMessageListsAndMapsRoundTrip() { + List messages = + List.of( + AutoContextTestSupport.userMessage("hello"), + AutoContextTestSupport.assistantMessage("world")); + Map> byUuid = Map.of("uuid", messages); + + Object serializedList = MsgUtils.serializeMsgList(messages); + Object serializedMap = MsgUtils.serializeMsgListMap(byUuid); + + assertInstanceOf(List.class, serializedList); + assertInstanceOf(Map.class, serializedMap); + + Object deserializedList = MsgUtils.deserializeToMsgList(serializedList); + Object deserializedMap = MsgUtils.deserializeToMsgListMap(serializedMap); + + assertInstanceOf(List.class, deserializedList); + assertInstanceOf(Map.class, deserializedMap); + @SuppressWarnings("unchecked") + List restoredMessages = (List) deserializedList; + assertEquals(2, restoredMessages.size()); + assertEquals("hello", restoredMessages.get(0).getTextContent()); + assertEquals("world", restoredMessages.get(1).getTextContent()); + assertEquals(MsgRole.USER, restoredMessages.get(0).getRole()); + assertEquals(MsgRole.ASSISTANT, restoredMessages.get(1).getRole()); + @SuppressWarnings("unchecked") + Map> restoredMap = (Map>) deserializedMap; + assertEquals(1, restoredMap.size()); + assertEquals("hello", restoredMap.get("uuid").get(0).getTextContent()); + assertEquals("world", restoredMap.get("uuid").get(1).getTextContent()); + + assertSame("plain", MsgUtils.serializeMsgList("plain")); + assertSame("plain", MsgUtils.deserializeToMsgList("plain")); + assertSame("plain", MsgUtils.serializeMsgListMap("plain")); + assertSame("plain", MsgUtils.deserializeToMsgListMap("plain")); + } + + @Test + void serializeAndDeserializeCompressionEventsSupportMetadataAndLegacyFields() { + CompressionEvent event = + new CompressionEvent( + "type", + 1L, + 2, + "prev", + "next", + "compressed", + Map.of("tokenBefore", 6, "tokenAfter", 2, "inputToken", 3, "time", 0.5d)); + Object serialized = MsgUtils.serializeCompressionEventList(List.of(event)); + Object deserialized = MsgUtils.deserializeToCompressionEventList(serialized); + + assertInstanceOf(List.class, deserialized); + @SuppressWarnings("unchecked") + List events = (List) deserialized; + assertEquals(1, events.size()); + assertEquals("type", events.get(0).getEventType()); + assertEquals(6, events.get(0).getTokenBefore()); + assertEquals(2, events.get(0).getTokenAfter()); + + Map legacyEvent = new HashMap<>(); + legacyEvent.put("eventType", "legacy"); + legacyEvent.put("timestamp", 2L); + legacyEvent.put("compressedMessageCount", 1); + legacyEvent.put("previousMessageId", "p"); + legacyEvent.put("nextMessageId", "n"); + legacyEvent.put("compressedMessageId", "c"); + legacyEvent.put("tokenBefore", 10); + legacyEvent.put("tokenAfter", 4); + legacyEvent.put("inputToken", 7); + legacyEvent.put("outputToken", 3); + legacyEvent.put("time", 1.25d); + List> legacy = List.of(legacyEvent); + @SuppressWarnings("unchecked") + List legacyEvents = + (List) MsgUtils.deserializeToCompressionEventList(legacy); + assertEquals(1, legacyEvents.size()); + assertEquals("legacy", legacyEvents.get(0).getEventType()); + assertEquals(10, legacyEvents.get(0).getTokenBefore()); + assertEquals(4, legacyEvents.get(0).getTokenAfter()); + assertEquals(7, legacyEvents.get(0).getCompressInputToken()); + assertEquals(3, legacyEvents.get(0).getCompressOutputToken()); + } + + @Test + void replaceMsgAndToolPredicatesHandleContentTypesAndCompressedMarkers() { + Msg plainAssistant = AutoContextTestSupport.assistantMessage("plain"); + Msg toolUse = toolUseMessage("create_plan", "plan-1"); + Msg toolResult = toolResultMessage("plan-1", "tool output"); + List messages = new ArrayList<>(List.of(plainAssistant, toolUse, toolResult)); + + MsgUtils.replaceMsg(messages, 1, 9, AutoContextTestSupport.assistantMessage("summary")); + assertEquals(2, messages.size()); + assertEquals("summary", messages.get(1).getTextContent()); + + MsgUtils.replaceMsg(messages, -1, 0, AutoContextTestSupport.assistantMessage("ignored")); + MsgUtils.replaceMsg(messages, 5, 6, AutoContextTestSupport.assistantMessage("ignored")); + assertEquals(2, messages.size()); + + Map compressedMeta = new HashMap<>(); + compressedMeta.put("_compress_meta", Map.of("compressed_current_round", true)); + Msg compressedCurrentRound = + Msg.builder() + .role(MsgRole.ASSISTANT) + .name("assistant") + .textContent("compressed") + .metadata(compressedMeta) + .build(); + + assertTrue(MsgUtils.isToolMessage(toolUse)); + assertTrue(MsgUtils.isToolMessage(toolResult)); + assertTrue( + MsgUtils.isToolMessage( + Msg.builder().role(MsgRole.TOOL).name("tool").textContent("ok").build())); + assertTrue(MsgUtils.isToolUseMessage(toolUse)); + assertFalse(MsgUtils.isToolUseMessage(toolResult)); + assertTrue(MsgUtils.isToolResultMessage(toolResult)); + assertFalse(MsgUtils.isToolResultMessage(plainAssistant)); + assertTrue(MsgUtils.isCompressedMessage(compressedCurrentRound)); + assertTrue(MsgUtils.isFinalAssistantResponse(plainAssistant)); + assertFalse(MsgUtils.isFinalAssistantResponse(toolUse)); + assertFalse(MsgUtils.isFinalAssistantResponse(compressedCurrentRound)); + } + + @Test + void planRelatedToolHelpersFilterMatchingCallsAndResults() { + Msg planToolUse = toolUseMessage("create_plan", "plan-1"); + Msg otherToolUse = toolUseMessage("search_web", "search-1"); + Msg planToolResult = toolResultMessage("plan-1", "plan result"); + Msg otherToolResult = toolResultMessage("search-1", "search result"); + List filtered = + MsgUtils.filterPlanRelatedToolCalls( + List.of( + AutoContextTestSupport.userMessage("start"), + planToolUse, + otherToolUse, + planToolResult, + otherToolResult)); + + assertTrue(MsgUtils.isPlanRelatedTool("create_plan")); + assertFalse(MsgUtils.isPlanRelatedTool("search_web")); + assertTrue(MsgUtils.containsPlanRelatedToolCall(planToolUse)); + assertFalse(MsgUtils.containsPlanRelatedToolCall(otherToolUse)); + assertEquals(3, filtered.size()); + assertEquals("start", filtered.get(0).getTextContent()); + assertEquals(otherToolUse, filtered.get(1)); + assertEquals(otherToolResult, filtered.get(2)); + } + + @Test + void calculateMessageCharCountCountsTextToolUseAndToolResultContent() { + ToolUseBlock toolUseBlock = + ToolUseBlock.builder() + .id("call-1") + .name("lookup") + .input(Map.of("city", "beijing")) + .content("raw") + .build(); + ToolResultBlock toolResultBlock = + ToolResultBlock.builder() + .id("call-1") + .name("lookup") + .output(List.of(TextBlock.builder().text("sunny").build())) + .build(); + Msg textMessage = AutoContextTestSupport.userMessage("hello"); + Msg toolUseMessage = + Msg.builder() + .role(MsgRole.ASSISTANT) + .name("assistant") + .content(toolUseBlock) + .build(); + Msg toolResultMessage = + Msg.builder().role(MsgRole.TOOL).name("tool").content(toolResultBlock).build(); + + int expectedToolUseChars = + "lookup".length() + + "call-1".length() + + JsonUtils.getJsonCodec().toJson(Map.of("city", "beijing")).length() + + "raw".length(); + int expectedToolResultChars = "lookup".length() + "call-1".length() + "sunny".length(); + + assertEquals(5, MsgUtils.calculateMessageCharCount(textMessage)); + assertEquals(expectedToolUseChars, MsgUtils.calculateMessageCharCount(toolUseMessage)); + assertEquals( + expectedToolResultChars, MsgUtils.calculateMessageCharCount(toolResultMessage)); + assertEquals( + 5 + expectedToolUseChars + expectedToolResultChars, + MsgUtils.calculateMessagesCharCount( + List.of(textMessage, toolUseMessage, toolResultMessage))); + } + + @Test + void finalAssistantResponseSupportsStructuredOutputMessages() { + Msg structuredAssistant = + Msg.builder() + .role(MsgRole.ASSISTANT) + .name("assistant") + .content(TextBlock.builder().text("done").build()) + .metadata(Map.of(MessageMetadataKeys.STRUCTURED_OUTPUT, Map.of("ok", true))) + .build(); + + assertTrue(MsgUtils.isFinalAssistantResponse(structuredAssistant)); + } + + private static Msg toolUseMessage(String toolName, String toolCallId) { + return Msg.builder() + .role(MsgRole.ASSISTANT) + .name("assistant") + .content( + ToolUseBlock.builder() + .id(toolCallId) + .name(toolName) + .input(Map.of("q", "value")) + .build()) + .build(); + } + + private static Msg toolResultMessage(String toolCallId, String text) { + return Msg.builder() + .role(MsgRole.TOOL) + .name("tool") + .content( + ToolResultBlock.builder() + .id(toolCallId) + .name("tool") + .output( + List.of( + TextBlock.builder().text(text).build())) + .build()) + .build(); + } +}