Merge pull request #43 from AnimeAIChat/develop

重构任务管理，连接优化

Author: Eric0308

go.mod

@@ -42,8 +42,8 @@ require (
 	golang.org/x/arch v0.17.0 // indirect
 	golang.org/x/crypto v0.38.0 // indirect
 	golang.org/x/net v0.40.0 // indirect
-	golang.org/x/sync v0.14.0 // indirect
 	golang.org/x/sys v0.33.0 // indirect
 	golang.org/x/text v0.25.0 // indirect
 	google.golang.org/protobuf v1.36.6 // indirect
+	gopkg.in/check.v1 v1.0.0-20201130134442-10cb98267c6c // indirect
 )

go.sum

@@ -53,8 +53,11 @@ github.com/klauspost/cpuid/v2 v2.0.9/go.mod h1:FInQzS24/EEf25PyTYn52gqo7WaD8xa02
 github.com/klauspost/cpuid/v2 v2.2.10 h1:tBs3QSyvjDyFTq3uoc/9xFpCuOsJQFNPiAhYdw2skhE=
 github.com/klauspost/cpuid/v2 v2.2.10/go.mod h1:hqwkgyIinND0mEev00jJYCxPNVRVXFQeu1XKlok6oO0=
 github.com/knz/go-libedit v1.10.1/go.mod h1:MZTVkCWyz0oBc7JOWP3wNAzd002ZbM/5hgShxwh4x8M=
+github.com/kr/pretty v0.2.1/go.mod h1:ipq/a2n7PKx3OHsz4KJII5eveXtPO4qwEXGdVfWzfnI=
 github.com/kr/pretty v0.3.1 h1:flRD4NNwYAUpkphVc1HcthR4KEIFJ65n8Mw5qdRn3LE=
 github.com/kr/pretty v0.3.1/go.mod h1:hoEshYVHaxMs3cyo3Yncou5ZscifuDolrwPKZanG3xk=
+github.com/kr/pty v1.1.1/go.mod h1:pFQYn66WHrOpPYNljwOMqo10TkYh1fy3cYio2l3bCsQ=
+github.com/kr/text v0.1.0/go.mod h1:4Jbv+DJW3UT/LiOwJeYQe1efqtUx/iVham/4vfdArNI=
 github.com/kr/text v0.2.0 h1:5Nx0Ya0ZqY2ygV366QzturHI13Jq95ApcVaJBhpS+AY=
 github.com/kr/text v0.2.0/go.mod h1:eLer722TekiGuMkidMxC/pM04lWEeraHUUmBw8l2grE=
 github.com/leodido/go-urn v1.4.0 h1:WT9HwE9SGECu3lg4d/dIA+jxlljEa1/ffXKmRjqdmIQ=
@@ -124,8 +127,9 @@ golang.org/x/tools v0.21.1-0.20240508182429-e35e4ccd0d2d h1:vU5i/LfpvrRCpgM/VPfJ
 golang.org/x/tools v0.21.1-0.20240508182429-e35e4ccd0d2d/go.mod h1:aiJjzUbINMkxbQROHiO6hDPo2LHcIPhhQsa9DLh0yGk=
 google.golang.org/protobuf v1.36.6 h1:z1NpPI8ku2WgiWnf+t9wTPsn6eP1L7ksHUlkfLvd9xY=
 google.golang.org/protobuf v1.36.6/go.mod h1:jduwjTPXsFjZGTmRluh+L6NjiWu7pchiJ2/5YcXBHnY=
-gopkg.in/check.v1 v0.0.0-20161208181325-20d25e280405 h1:yhCVgyC4o1eVCa2tZl7eS0r+SDo693bJlVdllGtEeKM=
 gopkg.in/check.v1 v0.0.0-20161208181325-20d25e280405/go.mod h1:Co6ibVJAznAaIkqp8huTwlJQCZ016jof/cbN4VW5Yz0=
+gopkg.in/check.v1 v1.0.0-20201130134442-10cb98267c6c h1:Hei/4ADfdWqJk1ZMxUNpqntNwaWcugrBjAiHlqqRiVk=
+gopkg.in/check.v1 v1.0.0-20201130134442-10cb98267c6c/go.mod h1:JHkPIbrfpd72SG/EVd6muEfDQjcINNoR0C8j2r3qZ4Q=
 gopkg.in/yaml.v3 v3.0.0-20200313102051-9f266ea9e77c/go.mod h1:K4uyk7z7BCEPqu6E+C64Yfv1cQ7kz7rIZviUmN+EgEM=
 gopkg.in/yaml.v3 v3.0.1 h1:fxVm/GzAzEWqLHuvctI91KS9hhNmmWOoWu0XTYJS7CA=
 gopkg.in/yaml.v3 v3.0.1/go.mod h1:K4uyk7z7BCEPqu6E+C64Yfv1cQ7kz7rIZviUmN+EgEM=

src/core/Auth/token.go src/core/auth/token.gosrc/core/Auth/token.go renamed to src/core/auth/token.go

