fix state_service redis

Author: zhilingluo

src/agentscope_runtime/engine/services/agent_state/redis_state_service.py

@@ -21,29 +21,68 @@ def __init__(
         self,
         redis_url: str = "redis://localhost:6379/0",
         redis_client: Optional[aioredis.Redis] = None,
+        socket_timeout: Optional[float] = 5.0,
+        socket_connect_timeout: Optional[float] = 5.0,
+        max_connections: Optional[int] = 50,
+        retry_on_timeout: bool = True,
+        ttl_seconds: Optional[int] = 3600,  # 1 hour in seconds
+        health_check_interval: Optional[float] = 30.0,
+        socket_keepalive: bool = True,
     ):
+        """
+        Initialize RedisStateService.
+
+        Args:
+            redis_url: Redis connection URL
+            redis_client: Optional pre-configured Redis client
+            socket_timeout: Socket timeout in seconds (default: 5.0)
+            socket_connect_timeout: Socket connect timeout in seconds
+            (default: 5.0)
+            max_connections: Maximum number of connections in the pool
+            (default: 50)
+            retry_on_timeout: Whether to retry on timeout (default: True)
+            ttl_seconds: Time-to-live in seconds for state data. If None,
+            data never expires (default: 3600, i.e., 1 hour)
+            health_check_interval: Interval in seconds for health checks on
+            idle connections (default: 30.0).
+                Connections idle longer than this will be checked before reuse.
+                Set to 0 to disable.
+            socket_keepalive: Enable TCP keepalive to prevent
+            silent disconnections (default: True)
+        """
         self._redis_url = redis_url
         self._redis = redis_client
-        self._health = False
+        self._socket_timeout = socket_timeout
+        self._socket_connect_timeout = socket_connect_timeout
+        self._max_connections = max_connections
+        self._retry_on_timeout = retry_on_timeout
+        self._ttl_seconds = ttl_seconds
+        self._health_check_interval = health_check_interval
+        self._socket_keepalive = socket_keepalive
 
     async def start(self) -> None:
-        """Initialize the Redis connection."""
+        """Starts the Redis connection with proper timeout and connection
+        pool settings."""
         if self._redis is None:
             self._redis = aioredis.from_url(
                 self._redis_url,
                 decode_responses=True,
+                socket_timeout=self._socket_timeout,
+                socket_connect_timeout=self._socket_connect_timeout,
+                max_connections=self._max_connections,
+                retry_on_timeout=self._retry_on_timeout,
+                health_check_interval=self._health_check_interval,
+                socket_keepalive=self._socket_keepalive,
             )
-        self._health = True
 
     async def stop(self) -> None:
-        """Close the Redis connection."""
+        """Closes the Redis connection."""
         if self._redis:
             await self._redis.close()
             self._redis = None
-        self._health = False
 
     async def health(self) -> bool:
-        """Service health check."""
+        """Checks the health of the service."""
         if not self._redis:
             return False
         try:
@@ -81,6 +120,11 @@ async def save_state(
                 round_id = 1
 
         await self._redis.hset(key, round_id, json.dumps(state))
+
+        # Set TTL for the state key if configured
+        if self._ttl_seconds is not None:
+            await self._redis.expire(key, self._ttl_seconds)
+
         return round_id
 
     async def export_state(
@@ -110,4 +154,9 @@ async def export_state(
 
         if state_json is None:
             return None
+
+        # Refresh TTL when accessing the state
+        if self._ttl_seconds is not None:
+            await self._redis.expire(key, self._ttl_seconds)
+
         return json.loads(state_json)

src/agentscope_runtime/engine/services/memory/redis_memory_service.py

@@ -197,6 +197,7 @@ async def search_memory(
             await self._redis.expire(key, ttl_seconds)
 
         return result
+
     async def get_query_text(self, message: Message) -> str:
         if message:
             if message.type == MessageType.MESSAGE:
@@ -233,7 +234,8 @@ async def list_memory(
                 # we need all previous messages for proper ordering)
                 # For now, we keep loading all for correctness
 
-        # Refresh TTL on active use to keep memory alive, mirroring get_session behavior
+        # Refresh TTL on active use to keep memory alive,
+        # mirroring get_session behavior
         if getattr(self, "_ttl_seconds", None):
             await self._redis.expire(key, self._ttl_seconds)
         return all_msgs[start_index:end_index]