Compare commits
3 Commits
| Author | SHA1 | Date | |
|---|---|---|---|
| 882c4af9f4 | |||
| a002e598f5 | |||
| 67cfd763c3 |
@ -7,6 +7,7 @@ import (
|
||||
"agent/src/workflow"
|
||||
"context"
|
||||
"fmt"
|
||||
"io"
|
||||
"net/http"
|
||||
"sync"
|
||||
)
|
||||
@ -18,6 +19,7 @@ type Agent struct {
|
||||
apicheck *keyChecker
|
||||
clients *common.ThreadSafeMap[*client, struct{}]
|
||||
}
|
||||
|
||||
var once sync.Once
|
||||
|
||||
func NewAgent(ctx context.Context, ollamaUrls []string) (*Agent, error) {
|
||||
@ -29,7 +31,7 @@ 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
|
||||
}
|
||||
@ -37,15 +39,12 @@ func NewAgent(ctx context.Context, ollamaUrls []string) (*Agent, error) {
|
||||
func (a *Agent) Start(port uint64, crtDir *string) {
|
||||
workflow.InitTaskPool(a.ctx)
|
||||
|
||||
go func(){
|
||||
go func() {
|
||||
sm := http.NewServeMux()
|
||||
// sm.HandleFunc("/simpletest2025", a.serveTestPage)
|
||||
// sm.HandleFunc("/doc", a.serveDoc)
|
||||
// sm.HandleFunc("/assets/", a.serveAssets)
|
||||
sm.HandleFunc("/ws", a.serveWs)
|
||||
sm.HandleFunc("/sse", a.serveSSE)
|
||||
|
||||
addr := fmt.Sprintf(":%d", port)
|
||||
log.Info("[agent] start websocket server", "addr", addr)
|
||||
log.Info("[agent] start SSE server", "addr", addr)
|
||||
|
||||
var err error
|
||||
if crtDir == nil {
|
||||
@ -60,6 +59,41 @@ 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 {
|
||||
@ -88,30 +122,12 @@ func (a *Agent) serveAssets(w http.ResponseWriter, r *http.Request) {
|
||||
http.ServeFile(w, r, "./"+r.URL.Path[1:])
|
||||
}
|
||||
|
||||
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()
|
||||
}()
|
||||
}
|
||||
// 给Agent添加addClient方法
|
||||
func (a *Agent) addClient(c *client) {
|
||||
a.clients.Set(c, struct{}{}) // 放进去
|
||||
log.Info("client added", "connected", a.clients.Size())
|
||||
}
|
||||
|
||||
func (a *Agent) removeClient(c *client) {
|
||||
a.clients.Delete(c)
|
||||
}
|
||||
|
||||
|
||||
|
||||
@ -1,189 +1,108 @@
|
||||
package agent
|
||||
|
||||
import (
|
||||
"agent/src/message"
|
||||
"agent/src/workflow"
|
||||
"context"
|
||||
"encoding/json"
|
||||
"errors"
|
||||
"fmt"
|
||||
"net/http"
|
||||
"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
|
||||
conn *websocket.Conn
|
||||
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 // 新增字段：用于存储响应的通道
|
||||
|
||||
wantResponseId uint64
|
||||
wantResponse chan []byte
|
||||
lock sync.Mutex
|
||||
once sync.Once
|
||||
}
|
||||
|
||||
func newClient(conn *websocket.Conn, a *Agent) *client {
|
||||
// 创建新的客户端对象
|
||||
func newClient(w http.ResponseWriter, f http.Flusher, a *Agent) *client {
|
||||
ctx, fnCancel := context.WithCancel(a.ctx)
|
||||
return &client{
|
||||
ctx : ctx,
|
||||
conn : conn,
|
||||
ctx: ctx,
|
||||
writer: w,
|
||||
flusher: f,
|
||||
outMsg: make(chan []byte, 32),
|
||||
agent : a,
|
||||
agent: a,
|
||||
fnCancel: fnCancel,
|
||||
rate: rate.NewLimiter(5, 1), // 每秒5次限制
|
||||
}
|
||||
}
|
||||
|
||||
// 设置API Key并标记为已认证
|
||||
func (c *client) setApiKey(apikey string) {
|
||||
c.apikey = apikey
|
||||
c.authenticated = true
|
||||
}
|
||||
|
||||
// 启动客户端的处理流程
|
||||
func (c *client) run() {
|
||||
go c.readPump()
|
||||
go c.writePump()
|
||||
}
|
||||
|
||||
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)
|
||||
}
|
||||
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
// writePump 处理SSE的消息推送
|
||||
func (c *client) writePump() {
|
||||
ticker := time.NewTicker(pingPeriod)
|
||||
defer ticker.Stop()
|
||||
defer c.Close()
|
||||
|
||||
done := c.ctx.Done()
|
||||
for {
|
||||
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() {
|
||||
once.Do(func(){
|
||||
c.once.Do(func() {
|
||||
c.fnCancel()
|
||||
c.agent.removeClient(c)
|
||||
c.conn.Close()
|
||||
log.Info("client close", "remote", c.conn.RemoteAddr().String(), "connected", c.agent.clients.Size())
|
||||
log.Info("client closed", "connected", c.agent.clients.Size())
|
||||
})
|
||||
}
|
||||
|
||||
// 将JSON数据发送到客户端
|
||||
func (c *client) WriteJson(data interface{}, block bool) error {
|
||||
b, err := json.MarshalIndent(data, "", "")
|
||||
b, err := json.Marshal(data)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
if block {
|
||||
select{
|
||||
case <-c.ctx.Done():
|
||||
return errors.New("connection closed")
|
||||
case c.outMsg <- b:
|
||||
return nil
|
||||
}
|
||||
} else {
|
||||
select {
|
||||
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
|
||||
return c.WriteText(string(b), block)
|
||||
}
|
||||
|
||||
// 将文本消息发送到客户端
|
||||
func (c *client) WriteText(msg string, block bool) error {
|
||||
if block {
|
||||
select {
|
||||
@ -201,43 +120,64 @@ func (c *client) WriteText(msg string, block bool) error {
|
||||
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
|
||||
return c.WriteText(msg, true)
|
||||
}
|
||||
|
||||
// 获取客户端的上下文
|
||||
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
|
||||
}
|
||||
|
||||
func (c *client) WantResponse(requestId uint64, ch chan []byte) {
|
||||
c.lock.Lock()
|
||||
c.wantResponseId, c.wantResponse = requestId, ch
|
||||
c.lock.Unlock()
|
||||
close(c.wantResponse)
|
||||
// 调用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
|
||||
}
|
||||
|
||||
// 根据请求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
|
||||
}
|
||||
Loading…
x
Reference in New Issue
Block a user