fix: ongoing chat stop issue

backend/open_webui/main.py

@@ -372,7 +372,11 @@ from open_webui.utils.auth import (
 from open_webui.utils.oauth import OAuthManager
 from open_webui.utils.security_headers import SecurityHeadersMiddleware
 
-from open_webui.tasks import stop_task, list_tasks  # Import from tasks.py
+from open_webui.tasks import (
+    list_task_ids_by_chat_id,
+    stop_task,
+    list_tasks,
+)  # Import from tasks.py
 
 from open_webui.utils.redis import get_sentinels_from_env
 
@@ -1196,7 +1200,7 @@ async def chat_action(
 @app.post("/api/tasks/stop/{task_id}")
 async def stop_task_endpoint(task_id: str, user=Depends(get_verified_user)):
     try:
-        result = await stop_task(task_id)  # Use the function from tasks.py
+        result = await stop_task(task_id)
         return result
     except ValueError as e:
         raise HTTPException(status_code=status.HTTP_404_NOT_FOUND, detail=str(e))
@@ -1204,7 +1208,19 @@ async def stop_task_endpoint(task_id: str, user=Depends(get_verified_user)):
 
 @app.get("/api/tasks")
 async def list_tasks_endpoint(user=Depends(get_verified_user)):
-    return {"tasks": list_tasks()}  # Use the function from tasks.py
+    return {"tasks": list_tasks()}
+
+
+@app.get("/api/tasks/chat/{chat_id}")
+async def list_tasks_by_chat_id_endpoint(chat_id: str, user=Depends(get_verified_user)):
+    chat = Chats.get_chat_by_id(chat_id)
+    if chat is None or chat.user_id != user.id:
+        return {"task_ids": []}
+
+    task_ids = list_task_ids_by_chat_id(chat_id)
+
+    print(f"Task IDs for chat {chat_id}: {task_ids}")
+    return {"task_ids": task_ids}
 
 
 ##################################

backend/open_webui/tasks.py

@@ -5,16 +5,23 @@ from uuid import uuid4
 
 # A dictionary to keep track of active tasks
 tasks: Dict[str, asyncio.Task] = {}
+chat_tasks = {}
 
 
-def cleanup_task(task_id: str):
+def cleanup_task(task_id: str, id=None):
     """
     Remove a completed or canceled task from the global `tasks` dictionary.
     """
     tasks.pop(task_id, None)  # Remove the task if it exists
 
+    # If an ID is provided, remove the task from the chat_tasks dictionary
+    if id and task_id in chat_tasks.get(id, []):
+        chat_tasks[id].remove(task_id)
+        if not chat_tasks[id]:  # If no tasks left for this ID, remove the entry
+            chat_tasks.pop(id, None)
 
-def create_task(coroutine):
+
+def create_task(coroutine, id=None):
     """
     Create a new asyncio task and add it to the global task dictionary.
     """
@@ -22,9 +29,15 @@ def create_task(coroutine):
     task = asyncio.create_task(coroutine)  # Create the task
 
     # Add a done callback for cleanup
-    task.add_done_callback(lambda t: cleanup_task(task_id))
-
+    task.add_done_callback(lambda t: cleanup_task(task_id, id))
     tasks[task_id] = task
+
+    # If an ID is provided, associate the task with that ID
+    if chat_tasks.get(id):
+        chat_tasks[id].append(task_id)
+    else:
+        chat_tasks[id] = [task_id]
+
     return task_id, task
 
 
@@ -42,6 +55,13 @@ def list_tasks():
     return list(tasks.keys())
 
 
+def list_task_ids_by_chat_id(id):
+    """
+    List all tasks associated with a specific ID.
+    """
+    return chat_tasks.get(id, [])
+
+
 async def stop_task(task_id: str):
     """
     Cancel a running task and remove it from the global task list.

backend/open_webui/utils/middleware.py

@@ -2245,7 +2245,9 @@ async def process_chat_response(
                 await response.background()
 
         # background_tasks.add_task(post_response_handler, response, events)
-        task_id, _ = create_task(post_response_handler(response, events))
+        task_id, _ = create_task(
+            post_response_handler(response, events), id=metadata["chat_id"]
+        )
         return {"status": True, "task_id": task_id}
 
     else: