diff --git a/agentscope-extensions/agentscope-extensions-session-redis/src/main/java/io/agentscope/core/session/redis/RedisClientAdapter.java b/agentscope-extensions/agentscope-extensions-session-redis/src/main/java/io/agentscope/core/session/redis/RedisClientAdapter.java index 1f91a9d5ea..98878f9379 100644 --- a/agentscope-extensions/agentscope-extensions-session-redis/src/main/java/io/agentscope/core/session/redis/RedisClientAdapter.java +++ b/agentscope-extensions/agentscope-extensions-session-redis/src/main/java/io/agentscope/core/session/redis/RedisClientAdapter.java @@ -92,6 +92,16 @@ public interface RedisClientAdapter { */ void addToSet(String key, String member); + /** + * Remove a member from a set. + * + * @param key the Redis set key + * @param member the member to remove + */ + default void removeFromSet(String key, String member) { + // Default no-op for custom adapters that do not track session keys in a Redis set. + } + /** * Get all members of a set. * diff --git a/agentscope-extensions/agentscope-extensions-session-redis/src/main/java/io/agentscope/core/session/redis/RedisSession.java b/agentscope-extensions/agentscope-extensions-session-redis/src/main/java/io/agentscope/core/session/redis/RedisSession.java index 7bb86e3414..14a3c26d19 100644 --- a/agentscope-extensions/agentscope-extensions-session-redis/src/main/java/io/agentscope/core/session/redis/RedisSession.java +++ b/agentscope-extensions/agentscope-extensions-session-redis/src/main/java/io/agentscope/core/session/redis/RedisSession.java @@ -348,6 +348,19 @@ public void delete(SessionKey sessionKey) { } } + @Override + public void delete(SessionKey sessionKey, String key) { + String sessionId = sessionKey.toIdentifier(); + String keysKey = getKeysKey(sessionId); + + try { + client.deleteKeys(getStateKey(sessionId, key)); + client.removeFromSet(keysKey, key); + } catch (Exception e) { + throw new RuntimeException("Failed to delete state: " + key, e); + } + } + @Override public Set listSessionKeys() { try { diff --git a/agentscope-extensions/agentscope-extensions-session-redis/src/main/java/io/agentscope/core/session/redis/jedis/JedisClientAdapter.java b/agentscope-extensions/agentscope-extensions-session-redis/src/main/java/io/agentscope/core/session/redis/jedis/JedisClientAdapter.java index 83e0bcfa51..ebd496f087 100644 --- a/agentscope-extensions/agentscope-extensions-session-redis/src/main/java/io/agentscope/core/session/redis/jedis/JedisClientAdapter.java +++ b/agentscope-extensions/agentscope-extensions-session-redis/src/main/java/io/agentscope/core/session/redis/jedis/JedisClientAdapter.java @@ -136,6 +136,11 @@ public void addToSet(String key, String member) { unifiedJedis.sadd(key, member); } + @Override + public void removeFromSet(String key, String member) { + unifiedJedis.srem(key, member); + } + @Override public Set getSetMembers(String key) { return unifiedJedis.smembers(key); diff --git a/agentscope-extensions/agentscope-extensions-session-redis/src/main/java/io/agentscope/core/session/redis/lettuce/LettuceClientAdapter.java b/agentscope-extensions/agentscope-extensions-session-redis/src/main/java/io/agentscope/core/session/redis/lettuce/LettuceClientAdapter.java index 8ca1304e23..89c9671934 100644 --- a/agentscope-extensions/agentscope-extensions-session-redis/src/main/java/io/agentscope/core/session/redis/lettuce/LettuceClientAdapter.java +++ b/agentscope-extensions/agentscope-extensions-session-redis/src/main/java/io/agentscope/core/session/redis/lettuce/LettuceClientAdapter.java @@ -217,6 +217,15 @@ public void addToSet(String key, String member) { } } + @Override + public void removeFromSet(String key, String member) { + if (commands != null) { + commands.srem(key, member); + } else { + clusterCommands.srem(key, member); + } + } + @Override public Set getSetMembers(String key) { if (commands != null) { diff --git a/agentscope-extensions/agentscope-extensions-session-redis/src/main/java/io/agentscope/core/session/redis/redisson/RedissonClientAdapter.java b/agentscope-extensions/agentscope-extensions-session-redis/src/main/java/io/agentscope/core/session/redis/redisson/RedissonClientAdapter.java index 542793afca..7b631279db 100644 --- a/agentscope-extensions/agentscope-extensions-session-redis/src/main/java/io/agentscope/core/session/redis/redisson/RedissonClientAdapter.java +++ b/agentscope-extensions/agentscope-extensions-session-redis/src/main/java/io/agentscope/core/session/redis/redisson/RedissonClientAdapter.java @@ -184,6 +184,12 @@ public void addToSet(String key, String member) { rSet.add(member); } + @Override + public void removeFromSet(String key, String member) { + RSet rSet = redissonClient.getSet(key, StringCodec.INSTANCE); + rSet.remove(member); + } + @Override public Set getSetMembers(String key) { RSet rSet = redissonClient.getSet(key, StringCodec.INSTANCE); diff --git a/agentscope-extensions/agentscope-extensions-session-redis/src/test/java/io/agentscope/core/session/redis/JedisSessionTest.java b/agentscope-extensions/agentscope-extensions-session-redis/src/test/java/io/agentscope/core/session/redis/JedisSessionTest.java index 5eb156ea4c..26e9676031 100644 --- a/agentscope-extensions/agentscope-extensions-session-redis/src/test/java/io/agentscope/core/session/redis/JedisSessionTest.java +++ b/agentscope-extensions/agentscope-extensions-session-redis/src/test/java/io/agentscope/core/session/redis/JedisSessionTest.java @@ -27,12 +27,17 @@ import static org.mockito.Mockito.verify; import static org.mockito.Mockito.when; +import io.agentscope.core.agent.Agent; +import io.agentscope.core.shutdown.GracefulShutdownManager; +import io.agentscope.core.shutdown.ShutdownInterruptedState; import io.agentscope.core.state.SessionKey; import io.agentscope.core.state.SimpleSessionKey; import io.agentscope.core.state.State; import java.util.ArrayList; +import java.util.HashMap; import java.util.HashSet; import java.util.List; +import java.util.Map; import java.util.Optional; import java.util.Set; import org.junit.jupiter.api.BeforeEach; @@ -221,6 +226,66 @@ void testDeleteSession() { verify(unifiedJedis).smembers("agentscope:session:session1:_keys"); } + @Test + @DisplayName("Should delete single state key and remove tracking metadata") + void testDeleteSingleStateKey() { + RedisSession session = + RedisSession.builder() + .jedisClient(unifiedJedis) + .keyPrefix("agentscope:session:") + .build(); + + SessionKey sessionKey = SimpleSessionKey.of("session1"); + session.delete(sessionKey, "shutdown_interrupted"); + + verify(unifiedJedis).del(new String[] {"agentscope:session:session1:shutdown_interrupted"}); + verify(unifiedJedis).srem("agentscope:session:session1:_keys", "shutdown_interrupted"); + } + + @Test + @DisplayName("Should not return single state after deleting it") + void testDeleteSingleStateKeyRemovesPersistedValue() { + RedisSession session = + RedisSession.builder() + .clientAdapter(new StatefulRedisClientAdapter()) + .keyPrefix("agentscope:session:") + .build(); + + SessionKey sessionKey = SimpleSessionKey.of("session1"); + session.save(sessionKey, "shutdown_interrupted", new TestState("interrupted", 1)); + + assertTrue(session.get(sessionKey, "shutdown_interrupted", TestState.class).isPresent()); + + session.delete(sessionKey, "shutdown_interrupted"); + + assertFalse(session.get(sessionKey, "shutdown_interrupted", TestState.class).isPresent()); + } + + @Test + @DisplayName("Should clear shutdown interrupted flag after first manager check") + void testManagerCheckAndClearShutdownInterruptedWithRedisSession() { + GracefulShutdownManager manager = GracefulShutdownManager.getInstance(); + manager.resetForTesting(); + try { + RedisSession session = + RedisSession.builder() + .clientAdapter(new StatefulRedisClientAdapter()) + .keyPrefix("agentscope:session:") + .build(); + SessionKey sessionKey = SimpleSessionKey.of("session1"); + Agent agent = mock(Agent.class); + when(agent.getAgentId()).thenReturn("agent-1"); + + session.save(sessionKey, "shutdown_interrupted", new ShutdownInterruptedState(true)); + manager.bindSession(agent, session, sessionKey); + + assertTrue(manager.checkAndClearShutdownInterrupted(agent)); + assertFalse(manager.checkAndClearShutdownInterrupted(agent)); + } finally { + manager.resetForTesting(); + } + } + @Test @DisplayName("Should list all session keys") void testListSessionKeys() { @@ -280,4 +345,76 @@ void testCloseShutsDownClient() { /** Simple test state record for testing. */ public record TestState(String value, int count) implements State {} + + private static class StatefulRedisClientAdapter implements RedisClientAdapter { + + private final Map values = new HashMap<>(); + private final Map> sets = new HashMap<>(); + + @Override + public void set(String key, String value) { + values.put(key, value); + } + + @Override + public String get(String key) { + return values.get(key); + } + + @Override + public void rightPushList(String key, String value) {} + + @Override + public List rangeList(String key, long start, long end) { + return List.of(); + } + + @Override + public long getListLength(String key) { + return 0; + } + + @Override + public void deleteKeys(String... keys) { + for (String key : keys) { + values.remove(key); + } + } + + @Override + public void addToSet(String key, String member) { + sets.computeIfAbsent(key, unused -> new HashSet<>()).add(member); + } + + @Override + public void removeFromSet(String key, String member) { + Set members = sets.get(key); + if (members != null) { + members.remove(member); + } + } + + @Override + public Set getSetMembers(String key) { + return sets.getOrDefault(key, Set.of()); + } + + @Override + public long getSetSize(String key) { + return getSetMembers(key).size(); + } + + @Override + public boolean keyExists(String key) { + return values.containsKey(key) || sets.containsKey(key); + } + + @Override + public Set findKeysByPattern(String pattern) { + return Set.of(); + } + + @Override + public void close() {} + } } diff --git a/agentscope-extensions/agentscope-extensions-session-redis/src/test/java/io/agentscope/core/session/redis/LettuceSessionTest.java b/agentscope-extensions/agentscope-extensions-session-redis/src/test/java/io/agentscope/core/session/redis/LettuceSessionTest.java index 8a8b8e51dd..14f178c7ce 100644 --- a/agentscope-extensions/agentscope-extensions-session-redis/src/test/java/io/agentscope/core/session/redis/LettuceSessionTest.java +++ b/agentscope-extensions/agentscope-extensions-session-redis/src/test/java/io/agentscope/core/session/redis/LettuceSessionTest.java @@ -206,6 +206,22 @@ void testDeleteSession() { verify(commands).smembers("agentscope:session:session1:_keys"); } + @Test + @DisplayName("Should delete single state key and remove tracking metadata") + void testDeleteSingleStateKey() { + RedisSession session = + RedisSession.builder() + .lettuceClient(redisClient) + .keyPrefix("agentscope:session:") + .build(); + + SessionKey sessionKey = SimpleSessionKey.of("session1"); + session.delete(sessionKey, "shutdown_interrupted"); + + verify(commands).del("agentscope:session:session1:shutdown_interrupted"); + verify(commands).srem("agentscope:session:session1:_keys", "shutdown_interrupted"); + } + @Test @DisplayName("Should list session keys correctly") void testListSessionKeys() { @@ -424,6 +440,23 @@ void testDeleteSessionCluster() { verify(clusterCommands).smembers("agentscope:session:clusterSession:_keys"); } + @Test + @DisplayName("Should delete single state key and remove tracking metadata in cluster mode") + void testDeleteSingleStateKeyCluster() { + RedisSession session = + RedisSession.builder() + .lettuceClusterClient(redisClusterClient) + .keyPrefix("agentscope:session:") + .build(); + + SessionKey sessionKey = SimpleSessionKey.of("clusterSession"); + session.delete(sessionKey, "shutdown_interrupted"); + + verify(clusterCommands).del("agentscope:session:clusterSession:shutdown_interrupted"); + verify(clusterCommands) + .srem("agentscope:session:clusterSession:_keys", "shutdown_interrupted"); + } + @Test @DisplayName("Should list session keys correctly in cluster mode") void testListSessionKeysCluster() { diff --git a/agentscope-extensions/agentscope-extensions-session-redis/src/test/java/io/agentscope/core/session/redis/RedissonSessionTest.java b/agentscope-extensions/agentscope-extensions-session-redis/src/test/java/io/agentscope/core/session/redis/RedissonSessionTest.java index 00dc98b5a9..adc3672f0e 100644 --- a/agentscope-extensions/agentscope-extensions-session-redis/src/test/java/io/agentscope/core/session/redis/RedissonSessionTest.java +++ b/agentscope-extensions/agentscope-extensions-session-redis/src/test/java/io/agentscope/core/session/redis/RedissonSessionTest.java @@ -251,6 +251,22 @@ void testDeleteSession() { verify(rSet).iterator(); } + @Test + @DisplayName("Should delete single state key and remove tracking metadata") + void testDeleteSingleStateKey() { + RedisSession session = + RedisSession.builder() + .redissonClient(redissonClient) + .keyPrefix("agentscope:session:") + .build(); + + SessionKey sessionKey = SimpleSessionKey.of("session1"); + session.delete(sessionKey, "shutdown_interrupted"); + + verify(keys).delete("agentscope:session:session1:shutdown_interrupted"); + verify(rSet).remove("shutdown_interrupted"); + } + @Test @DisplayName("Should list all session keys") void testListSessionKeys() {