Merge pull request #62 from usnavy13/session-fix

State size threshold for Redis with MinIO overflow

Author: usnavy13

src/api/files.py

@@ -115,18 +115,18 @@ async def upload_file(
             # Read file content
             content = await file.read()
 
-            # Store file directly
+            # Sanitize filename to match what will be used in container
+            sanitized_name = OutputProcessor.sanitize_filename(file.filename)
+
+            # Store with sanitized name so MinIO, sandbox, and cleanup all use the same name
             file_id = await file_service.store_uploaded_file(
                 session_id=session_id,
-                filename=file.filename,
+                filename=sanitized_name,
                 content=content,
                 content_type=file.content_type,
                 is_agent_file=is_agent_file,
             )
 
-            # Sanitize filename to match what will be used in container
-            sanitized_name = OutputProcessor.sanitize_filename(file.filename)
-
             uploaded_files.append(
                 {
                     "id": file_id,

src/config/__init__.py

@@ -201,6 +201,12 @@ class Settings(BaseSettings):
     state_capture_on_error: bool = Field(
         default=False, description="Capture and persist state even when execution fails"
     )
+    state_max_redis_size_mb: int = Field(
+        default=100,
+        ge=1,
+        le=500,
+        description="Max state size (MB, raw bytes) for Redis storage. Larger states go directly to MinIO",
+    )
 
     # State Archival Configuration - Hybrid Redis + MinIO storage
     state_archive_enabled: bool = Field(

src/core/pool.py

@@ -40,7 +40,7 @@ def _initialize(self) -> None:
                 redis_url,
                 max_connections=20,  # Shared across all services
                 decode_responses=True,
-                socket_timeout=5.0,
+                socket_timeout=30.0,
                 socket_connect_timeout=5.0,
                 retry_on_timeout=True,
             )

src/services/orchestrator.py

@@ -520,11 +520,51 @@ async def _save_state(self, ctx: ExecutionContext) -> None:
 
         if ctx.new_state:
             try:
-                success, state_hash = await self.state_service.save_state(
-                    ctx.session_id,
-                    ctx.new_state,
-                    ttl_seconds=settings.state_ttl_seconds,
-                )
+                import base64
+
+                raw_size = len(base64.b64decode(ctx.new_state))
+                max_redis_bytes = settings.state_max_redis_size_mb * 1024 * 1024
+
+                if raw_size > max_redis_bytes:
+                    # Large state: store blob in MinIO, pointer in Redis
+                    logger.info(
+                        "State exceeds Redis threshold, storing in MinIO",
+                        session_id=ctx.session_id[:12],
+                        state_size_mb=round(raw_size / 1024 / 1024, 1),
+                        threshold_mb=settings.state_max_redis_size_mb,
+                    )
+                    archived = False
+                    if self.state_archival_service:
+                        archived = await self.state_archival_service.archive_state(
+                            ctx.session_id, ctx.new_state
+                        )
+                    if archived:
+                        success, state_hash = (
+                            await self.state_service.save_state_pointer(
+                                ctx.session_id,
+                                ctx.new_state,
+                                ttl_seconds=settings.state_ttl_seconds,
+                            )
+                        )
+                    else:
+                        # MinIO failed, fall back to Redis anyway
+                        logger.warning(
+                            "MinIO archival failed, falling back to Redis",
+                            session_id=ctx.session_id[:12],
+                        )
+                        success, state_hash = await self.state_service.save_state(
+                            ctx.session_id,
+                            ctx.new_state,
+                            ttl_seconds=settings.state_ttl_seconds,
+                        )
+                else:
+                    # Normal path: store in Redis
+                    success, state_hash = await self.state_service.save_state(
+                        ctx.session_id,
+                        ctx.new_state,
+                        ttl_seconds=settings.state_ttl_seconds,
+                    )
+
                 if success:
                     ctx.new_state_hash = state_hash
 

src/services/state.py

@@ -180,6 +180,71 @@ async def save_state(
             )
             return False, None
 
+    async def save_state_pointer(
+        self,
+        session_id: str,
+        state_b64: str,
+        ttl_seconds: Optional[int] = None,
+    ) -> Tuple[bool, Optional[str]]:
+        """Save only hash and metadata to Redis (state blob stored in MinIO).
+
+        Used when state exceeds the Redis size threshold. The full state
+        is stored in MinIO; Redis only holds the hash and metadata for
+        fast lookups. The orchestrator's _load_state MinIO fallback
+        handles retrieval.
+
+        Args:
+            session_id: Session identifier
+            state_b64: Base64-encoded state (used to compute hash/size, not stored in Redis)
+            ttl_seconds: TTL in seconds (default from settings)
+
+        Returns:
+            Tuple of (success: bool, state_hash: Optional[str])
+        """
+        if not state_b64:
+            return True, None
+
+        if ttl_seconds is None:
+            ttl_seconds = settings.state_ttl_seconds
+
+        try:
+            raw_bytes = base64.b64decode(state_b64)
+            state_hash = self.compute_hash(raw_bytes)
+            now = datetime.now(timezone.utc)
+
+            pipe = self.redis.pipeline(transaction=True)
+
+            # Save hash (small — just the SHA256 string)
+            pipe.setex(self._hash_key(session_id), ttl_seconds, state_hash)
+
+            # Save metadata with storage location marker
+            meta = json.dumps(
+                {
+                    "size_bytes": len(raw_bytes),
+                    "hash": state_hash,
+                    "created_at": now.isoformat(),
+                    "storage": "minio",
+                }
+            )
+            pipe.setex(self._meta_key(session_id), ttl_seconds, meta)
+
+            await pipe.execute()
+
+            logger.info(
+                "Saved state pointer to Redis (blob in MinIO)",
+                session_id=session_id[:12],
+                state_size=len(raw_bytes),
+                hash=state_hash[:12],
+            )
+            return True, state_hash
+        except Exception as e:
+            logger.error(
+                "Failed to save state pointer",
+                session_id=session_id[:12],
+                error=str(e),
+            )
+            return False, None
+
     async def delete_state(self, session_id: str) -> bool:
         """Delete state for a session.
 

src/services/state_archival.py

@@ -187,10 +187,26 @@ async def restore_state(self, session_id: str) -> Optional[str]:
 
             state_data = state_bytes.decode("utf-8")
 
-            # Restore to Redis for fast access
-            await self.state_service.save_state(
-                session_id, state_data, ttl_seconds=settings.state_ttl_seconds
-            )
+            # Only restore to Redis if under the size threshold
+            import base64 as _b64
+
+            raw_size = len(_b64.b64decode(state_data))
+            max_redis_bytes = settings.state_max_redis_size_mb * 1024 * 1024
+
+            if raw_size <= max_redis_bytes:
+                await self.state_service.save_state(
+                    session_id, state_data, ttl_seconds=settings.state_ttl_seconds
+                )
+            else:
+                # Too large for Redis — save only pointer
+                await self.state_service.save_state_pointer(
+                    session_id, state_data, ttl_seconds=settings.state_ttl_seconds
+                )
+                logger.info(
+                    "State too large for Redis, kept in MinIO only",
+                    session_id=session_id[:12],
+                    state_size_mb=round(raw_size / 1024 / 1024, 1),
+                )
 
             logger.info(
                 "Restored state from MinIO",
@@ -504,10 +520,16 @@ async def restore_state_by_hash(self, state_hash: str) -> Optional[str]:
 
             state_data = state_bytes.decode("utf-8")
 
-            # Restore to Redis for fast access
-            await self.state_service.save_state_by_hash(
-                state_hash, state_data, ttl_seconds=settings.state_ttl_seconds
-            )
+            # Only restore to Redis if under the size threshold
+            import base64 as _b64
+
+            raw_size = len(_b64.b64decode(state_data))
+            max_redis_bytes = settings.state_max_redis_size_mb * 1024 * 1024
+
+            if raw_size <= max_redis_bytes:
+                await self.state_service.save_state_by_hash(
+                    state_hash, state_data, ttl_seconds=settings.state_ttl_seconds
+                )
 
             logger.debug(
                 "Restored state by hash from MinIO",