SSE请求方式修改

This commit is contained in:
xingchen 2025-04-27 16:17:36 +08:00
parent a002e598f5
commit 882c4af9f4
2 changed files with 34 additions and 14 deletions

View File

@ -7,6 +7,7 @@ import (
"agent/src/workflow"
"context"
"fmt"
"io"
"net/http"
"sync"
)
@ -61,28 +62,36 @@ 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)
http.Error(w, "Streaming unsupported", http.StatusInternalServerError)
return
}
log.Info("[agent] new SSE connection", "remote", r.RemoteAddr)
// 创建新 client
client := newClient(w, flusher, a)
a.addClient(client)
// 设置SSE头
// 设置 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", "*")
c := newClient(w, flusher, a)
// 启动写消息循环
client.run()
// check api key
if a.apicheck.check(c, r) {
c.run()
a.clients.Set(c, struct{}{})
} else {
http.Error(w, "Unauthorized", http.StatusUnauthorized)
// 读取请求体
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) {
@ -113,6 +122,12 @@ 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) removeClient(c *client) {
a.clients.Delete(c)
}

View File

@ -158,10 +158,15 @@ func (c *client) processSSERequest(msg []byte) {
// 设置希望接收响应的请求ID及对应的通道
func (c *client) WantResponse(requestId uint64, ch chan []byte) {
c.lock.Lock()
c.wantResponseId, c.wantResponse = requestId, ch
c.lock.Unlock()
// 确保在设置 wantResponse 之后关闭旧的通道
defer c.lock.Unlock()
// 关闭之前的通道(如果有的话)
if c.wantResponse != nil {
close(c.wantResponse)
}
c.wantResponseId = requestId
c.wantResponse = ch
}
// 根据请求ID发送响应消息