@@ -1,4 +1,4 @@
-package Auth
+package auth
 
 import (
 	"errors"

src/core/connection.go

@@ -36,13 +36,14 @@ type configGetter interface {
 // ConnectionHandler 连接处理器结构
 type ConnectionHandler struct {
 	// 确保实现 AsrEventListener 接口
-	_         providers.AsrEventListener
-	config    *configs.Config
-	logger    *utils.Logger
-	conn      Conn
-	closeOnce sync.Once
-	taskMgr   *task.TaskManager
-	providers struct {
+	_                providers.AsrEventListener
+	config           *configs.Config
+	logger           *utils.Logger
+	conn             Conn
+	closeOnce        sync.Once
+	taskMgr          *task.TaskManager
+	safeCallbackFunc func(func(*ConnectionHandler)) func()
+	providers        struct {
 		asr   providers.ASRProvider
 		llm   providers.LLMProvider
 		tts   providers.TTSProvider
@@ -106,6 +107,8 @@ type ConnectionHandler struct {
 	// functions
 	functionRegister *function.FunctionRegistry
 	mcpManager       *mcp.Manager
+
+	ctx context.Context
 }
 
 // NewConnectionHandler 创建新的连接处理器
@@ -114,6 +117,7 @@ func NewConnectionHandler(
 	providerSet *pool.ProviderSet,
 	logger *utils.Logger,
 	req *http.Request,
+	ctx context.Context,
 ) *ConnectionHandler {
 	handler := &ConnectionHandler{
 		config:           config,
@@ -143,6 +147,8 @@ func NewConnectionHandler(
 		serverAudioChannels:      1,
 		serverAudioFrameDuration: 60,
 
+		ctx: ctx,
+
 		headers: make(map[string]string),
 	}
 
@@ -187,6 +193,37 @@ func NewConnectionHandler(
 	return handler
 }
 
+func (h *ConnectionHandler) SetTaskCallback(callback func(func(*ConnectionHandler)) func()) {
+	h.safeCallbackFunc = callback
+}
+
+func (h *ConnectionHandler) SubmitTask(taskType string, params map[string]interface{}) {
+	_task, id := task.NewTask(h.ctx, "", params)
+	h.logger.Info(fmt.Sprintf("提交任务: %s, ID: %s, 参数: %v", _task.Type, id, params))
+	// 创建安全回调用于任务完成时调用
+	var taskCallback func(result interface{})
+	if h.safeCallbackFunc != nil {
+		taskCallback = func(result interface{}) {
+			fmt.Print("任务完成回调: ")
+			safeCallback := h.safeCallbackFunc(func(handler *ConnectionHandler) {
+				// 处理任务完成逻辑
+				handler.handleTaskComplete(_task, id, result)
+			})
+			// 执行安全回调
+			if safeCallback != nil {
+				safeCallback()
+			}
+		}
+	}
+	cb := task.NewCallBack(taskCallback)
+	_task.Callback = cb
+	h.taskMgr.SubmitTask(h.sessionID, _task)
+}
+
+func (h *ConnectionHandler) handleTaskComplete(task *task.Task, id string, result interface{}) {
+	h.logger.Info(fmt.Sprintf("任务 %s 完成，ID: %s, %v", task.Type, id, result))
+}
+
 // Handle 处理WebSocket连接
 func (h *ConnectionHandler) Handle(conn Conn) {
 	defer conn.Close()

src/core/connection_handlemsg.go

@@ -10,7 +10,6 @@ import (
 	"xiaozhi-server-go/src/core/image"
 	"xiaozhi-server-go/src/core/providers"
 	"xiaozhi-server-go/src/core/utils"
-	"xiaozhi-server-go/src/task"
 )
 
 // handleMessage 处理接收到的消息
@@ -102,17 +101,6 @@ func (h *ConnectionHandler) handleVisionMessage(msgMap map[string]interface{}) e
 	// 处理视觉消息
 	cmd := msgMap["cmd"].(string)
 	if cmd == "gen_pic" {
-		text := msgMap["text"].(string)
-		params := map[string]interface{}{
-			"prompt":    text,
-			"size":      "1024x1024",
-			"quality":   "standard",
-			"api_key":   h.config.LLM["ChatGLMLLM"].APIKey,
-			"client_id": h.sessionID,
-		}
-		task, id := task.NewTask(task.TaskTypeImageGen, params, task.NewMessageCallback(h.conn, "vision", cmd))
-		h.taskMgr.SubmitTask(h.sessionID, task)
-		h.logger.Info(fmt.Sprintf("生成图片任务提交成功: %s, %s", text, id))
 	} else if cmd == "gen_video" {
 	} else if cmd == "read_img" {
 	}

src/core/mcp/manager.go

@@ -144,6 +144,8 @@ func (m *Manager) BindConnection(conn Conn, fh types.FunctionRegistryInterface,
 	} else {
 		// 重新绑定连接而不是重新创建
 		m.XiaoZhiMCPClient.SetConnection(conn)
+		m.XiaoZhiMCPClient.SetID(deviceID, clientID)
+		m.XiaoZhiMCPClient.SetToken(token)
 		if !m.XiaoZhiMCPClient.IsReady() {
 			if err := m.XiaoZhiMCPClient.Start(context.Background()); err != nil {
 				return fmt.Errorf("重启XiaoZhi MCP客户端失败: %v", err)

src/core/mcp/xiaozhi_client.go

@@ -8,7 +8,7 @@ import (
 	"sync"
 	"time"
 
-	"xiaozhi-server-go/src/core/Auth"
+	"xiaozhi-server-go/src/core/auth"
 	"xiaozhi-server-go/src/core/utils"
 
 	"github.com/sashabaranov/go-openai"
@@ -75,7 +75,7 @@ func (c *XiaoZhiMCPClient) SetID(deviceID string, clientID string) {
 }
 
 func (c *XiaoZhiMCPClient) SetToken(token string) {
-	auth := Auth.NewAuthToken(token)
+	auth := auth.NewAuthToken(token)
 	visionToken, err := auth.GenerateToken(c.deviceID)
 
 	if err != nil {

src/core/websocket_conn.go

@@ -0,0 +1,110 @@
+package core
+
+import (
+	"errors"
+	"sync"
+	"sync/atomic"
+	"time"
+
+	"github.com/gorilla/websocket"
+)
+
+var (
+	ErrConnectionClosed = errors.New("websocket connection is closed")
+)
+
+// websocketConn 封装gorilla/websocket的连接实现
+type websocketConn struct {
+	conn       *websocket.Conn
+	writeMu    sync.Mutex // 写操作互斥锁
+	closed     int32      // 原子操作标记连接状态 (0=open, 1=closed)
+	lastActive int64      // 最后活跃时间戳（原子操作）
+}
+
+func (w *websocketConn) ReadMessage() (messageType int, p []byte, err error) {
+	if atomic.LoadInt32(&w.closed) == 1 {
+		return 0, nil, ErrConnectionClosed
+	}
+
+	// 设置读取超时
+	w.conn.SetReadDeadline(time.Now().Add(60 * time.Second))
+
+	messageType, p, err = w.conn.ReadMessage()
+	if err != nil {
+		// 如果读取出错，标记连接为已关闭
+		atomic.StoreInt32(&w.closed, 1)
+		return 0, nil, err
+	}
+
+	// 更新最后活跃时间
+	atomic.StoreInt64(&w.lastActive, time.Now().Unix())
+
+	return messageType, p, nil
+}
+
+func (w *websocketConn) WriteMessage(messageType int, data []byte) error {
+	// 检查连接状态
+	if atomic.LoadInt32(&w.closed) == 1 {
+		return ErrConnectionClosed
+	}
+
+	// 使用写锁确保写操作的串行化
+	w.writeMu.Lock()
+	defer w.writeMu.Unlock()
+
+	// 双重检查，防止在获取锁的过程中连接被关闭
+	if atomic.LoadInt32(&w.closed) == 1 {
+		return ErrConnectionClosed
+	}
+
+	// 设置写入超时
+	w.conn.SetWriteDeadline(time.Now().Add(30 * time.Second))
+
+	err := w.conn.WriteMessage(messageType, data)
+	if err != nil {
+		// 如果写入出错，标记连接为已关闭
+		atomic.StoreInt32(&w.closed, 1)
+		return err
+	}
+
+	// 更新最后活跃时间
+	atomic.StoreInt64(&w.lastActive, time.Now().Unix())
+
+	return nil
+}
+
+func (w *websocketConn) Close() error {
+	// 使用原子操作避免重复关闭
+	if !atomic.CompareAndSwapInt32(&w.closed, 0, 1) {
+		return nil // 已经关闭过了
+	}
+
+	w.writeMu.Lock()
+	defer w.writeMu.Unlock()
+
+	// 尝试发送关闭帧（不强制要求成功）
+	closeMsg := websocket.FormatCloseMessage(websocket.CloseNormalClosure, "connection closed")
+	w.conn.SetWriteDeadline(time.Now().Add(5 * time.Second))
+	w.conn.WriteMessage(websocket.CloseMessage, closeMsg)
+
+	return w.conn.Close()
+}
+
+// IsClosed 检查连接是否已关闭
+func (w *websocketConn) IsClosed() bool {
+	return atomic.LoadInt32(&w.closed) == 1
+}
+
+// GetLastActiveTime 获取最后活跃时间
+func (w *websocketConn) GetLastActiveTime() time.Time {
+	timestamp := atomic.LoadInt64(&w.lastActive)
+	return time.Unix(timestamp, 0)
+}
+
+// IsStale 检查连接是否过期（基于最后活跃时间）
+func (w *websocketConn) IsStale(timeout time.Duration) bool {
+	if w.IsClosed() {
+		return true
+	}
+	return time.Since(w.GetLastActiveTime()) > timeout
+}