src/agent/agent.go

@@ -7,7 +7,6 @@ import (
 	"agent/src/workflow"
 	"context"
 	"fmt"
-	"io"
 	"net/http"
 	"sync"
 )
@@ -15,11 +14,10 @@ import (
 const AgentWorkers = 4
 
 type Agent struct {
-	ctx      context.Context
-	apicheck *keyChecker
-	clients  *common.ThreadSafeMap[*client, struct{}]
+	ctx        context.Context
+	apicheck   *keyChecker
+	clients    *common.ThreadSafeMap[*client, struct{}]
 }
-
 var once sync.Once
 
 func NewAgent(ctx context.Context, ollamaUrls []string) (*Agent, error) {
@@ -31,20 +29,23 @@ func NewAgent(ctx context.Context, ollamaUrls []string) (*Agent, error) {
 		return nil, err
 	}
 	return &Agent{
-		ctx:     ctx,
+		ctx    : ctx,
 		clients: new(common.ThreadSafeMap[*client, struct{}]).Init(nil, false),
 	}, nil
 }
 
 func (a *Agent) Start(port uint64, crtDir *string) {
 	workflow.InitTaskPool(a.ctx)
-
-	go func() {
+	
+	go func(){
 		sm := http.NewServeMux()
-		sm.HandleFunc("/sse", a.serveSSE)
+		// sm.HandleFunc("/simpletest2025", a.serveTestPage)
+		// sm.HandleFunc("/doc", a.serveDoc)
+		// sm.HandleFunc("/assets/", a.serveAssets)
+		sm.HandleFunc("/ws", a.serveWs)
 
 		addr := fmt.Sprintf(":%d", port)
-		log.Info("[agent] start SSE server", "addr", addr)
+		log.Info("[agent] start websocket server", "addr", addr)
 
 		var err error
 		if crtDir == nil {
@@ -59,41 +60,6 @@ func (a *Agent) Start(port uint64, crtDir *string) {
 	}()
 }
 
-func (a *Agent) serveSSE(w http.ResponseWriter, r *http.Request) {
-	flusher, ok := w.(http.Flusher)
-	if !ok {
-		http.Error(w, "Streaming unsupported", http.StatusInternalServerError)
-		return
-	}
-
-	// 创建新 client
-	client := newClient(w, flusher, a)
-	a.addClient(client)
-
-	// 设置 SSE 必备的响应头
-	w.Header().Set("Content-Type", "text/event-stream")
-	w.Header().Set("Cache-Control", "no-cache")
-	w.Header().Set("Connection", "keep-alive")
-	w.Header().Set("Access-Control-Allow-Origin", "*")
-
-	// 启动写消息循环
-	client.run()
-
-	// 读取请求体
-	msg, err := io.ReadAll(r.Body)
-	if err != nil {
-		client.SendText("Invalid request body")
-		return
-	}
-	defer r.Body.Close()
-
-	// 处理请求
-	client.processSSERequest(msg)
-
-	// 阻塞直到 client 关闭
-	<-client.GetCtx().Done()
-}
-
 func (a *Agent) serveTestPage(w http.ResponseWriter, r *http.Request) {
 	log.Info("serveHome", "url", r.URL)
 	if r.Method != http.MethodGet {
@@ -122,12 +88,30 @@ func (a *Agent) serveAssets(w http.ResponseWriter, r *http.Request) {
 	http.ServeFile(w, r, "./"+r.URL.Path[1:])
 }
 
-// 给Agent添加addClient方法
-func (a *Agent) addClient(c *client) {
-	a.clients.Set(c, struct{}{}) // 放进去
-	log.Info("client added", "connected", a.clients.Size())
+func (a *Agent) serveWs(w http.ResponseWriter, r *http.Request) {
+	// TODO IP连接数限制
+	conn, err := upgrader.Upgrade(w, r, nil)
+	if err != nil {
+		log.Error("[agent] serveWs", "err", err)
+		return
+	}
+	log.Info("[agent] new connection", "remote", conn.RemoteAddr())
+	c := newClient(conn, a)
+
+	// check api key
+	if a.apicheck.check(c, r) {
+		c.run()
+		a.clients.Set(c, struct{}{})
+	} else {
+		go func() {
+			c.SendText("HTTP/1.1 401 Unauthorized\r\n\r\n")
+			conn.Close()
+		}()
+	}
 }
 
 func (a *Agent) removeClient(c *client) {
 	a.clients.Delete(c)
 }
+
+

src/agent/client.go

@@ -1,183 +1,243 @@
 package agent
 
 import (
-	"agent/src/message"
-	"agent/src/workflow"
 	"context"
 	"encoding/json"
 	"errors"
-	"net/http"
+	"fmt"
 	"sync"
 	"time"
 
+	"agent/src/message"
 	"agent/src/utils/log"
+	"agent/src/workflow"
 
+	"github.com/gorilla/websocket"
 	"golang.org/x/time/rate"
 )
 
+var upgrader = websocket.Upgrader{
+	CheckOrigin: func(r *http.Request) bool {
+		return true // 允许所有来源（仅限开发环境）
+	},
+	ReadBufferSize:  1024,
+	WriteBufferSize: 1024,
+}
+
+
 const (
+	// Time allowed to write a message to the peer.
 	writeWait = 15 * time.Second
+
+	// Time allowed to read the next pong message from the peer.
+	pongWait = 3 * 60 * time.Second
+
+	// Send pings to peer with this period. Must be less than pongWait.
+	pingPeriod = (pongWait * 9) / 10
+
+	// Maximum message size allowed from peer.
+	maxMessageSize = 10*1024*1024
 )
 
-// client 用于管理每个与客户端的连接
 type client struct {
-	ctx            context.Context
-	writer         http.ResponseWriter
-	flusher        http.Flusher
-	outMsg         chan []byte
-	authenticated  bool
-	apikey         string
-	agent          *Agent
-	fnCancel       context.CancelFunc
-	rate           *rate.Limiter
-	wantResponseId uint64      // 新增字段：存储当前请求的ID
-	wantResponse   chan []byte // 新增字段：用于存储响应的通道
-
-	lock sync.Mutex
-	once sync.Once
+	ctx    context.Context
+	conn   *websocket.Conn
+	outMsg chan []byte
+	authenticated bool
+	apikey string
+	agent  *Agent
+	fnCancel context.CancelFunc
+	rate   *rate.Limiter
+	
+	wantResponseId uint64
+	wantResponse   chan []byte
+	lock           sync.Mutex
+	once           sync.Once
 }
 
-// 创建新的客户端对象
-func newClient(w http.ResponseWriter, f http.Flusher, a *Agent) *client {
+func newClient(conn *websocket.Conn, a *Agent) *client {
 	ctx, fnCancel := context.WithCancel(a.ctx)
 	return &client{
-		ctx:      ctx,
-		writer:   w,
-		flusher:  f,
-		outMsg:   make(chan []byte, 32),
-		agent:    a,
+		ctx   : ctx,
+		conn  : conn,
+		outMsg: make(chan []byte, 32),
+		agent : a,
 		fnCancel: fnCancel,
-		rate:     rate.NewLimiter(5, 1), // 每秒5次限制
+		rate: rate.NewLimiter(5, 1), // 每秒5次限制
 	}
 }
-
-// 设置API Key并标记为已认证
-func (c *client) setApiKey(apikey string) {
+func (c *client) setApiKey(apikey string) { 
 	c.apikey = apikey
 	c.authenticated = true
 }
 
-// 启动客户端的处理流程
 func (c *client) run() {
+	go c.readPump()
 	go c.writePump()
 }
 
-// writePump 处理SSE的消息推送
+func (c *client) readPump() {
+	defer c.Close()
+
+	c.conn.SetReadLimit(maxMessageSize)
+	c.conn.SetReadDeadline(time.Now().Add(pongWait))
+	c.conn.SetPongHandler(func(string) error { c.conn.SetReadDeadline(time.Now().Add(pongWait)); return nil })
+	done := c.ctx.Done()
+	for {
+		select {			
+		case <-done:
+			return
+		default:
+		}
+
+		_, msg, err := c.conn.ReadMessage()
+		if err != nil {
+			if websocket.IsUnexpectedCloseError(err, websocket.CloseGoingAway, websocket.CloseAbnormalClosure) {
+				log.Error(fmt.Sprintf("error: %v", err))
+			}
+			break
+		}
+
+		if c.rate.Allow() {
+			workflow.Handle(c, msg)
+		} else {
+			c.WriteJson(&message.ReponseError{
+				Error: "maximum 5 requests per second",
+			}, true)
+		}
+
+	}
+}
+
+
 func (c *client) writePump() {
+	ticker := time.NewTicker(pingPeriod)
+	defer ticker.Stop()
 	defer c.Close()
 
 	done := c.ctx.Done()
 	for {
-		select {
+		select {			
 		case message, ok := <-c.outMsg:
+			c.conn.SetWriteDeadline(time.Now().Add(writeWait))
 			if !ok {
+				c.conn.WriteMessage(websocket.CloseMessage, []byte{})
+				return
+			}
+
+			w, err := c.conn.NextWriter(websocket.TextMessage)
+			if err != nil {
+				return
+			}
+			w.Write(message)
+
+			if err := w.Close(); err != nil {
+				return
+			}
+		case <-ticker.C:
+			c.conn.SetWriteDeadline(time.Now().Add(writeWait))
+			if err := c.conn.WriteMessage(websocket.PingMessage, nil); err != nil {
 				return
 			}
-			// 向客户端发送消息，SSE格式
-			c.writer.Write([]byte("data: "))
-			c.writer.Write(message)
-			c.writer.Write([]byte("\n\n"))
-			c.flusher.Flush()
 		case <-done:
 			return
 		}
 	}
 }
 
-// 关闭客户端连接
 func (c *client) Close() {
-	c.once.Do(func() {
+	once.Do(func(){
 		c.fnCancel()
 		c.agent.removeClient(c)
-		log.Info("client closed", "connected", c.agent.clients.Size())
+		c.conn.Close()
+		log.Info("client close", "remote", c.conn.RemoteAddr().String(), "connected", c.agent.clients.Size())
 	})
 }
 
-// 将JSON数据发送到客户端
 func (c *client) WriteJson(data interface{}, block bool) error {
-	b, err := json.Marshal(data)
+	b, err := json.MarshalIndent(data, "", "")
 	if err != nil {
 		return err
 	}
-	return c.WriteText(string(b), block)
-}
-
-// 将文本消息发送到客户端
-func (c *client) WriteText(msg string, block bool) error {
 	if block {
-		select {
-		case <-c.ctx.Done():
-			return errors.New("connection closed")
-		case c.outMsg <- []byte(msg):
-			return nil
+		select{
+			case <-c.ctx.Done():
+				return errors.New("connection closed")
+			case c.outMsg <- b:
+				return nil
 		}
 	} else {
 		select {
-		case c.outMsg <- []byte(msg):
-		default:
+			case <-c.ctx.Done():
+				return errors.New("connection closed")
+			case c.outMsg <- b:
+				return nil
+			default:
+				go func(){
+					select{
+						case <-c.ctx.Done():
+						case c.outMsg <- b:
+					}
+				}()
 		}
 	}
 	return nil
 }
 
-// 发送纯文本消息
-func (c *client) SendText(msg string) error {
-	return c.WriteText(msg, true)
+func (c *client) WriteText(msg string, block bool) error {
+	if block {
+		select {
+			case <-c.ctx.Done():
+				return errors.New("connection closed")
+			case c.outMsg <- []byte(msg):
+				return nil
+		}
+	} else {
+		select {
+			case c.outMsg <- []byte(msg):
+			default:
+		}
+	}
+	return nil
+}
+
+// send string, block util flush
+func (c *client) SendText(msg string) error {
+	c.conn.SetWriteDeadline(time.Now().Add(writeWait))
+	w, err := c.conn.NextWriter(websocket.TextMessage)
+	if err != nil {
+		return err
+	}
+	_, err = w.Write([]byte(msg))
+	if err != nil {
+		return err
+	}
+	err = w.Close()
+	return err
 }
 
-// 获取客户端的上下文
 func (c *client) GetCtx() context.Context {
 	return c.ctx
 }
 
-// 获取客户端的消息通道
 func (c *client) GetChannel() chan<- []byte {
 	return c.outMsg
 }
 
-// 处理接收到的消息请求
-func (c *client) processSSERequest(msg []byte) {
-	// 检查速率限制
-	if !c.rate.Allow() {
-		err := c.WriteJson(&message.ReponseError{
-			Error: "maximum 5 requests per second",
-		}, true)
-		if err != nil {
-			log.Error("Failed to send rate limit response", "error", err)
-		}
-		return
-	}
 
-	// 调用workflow.Handle来处理请求
-	err := workflow.Handle(c, msg)
-	if err != nil {
-		log.Error("Failed to handle workflow", "error", err)
-	}
-}
-
-// 设置希望接收响应的请求ID及对应的通道
 func (c *client) WantResponse(requestId uint64, ch chan []byte) {
 	c.lock.Lock()
-	defer c.lock.Unlock()
-
-	// 关闭之前的通道（如果有的话）
-	if c.wantResponse != nil {
-		close(c.wantResponse)
-	}
-
-	c.wantResponseId = requestId
-	c.wantResponse = ch
+	c.wantResponseId, c.wantResponse = requestId, ch
+	c.lock.Unlock()
+	close(c.wantResponse)
 }
 
-// 根据请求ID发送响应消息
 func (c *client) NewWantResponse(requestId uint64, msg []byte) error {
 	c.lock.Lock()
 	defer c.lock.Unlock()
-	// 检查请求ID是否匹配，确保没有不匹配的消息
 	if requestId != c.wantResponseId || c.wantResponse == nil {
 		return errors.New("unexpected msg")
 	}
-	// 将消息发送到对应的通道
 	c.wantResponse <- msg
 	return nil
-}
+}