Merge branch 'rebased-chore-ai-search' of https://github.com/infinilabs/coco-server into rebased-chore-ai-search

Author: yaojiping

modules/assistant/api/session.go

@@ -411,9 +411,10 @@ func (h APIHandler) getChatHistoryBySession(w http.ResponseWriter, req *http.Req
 func (h APIHandler) cancelReplyMessage(w http.ResponseWriter, req *http.Request, ps httprouter.Params) {
 	sessionID := ps.MustGetParameter("session_id")
 	messageID := h.GetParameterOrDefault(req, "message_id", "")
+	lang := h.GetParameterOrDefault(req, "lang", "")
 	log.Info("cancel reply to message: ", messageID, ", session: ", sessionID)
 	taskID := service.GetReplyMessageTaskID(sessionID, messageID)
-	service.StopMessageReplyTask(taskID)
+	service.StopMessageReplyTask(taskID, lang)
 	h.WriteAckOKJSON(w)
 }
 

modules/assistant/common/message_task.go

@@ -3,4 +3,10 @@ package common
 type MessageTask struct {
 	SessionID  string
 	CancelFunc func()
+	// CancelLang is the user's UI language (e.g. "zh-CN", "en"), set by the
+	// cancel API right before invoking CancelFunc. The async processor reads
+	// this from the InflightMessages map so it can persist a localized
+	// "task cancelled" message into the reply — the backend has no i18n
+	// system, so the frontend supplies the language at cancel time.
+	CancelLang string
 }

modules/assistant/service/background_job.go

@@ -28,6 +28,19 @@ import (
 	"infini.sh/framework/core/util"
 )
 
+// getCancelledMessage returns a localized "task cancelled" message based on
+// the CancelLang stored in the inflight task. The cancel API writes the user's
+// UI language into the task right before triggering cancellation, so we can
+// read it here and produce an appropriate message without a full i18n system.
+func getCancelledMessage(taskID string) string {
+	if v, ok := InflightMessages.Load(taskID); ok {
+		if task, ok := v.(common2.MessageTask); ok && strings.HasPrefix(task.CancelLang, "zh") {
+			return "~~该任务已被取消。~~"
+		}
+	}
+	return "~~This task has been cancelled.~~"
+}
+
 func CreateAssistantReplyMessage(sessionID, assistantID, requestMessageID string) *core.ChatMessage {
 	msg := &core.ChatMessage{
 		SessionID:      sessionID,
@@ -71,6 +84,8 @@ func ProcessMessageAsync(ctx context.Context, userID string, reqMsg, replyMsg *c
 	)
 
 	defer func() {
+		taskID := GetReplyMessageTaskID(params.SessionID, reqMsg.ID)
+
 		if !global.Env().IsDebug {
 			if r := recover(); r != nil {
 				var v string
@@ -85,8 +100,11 @@ func ProcessMessageAsync(ctx context.Context, userID string, reqMsg, replyMsg *c
 
 				// If the context was cancelled (user switched chat or clicked cancel),
 				// treat it as a graceful stop — no error message to the user.
-				if ctx.Err() != nil {
+				if ctx.Err() != nil || strings.Contains(v, "context canceled") || strings.Contains(v, "context deadline exceeded") {
 					log.Infof("async processing cancelled by user: %v", v)
+					if replyMsg.Message == "" {
+						replyMsg.Message = getCancelledMessage(taskID)
+					}
 				} else {
 					msg := fmt.Sprintf("⚠️ error in async processing message reply, %v", v)
 					if replyMsg.Message != "" {
@@ -102,13 +120,19 @@ func ProcessMessageAsync(ctx context.Context, userID string, reqMsg, replyMsg *c
 		}
 
 		if err != nil {
-			log.Errorf("Failed to process message reply: %v", err)
-			replyMsg.Message += err.Error()
+			if ctx.Err() != nil || strings.Contains(err.Error(), "context canceled") || strings.Contains(err.Error(), "context deadline exceeded") {
+				log.Infof("async processing cancelled: %v", err)
+				if replyMsg.Message == "" {
+					replyMsg.Message = getCancelledMessage(taskID)
+				}
+			} else {
+				log.Errorf("Failed to process message reply: %v", err)
+				replyMsg.Message += err.Error()
+			}
 		}
 
 		finalizeProcessing(ctx, params.SessionID, replyMsg, sender)
 		// clear the inflight message task
-		taskID := GetReplyMessageTaskID(params.SessionID, reqMsg.ID)
 		InflightMessages.Delete(taskID)
 
 		log.Info("finished async processing message")

modules/assistant/service/sessioin.go

@@ -5,6 +5,7 @@ import (
 	"errors"
 	"fmt"
 	"net/http"
+	"strings"
 	"sync"
 
 	"github.com/cihub/seelog"
@@ -91,17 +92,48 @@ func GetChatHistoryBySessionInternal(sessionID string, size int) ([]core.ChatMes
 	return docs, nil
 }
 
-func StopMessageReplyTask(taskID string) {
+// StopMessageReplyTask cancels the in-flight async reply identified by taskID.
+// lang is the caller's UI language (e.g. "zh-CN"); it is written into the task
+// before calling CancelFunc so that the async processor (background_job.go) can
+// read it back and persist a localized cancellation message. This is necessary
+// because the backend has no i18n framework — the frontend is the only place
+// that knows the user's preferred language.
+func StopMessageReplyTask(taskID string, lang string) {
 	v, ok := InflightMessages.Load(taskID)
 	if ok {
 		v1, ok := v.(common.MessageTask)
 		if ok {
 			seelog.Debug("stop task:", v1)
+			// Store lang back before cancelling so the deferred cleanup
+			// in background_job.go can read it from InflightMessages.
+			v1.CancelLang = lang
+			InflightMessages.Store(taskID, v1)
 			v1.CancelFunc()
 		}
-	} else {
-		_ = seelog.Warnf("task id [%s] was not found", taskID)
+		return
+	}
+
+	// If the exact taskID was not found and it looks like a bare sessionID
+	// (no underscore), try to find any inflight task for this session.
+	if !strings.Contains(taskID, "_") {
+		InflightMessages.Range(func(key, value any) bool {
+			k, _ := key.(string)
+			if strings.HasPrefix(k, taskID+"_") {
+				task, ok := value.(common.MessageTask)
+				if ok && task.CancelFunc != nil {
+					seelog.Debugf("stop task by session prefix: %s", k)
+					task.CancelLang = lang
+					InflightMessages.Store(k, task)
+					task.CancelFunc()
+				}
+				return false // stop iteration after first match
+			}
+			return true
+		})
+		return
 	}
+
+	_ = seelog.Warnf("task id [%s] was not found", taskID)
 }
 
 func StopAllMessageReplyTasks() int {