feat(chat stream): 支持流式输出恢复 (#291)

Co-authored-by: hongyu9 <hongyu9@staff.weibo.com>

Author: Micro66

backend/app/api/endpoints/adapter/chat.py

@@ -12,6 +12,7 @@
 from app.api.dependencies import get_db
 from app.core import security
 from app.core.config import settings
+from app.models.subtask import Subtask, SubtaskRole, SubtaskStatus
 from app.models.user import User
 from app.schemas.task import (
     TaskCreate,
@@ -44,6 +45,22 @@ async def call_executor_cancel(task_id: int):
         )
 
 
+async def call_chat_shell_cancel(subtask_id: int, partial_content: str = ""):
+    """Background task to cancel Chat Shell streaming via session manager"""
+    try:
+        from app.services.chat.session_manager import session_manager
+        
+        success = await session_manager.cancel_stream(subtask_id)
+        if success:
+            logger.info(f"Chat Shell stream cancelled successfully for subtask {subtask_id}")
+        else:
+            logger.warning(f"Failed to cancel Chat Shell stream for subtask {subtask_id}")
+    except Exception as e:
+        logger.error(
+            f"Error cancelling Chat Shell stream for subtask {subtask_id}: {str(e)}"
+        )
+
+
 @router.post("", response_model=dict)
 def create_task_id(
     current_user: User = Depends(security.get_current_user),
@@ -170,7 +187,10 @@ async def cancel_task(
     current_user: User = Depends(security.get_current_user),
     db: Session = Depends(get_db),
 ):
-    """Cancel a running task by calling executor_manager"""
+    """Cancel a running task by calling executor_manager or Chat Shell cancel"""
+    from app.models.kind import Kind
+    from app.schemas.kind import Task
+    
     # Verify user owns this task
     task = task_kinds_service.get_task_detail(
         db=db, task_id=task_id, user_id=current_user.id
@@ -196,24 +216,101 @@ async def cancel_task(
         logger.info(f"Task {task_id} is already being cancelled")
         return {"message": "Task is already being cancelled", "status": "CANCELLING"}
 
-    # Update task status to CANCELLING immediately
-    try:
-        task_kinds_service.update_task(
-            db=db,
-            task_id=task_id,
-            obj_in=TaskUpdate(status="CANCELLING"),
-            user_id=current_user.id,
-        )
-        logger.info(
-            f"Task {task_id} status updated to CANCELLING by user {current_user.id}"
+    # Check if this is a Chat Shell task by looking at the source label
+    is_chat_shell = False
+    task_kind = (
+        db.query(Kind)
+        .filter(
+            Kind.id == task_id,
+            Kind.user_id == current_user.id,
+            Kind.kind == "Task",
+            Kind.is_active == True,
         )
-    except Exception as e:
-        logger.error(f"Failed to update task {task_id} status to CANCELLING: {str(e)}")
-        raise HTTPException(
-            status_code=500, detail=f"Failed to update task status: {str(e)}"
+        .first()
+    )
+    
+    if task_kind and task_kind.json:
+        task_crd = Task.model_validate(task_kind.json)
+        if task_crd.metadata.labels:
+            source = task_crd.metadata.labels.get("source", "")
+            is_chat_shell = source == "chat_shell"
+    
+    logger.info(f"Task {task_id} is_chat_shell={is_chat_shell}")
+
+    if is_chat_shell:
+        # For Chat Shell tasks, find the running subtask and cancel via session manager
+        running_subtask = (
+            db.query(Subtask)
+            .filter(
+                Subtask.task_id == task_id,
+                Subtask.user_id == current_user.id,
+                Subtask.role == SubtaskRole.ASSISTANT,
+                Subtask.status == SubtaskStatus.RUNNING,
+            )
+            .first()
         )
+        
+        if running_subtask:
+            # Cancel the Chat Shell stream
+            background_tasks.add_task(call_chat_shell_cancel, running_subtask.id)
+            
+            # Update subtask status to COMPLETED (not CANCELLED, to show partial content)
+            from datetime import datetime
+            running_subtask.status = SubtaskStatus.COMPLETED
+            running_subtask.progress = 100
+            running_subtask.completed_at = datetime.now()
+            running_subtask.updated_at = datetime.now()
+            running_subtask.error_message = ""
+            db.commit()
+            
+            # Update task status to COMPLETED (not CANCELLING, for Chat Shell)
+            try:
+                task_kinds_service.update_task(
+                    db=db,
+                    task_id=task_id,
+                    obj_in=TaskUpdate(status="COMPLETED"),
+                    user_id=current_user.id,
+                )
+                logger.info(
+                    f"Chat Shell task {task_id} cancelled and marked as COMPLETED"
+                )
+            except Exception as e:
+                logger.error(f"Failed to update Chat Shell task {task_id} status: {str(e)}")
+            
+            return {"message": "Chat stopped successfully", "status": "COMPLETED"}
+        else:
+            # No running subtask found, just mark task as completed
+            try:
+                task_kinds_service.update_task(
+                    db=db,
+                    task_id=task_id,
+                    obj_in=TaskUpdate(status="COMPLETED"),
+                    user_id=current_user.id,
+                )
+            except Exception as e:
+                logger.error(f"Failed to update task {task_id} status: {str(e)}")
+            
+            return {"message": "No running stream to cancel", "status": "COMPLETED"}
+    else:
+        # For non-Chat Shell tasks, use executor_manager
+        # Update task status to CANCELLING immediately
+        try:
+            task_kinds_service.update_task(
+                db=db,
+                task_id=task_id,
+                obj_in=TaskUpdate(status="CANCELLING"),
+                user_id=current_user.id,
+            )
+            logger.info(
+                f"Task {task_id} status updated to CANCELLING by user {current_user.id}"
+            )
+        except Exception as e:
+            logger.error(f"Failed to update task {task_id} status to CANCELLING: {str(e)}")
+            raise HTTPException(
+                status_code=500, detail=f"Failed to update task status: {str(e)}"
+            )
 
-    # Call executor_manager in the background
-    background_tasks.add_task(call_executor_cancel, task_id)
+        # Call executor_manager in the background
+        background_tasks.add_task(call_executor_cancel, task_id)
 
-    return {"message": "Cancel request accepted", "status": "CANCELLING"}
+        return {"message": "Cancel request accepted", "status": "CANCELLING"}

backend/app/api/endpoints/adapter/tasks.py

@@ -52,6 +52,12 @@ class Settings(BaseSettings):
     CHAT_HISTORY_MAX_MESSAGES: int = 50  # Maximum messages to keep in history
     CHAT_API_TIMEOUT_SECONDS: int = 300  # LLM API call timeout (5 minutes)
 
+    # Streaming incremental save configuration
+    STREAMING_REDIS_SAVE_INTERVAL: float = 1.0  # Redis save interval (seconds)
+    STREAMING_DB_SAVE_INTERVAL: float = 5.0  # Database save interval (seconds)
+    STREAMING_REDIS_TTL: int = 300  # Redis streaming cache TTL (seconds)
+    STREAMING_MIN_CHARS_TO_SAVE: int = 50  # Minimum characters to save on disconnect
+
     # Task append expiration (hours)
     APPEND_CHAT_TASK_EXPIRE_HOURS: int = 2
     APPEND_CODE_TASK_EXPIRE_HOURS: int = 24

backend/app/core/config.py

@@ -10,14 +10,21 @@
 from datetime import datetime
 from enum import Enum as PyEnum
 
-from sqlalchemy import Column, DateTime, ForeignKey, Integer, String, Text
+from sqlalchemy import Column, DateTime, ForeignKey, Integer, LargeBinary, String, Text
 from sqlalchemy import Enum as SQLEnum
 from sqlalchemy.dialects.mysql import LONGBLOB, LONGTEXT
 from sqlalchemy.sql import func
 
 from app.db.base import Base
 
 
+# Type adapter for binary data - uses LONGBLOB for MySQL, LargeBinary for others
+BinaryDataType = LargeBinary().with_variant(LONGBLOB, "mysql")
+
+# Type adapter for long text - uses LONGTEXT for MySQL, Text for others
+LongTextType = Text().with_variant(LONGTEXT, "mysql")
+
+
 class AttachmentStatus(str, PyEnum):
     """Attachment processing status."""
     UPLOADING = "uploading"
@@ -56,18 +63,18 @@ class SubtaskAttachment(Base):
     file_size = Column(Integer, nullable=False)  # File size in bytes
     mime_type = Column(String(100), nullable=False)
 
-    # Binary data storage (LONGBLOB for MySQL - supports up to 4GB)
-    binary_data = Column(LONGBLOB, nullable=False)
+    # Binary data storage (LONGBLOB for MySQL, LargeBinary for SQLite - supports up to 4GB)
+    binary_data = Column(BinaryDataType, nullable=False)
 
-    # Image base64 encoding (for vision models, LONGTEXT for large images)
+    # Image base64 encoding (for vision models, LONGTEXT for MySQL, Text for SQLite)
     # Note: MySQL doesn't allow default values for TEXT/BLOB columns, so nullable=True
     # Empty string or None means no image data
-    image_base64 = Column(LONGTEXT, nullable=True, default="")
+    image_base64 = Column(LongTextType, nullable=True, default="")
 
-    # Extracted text content (LONGTEXT for MySQL - supports up to 4GB)
+    # Extracted text content (LONGTEXT for MySQL, Text for SQLite - supports up to 4GB)
     # Note: MySQL doesn't allow default values for TEXT/BLOB columns, so nullable=True
     # Empty string or None means no extracted text
-    extracted_text = Column(LONGTEXT, nullable=True, default="")
+    extracted_text = Column(LongTextType, nullable=True, default="")
     text_length = Column(Integer, nullable=False, default=0)  # Character count of extracted text
 
     # Processing status