diff --git a/agentscope-core/src/main/java/io/agentscope/core/ReActAgent.java b/agentscope-core/src/main/java/io/agentscope/core/ReActAgent.java index ba970eafe..4dd8fcfba 100644 --- a/agentscope-core/src/main/java/io/agentscope/core/ReActAgent.java +++ b/agentscope-core/src/main/java/io/agentscope/core/ReActAgent.java @@ -623,6 +623,16 @@ protected void afterAgentExecution() { unbindRuntimeContextFromHooks(); } + @Override + protected void afterCallScopeExecution(Object callScope) { + if (stateStore == null) { + return; + } + CallExecution scope = (CallExecution) callScope; + stateCache.remove(scope.slotKey, scope.state); + permissionEngineCache.remove(scope.slotKey, scope.permissionEngine); + } + private RuntimeContext buildMergedRuntimeContext(RuntimeContext run) { if (run == null) { if (toolExecutionContext != null) { @@ -3758,6 +3768,38 @@ public void saveAgentState(String userId, String sessionId) { } } + /** + * Evicts the in-memory state and permission engine for one session. + * + *

This does not delete persisted state. A later access reloads the session from the + * configured {@link AgentStateStore}, or creates fresh in-memory state when no store is + * configured. Call this only after any in-flight call for the session has terminated. + * + * @param userId user identity for the slot (may be {@code null}) + * @param sessionId session identity (falls back to the default session id when blank) + */ + public void evictSession(String userId, String sessionId) { + String sid = (sessionId == null || sessionId.isBlank()) ? defaultSessionId : sessionId; + String slot = slotKey(userId, sid); + stateCache.remove(slot); + permissionEngineCache.remove(slot); + } + + /** + * Evicts all in-memory state and permission engines belonging to one user. + * + *

This does not delete persisted state. Call this only after any in-flight calls for the + * user have terminated. + * + * @param userId user identity ({@code null} or blank selects anonymous sessions) + */ + public void evictUser(String userId) { + String normalizedUser = userId == null || userId.isBlank() ? "__anon__" : userId; + String prefix = normalizedUser + "/"; + stateCache.keySet().removeIf(slot -> slot.startsWith(prefix)); + permissionEngineCache.keySet().removeIf(slot -> slot.startsWith(prefix)); + } + /** Returns the {@link AgentStateStore} configured for state persistence, or {@code null}. */ public AgentStateStore getStateStore() { return stateStore; @@ -3841,8 +3883,8 @@ public static Builder builder() { @Override public void close() { - // No-op for the core ReActAgent. Subclasses / wrappers (HarnessAgent) may release - // additional resources here. + stateCache.clear(); + permissionEngineCache.clear(); } // ==================== Builder ==================== diff --git a/agentscope-core/src/main/java/io/agentscope/core/agent/AgentBase.java b/agentscope-core/src/main/java/io/agentscope/core/agent/AgentBase.java index a6b0248b6..b1d51a2a9 100644 --- a/agentscope-core/src/main/java/io/agentscope/core/agent/AgentBase.java +++ b/agentscope-core/src/main/java/io/agentscope/core/agent/AgentBase.java @@ -319,7 +319,11 @@ private Mono runLifecycleBody( .onErrorResume( createErrorHandler( msgs.toArray(new Msg[0])))); - return scope == null ? body : body.contextWrite(c -> c.put(CALL_SCOPE_KEY, scope)); + if (scope == null) { + return body; + } + return body.contextWrite(c -> c.put(CALL_SCOPE_KEY, scope)) + .doFinally(signal -> afterCallScopeExecution(scope)); } /** @@ -595,6 +599,15 @@ protected Object beforeAgentExecution(List msgs, RuntimeContext rc) { */ protected void afterAgentExecution() {} + /** + * Invoked when a call-specific scope terminates after success, error, or cancellation. + * Subclasses can release resources owned by the exact scope without consulting shared + * instance fields that may already refer to another concurrent call. + * + * @param callScope the scope returned by {@link #beforeAgentExecution(List, RuntimeContext)} + */ + protected void afterCallScopeExecution(Object callScope) {} + /** * Pushes {@code ctx} to all {@link RuntimeContextAware} hooks registered for this agent. The * per-call {@link RuntimeContext} itself is no longer stored on a shared instance field; it diff --git a/agentscope-core/src/test/java/io/agentscope/core/agent/AgentBaseTest.java b/agentscope-core/src/test/java/io/agentscope/core/agent/AgentBaseTest.java index 9e2919643..cdfb9bb2c 100644 --- a/agentscope-core/src/test/java/io/agentscope/core/agent/AgentBaseTest.java +++ b/agentscope-core/src/test/java/io/agentscope/core/agent/AgentBaseTest.java @@ -15,6 +15,7 @@ */ package io.agentscope.core.agent; +import static org.junit.jupiter.api.Assertions.assertDoesNotThrow; import static org.junit.jupiter.api.Assertions.assertEquals; import static org.junit.jupiter.api.Assertions.assertNotEquals; import static org.junit.jupiter.api.Assertions.assertNotNull; @@ -110,6 +111,18 @@ protected Mono handleInterrupt(InterruptContext context, Msg... originalArg } } + static class ScopedTestAgent extends TestAgent { + + ScopedTestAgent(String name) { + super(name); + } + + @Override + protected Object beforeAgentExecution(List msgs, RuntimeContext rc) { + return new Object(); + } + } + @BeforeEach void setUp() { agent = new TestAgent(TestConstants.TEST_AGENT_NAME); @@ -153,6 +166,19 @@ void testSingleMessageInput() { TestConstants.TEST_ASSISTANT_RESPONSE, text, "Response text should match expected"); } + @Test + @DisplayName("Should support call scope with the default terminal callback") + void testDefaultAfterCallScopeExecution() { + ScopedTestAgent scopedAgent = new ScopedTestAgent(TestConstants.TEST_AGENT_NAME); + Msg userMsg = TestUtils.createUserMessage("User", TestConstants.TEST_USER_INPUT); + + assertDoesNotThrow( + () -> + scopedAgent + .call(userMsg) + .block(Duration.ofMillis(TestConstants.DEFAULT_TEST_TIMEOUT_MS))); + } + @Test @DisplayName("Should handle multiple message input") void testMultipleMessageInput() { diff --git a/agentscope-core/src/test/java/io/agentscope/core/agent/ReActAgentPerSessionStateTest.java b/agentscope-core/src/test/java/io/agentscope/core/agent/ReActAgentPerSessionStateTest.java index 4f7a8a09e..6c14de51e 100644 --- a/agentscope-core/src/test/java/io/agentscope/core/agent/ReActAgentPerSessionStateTest.java +++ b/agentscope-core/src/test/java/io/agentscope/core/agent/ReActAgentPerSessionStateTest.java @@ -19,6 +19,7 @@ import static org.junit.jupiter.api.Assertions.assertFalse; import static org.junit.jupiter.api.Assertions.assertNotSame; import static org.junit.jupiter.api.Assertions.assertSame; +import static org.junit.jupiter.api.Assertions.assertThrows; import static org.junit.jupiter.api.Assertions.assertTrue; import io.agentscope.core.ReActAgent; @@ -38,9 +39,11 @@ import io.agentscope.core.state.InMemoryAgentStateStore; import io.agentscope.core.state.legacy.ToolkitState; import io.agentscope.core.tool.Toolkit; +import java.lang.reflect.Field; import java.time.Duration; import java.util.ArrayList; import java.util.List; +import java.util.Map; import java.util.concurrent.CompletableFuture; import java.util.concurrent.CountDownLatch; import java.util.concurrent.TimeUnit; @@ -48,6 +51,7 @@ import java.util.stream.IntStream; import org.junit.jupiter.api.DisplayName; import org.junit.jupiter.api.Test; +import reactor.core.Disposable; import reactor.core.publisher.Flux; import reactor.core.publisher.Mono; import reactor.core.scheduler.Schedulers; @@ -72,6 +76,42 @@ protected Flux doStream( } } + private static final class FailingModel extends ChatModelBase { + @Override + public String getModelName() { + return "failing"; + } + + @Override + protected Flux doStream( + List messages, List tools, GenerateOptions options) { + return Flux.error(new IllegalStateException("model failed")); + } + } + + private static final class NeverModel extends ChatModelBase { + private final CountDownLatch subscribed; + + private NeverModel(CountDownLatch subscribed) { + this.subscribed = subscribed; + } + + @Override + public String getModelName() { + return "never"; + } + + @Override + protected Flux doStream( + List messages, List tools, GenerateOptions options) { + return Flux.defer( + () -> { + subscribed.countDown(); + return Flux.never(); + }); + } + } + private ReActAgent agent(InMemoryAgentStateStore store) { return ReActAgent.builder() .name("asst") @@ -179,6 +219,170 @@ void savePersistsPerSlot() { assertEquals("", other.getSummary()); } + @Test + @DisplayName("completed calls release persistent per-session caches") + void completedCallsReleasePersistentCaches() { + ReActAgent agent = agent(new InMemoryAgentStateStore()); + agent.getAgentState(); + int initialStateCacheSize = cacheSize(agent, "stateCache"); + int initialPermissionCacheSize = cacheSize(agent, "permissionEngineCache"); + + for (int i = 0; i < 32; i++) { + RuntimeContext ctx = + RuntimeContext.builder().userId("user").sessionId("session-" + i).build(); + agent.call(List.of(userMsg("hello-" + i)), ctx).block(Duration.ofSeconds(5)); + } + + assertEquals(initialStateCacheSize, cacheSize(agent, "stateCache")); + assertEquals(initialPermissionCacheSize, cacheSize(agent, "permissionEngineCache")); + } + + @Test + @DisplayName("failed calls release persistent per-session caches") + void failedCallsReleasePersistentCaches() { + ReActAgent agent = + ReActAgent.builder() + .name("asst") + .sysPrompt("hi") + .model(new FailingModel()) + .stateStore(new InMemoryAgentStateStore()) + .build(); + RuntimeContext ctx = RuntimeContext.builder().userId("user").sessionId("failed").build(); + agent.getAgentState(); + int initialStateCacheSize = cacheSize(agent, "stateCache"); + int initialPermissionCacheSize = cacheSize(agent, "permissionEngineCache"); + + assertThrows( + IllegalStateException.class, + () -> agent.call(List.of(userMsg("hello")), ctx).block(Duration.ofSeconds(5))); + + assertEquals(initialStateCacheSize, cacheSize(agent, "stateCache")); + assertEquals(initialPermissionCacheSize, cacheSize(agent, "permissionEngineCache")); + } + + @Test + @DisplayName("cancelled calls release persistent per-session caches") + void cancelledCallsReleasePersistentCaches() throws Exception { + CountDownLatch subscribed = new CountDownLatch(1); + ReActAgent agent = + ReActAgent.builder() + .name("asst") + .sysPrompt("hi") + .model(new NeverModel(subscribed)) + .stateStore(new InMemoryAgentStateStore()) + .build(); + RuntimeContext ctx = RuntimeContext.builder().userId("user").sessionId("cancelled").build(); + agent.getAgentState(); + int initialStateCacheSize = cacheSize(agent, "stateCache"); + int initialPermissionCacheSize = cacheSize(agent, "permissionEngineCache"); + + Disposable call = agent.call(List.of(userMsg("hello")), ctx).subscribe(); + assertTrue(subscribed.await(5, TimeUnit.SECONDS), "model stream should start"); + call.dispose(); + + assertEquals(initialStateCacheSize, cacheSize(agent, "stateCache")); + assertEquals(initialPermissionCacheSize, cacheSize(agent, "permissionEngineCache")); + } + + @Test + @DisplayName("session eviction removes only the selected in-memory slot") + void evictSessionRemovesOnlySelectedSlot() { + ReActAgent agent = + ReActAgent.builder().name("asst").sysPrompt("hi").model(new NoopModel()).build(); + RuntimeContext first = RuntimeContext.builder().userId("u").sessionId("first").build(); + RuntimeContext second = RuntimeContext.builder().userId("u").sessionId("second").build(); + agent.call(List.of(userMsg("first")), first).block(Duration.ofSeconds(5)); + agent.call(List.of(userMsg("second")), second).block(Duration.ofSeconds(5)); + AgentState firstState = agent.getAgentState(first); + AgentState secondState = agent.getAgentState(second); + + agent.evictSession("u", "first"); + + assertNotSame(firstState, agent.getAgentState(first)); + assertSame(secondState, agent.getAgentState(second)); + + AgentState defaultState = agent.getAgentState(null, agent.getDefaultSessionId()); + agent.evictSession(null, null); + assertNotSame(defaultState, agent.getAgentState(null, agent.getDefaultSessionId())); + + AgentState blankSessionState = agent.getAgentState(null, agent.getDefaultSessionId()); + agent.evictSession(null, " "); + assertNotSame(blankSessionState, agent.getAgentState(null, agent.getDefaultSessionId())); + } + + @Test + @DisplayName("user eviction removes all and only that user's in-memory slots") + void evictUserRemovesOnlySelectedUsersSlots() { + ReActAgent agent = + ReActAgent.builder().name("asst").sysPrompt("hi").model(new NoopModel()).build(); + agent.call( + List.of(userMsg("first")), + RuntimeContext.builder().userId("u1").sessionId("first").build()) + .block(Duration.ofSeconds(5)); + agent.call( + List.of(userMsg("second")), + RuntimeContext.builder().userId("u1").sessionId("second").build()) + .block(Duration.ofSeconds(5)); + agent.call( + List.of(userMsg("other")), + RuntimeContext.builder().userId("u2").sessionId("first").build()) + .block(Duration.ofSeconds(5)); + AgentState first = agent.getAgentState("u1", "first"); + AgentState second = agent.getAgentState("u1", "second"); + AgentState other = agent.getAgentState("u2", "first"); + + agent.evictUser("u1"); + + assertNotSame(first, agent.getAgentState("u1", "first")); + assertNotSame(second, agent.getAgentState("u1", "second")); + assertSame(other, agent.getAgentState("u2", "first")); + + AgentState anonymous = agent.getAgentState(null, "anonymous"); + agent.evictUser(null); + assertNotSame(anonymous, agent.getAgentState(null, "anonymous")); + + AgentState blankUser = agent.getAgentState(null, "blank-user"); + agent.evictUser(" "); + assertNotSame(blankUser, agent.getAgentState(null, "blank-user")); + } + + @Test + @DisplayName("close clears all in-memory session caches") + void closeClearsSessionCaches() { + ReActAgent agent = + ReActAgent.builder().name("asst").sysPrompt("hi").model(new NoopModel()).build(); + agent.call( + List.of(userMsg("hello")), + RuntimeContext.builder().userId("u").sessionId("session").build()) + .block(Duration.ofSeconds(5)); + assertTrue(cacheSize(agent, "stateCache") > 0); + assertTrue(cacheSize(agent, "permissionEngineCache") > 0); + + agent.close(); + + assertEquals(0, cacheSize(agent, "stateCache")); + assertEquals(0, cacheSize(agent, "permissionEngineCache")); + } + + @Test + @DisplayName("calls without a state store retain in-memory per-session state") + void callsWithoutStateStoreRetainSessionCaches() { + ReActAgent agent = + ReActAgent.builder().name("asst").sysPrompt("hi").model(new NoopModel()).build(); + agent.getAgentState(); + int initialStateCacheSize = cacheSize(agent, "stateCache"); + int initialPermissionCacheSize = cacheSize(agent, "permissionEngineCache"); + + for (int i = 0; i < 3; i++) { + RuntimeContext ctx = + RuntimeContext.builder().userId("user").sessionId("session-" + i).build(); + agent.call(List.of(userMsg("hello-" + i)), ctx).block(Duration.ofSeconds(5)); + } + + assertEquals(initialStateCacheSize + 3, cacheSize(agent, "stateCache")); + assertEquals(initialPermissionCacheSize + 3, cacheSize(agent, "permissionEngineCache")); + } + @Test @DisplayName("user interrupt persists recovery state to the store") void userInterruptPersistsRecoveryState() throws Exception { @@ -276,6 +480,16 @@ private static List allText(AgentState state) { return out; } + private static int cacheSize(ReActAgent agent, String fieldName) { + try { + Field field = ReActAgent.class.getDeclaredField(fieldName); + field.setAccessible(true); + return ((Map) field.get(agent)).size(); + } catch (ReflectiveOperationException e) { + throw new AssertionError("Failed to inspect " + fieldName, e); + } + } + @Test @DisplayName( "concurrent calls to distinct sessions run in parallel without cross-contamination")