diff --git a/relay/helper/stream_scanner.go b/relay/helper/stream_scanner.go index bdc7437..8a66be7 100644 --- a/relay/helper/stream_scanner.go +++ b/relay/helper/stream_scanner.go @@ -30,13 +30,22 @@ const ( // DataHandlerFunc data 处理回调 type DataHandlerFunc func(data string, sr *StreamResult) -// StreamScannerHandler 三协程 SSE 流扫描器 -// Scanner 协程:读取上游行,过滤 data: 行,推入 dataChan +// StreamScannerHandler 三协程 SSE 流扫描器(对齐 new-api 架构) +// +// Scanner 协程:读取上游行,检测 [DONE] 终止标记并立即退出(不转发 [DONE] 到 dataChan) // DataHandler 协程:从 dataChan 读取,调用 handler 回调,持有 writeMutex // Ping 协程:定期写 SSE 保活注释,持有 writeMutex // -// 超时策略:per-chunk 空闲超时(参考 new-api),每收到一个有效 chunk 就重置计时器。 -// 只要上游持续发送数据,流可以无限延续,适用于长时间思考的模型。 +// 超时策略:per-chunk 空闲超时,scanner.Scan() 每返回一行就通过 resetChan 通知主协程重置 timer。 +// scanner.Scan() 阻塞期间不重置,超时从最后收到的字节开始算。 +// 对齐 new-api 的 ticker.Reset(streamingTimeout) 行为,但用 Timer+resetChan 避免 ticker 跨协程操作。 +// +// 终止流程(对齐 new-api): +// 1. Scanner 检测到 [DONE] → return → close(dataChan) → safeSendBool(stopChan) +// 2. DataHandler range 退出 → safeSendBool(stopChan) +// 3. 主 goroutine 收到 stopChan → 进入 shutdown +// 4. shutdown: safeSendBool(stopChan) 通知剩余协程 → wg.Wait(最多 5s) → resp.Body.Close() +// 5. [DONE] 由 handler 层(如 openai/stream.go)负责写入客户端 func StreamScannerHandler( ctx context.Context, resp *http.Response, @@ -46,24 +55,38 @@ func StreamScannerHandler( ) { stopChan := make(chan bool, stopChanSize) dataChan := make(chan string, dataChanSize) - resetChan := make(chan struct{}, resetChanSize) // 收到有效数据时通知主 goroutine 重置超时 + resetChan := make(chan struct{}, resetChanSize) // scanner 每读一行通知主协程重置超时 var writeMutex sync.Mutex var wg sync.WaitGroup streamTimeout := getStreamTimeout(info) + // 确保 resp.Body 最终被关闭(对齐 new-api),解锁被 scanner.Scan() 阻塞的协程 + defer func() { + if resp.Body != nil { + resp.Body.Close() + } + }() + SetEventStreamHeaders(writer) // Scanner 协程 wg.Add(1) go func() { - defer wg.Done() - defer close(dataChan) + defer func() { + close(dataChan) + wg.Done() + if r := recover(); r != nil { + info.StreamStatus.SetEndReason(common.StreamEndReasonError, fmt.Errorf("scanner panic: %v", r)) + } + safeSendBool(stopChan) + }() scanner := bufio.NewScanner(resp.Body) buf := make([]byte, 0, scannerInitialBuf) scanner.Buffer(buf, scannerMaxBuf) for scanner.Scan() { + // 检查停止信号 select { case <-stopChan: return @@ -73,13 +96,22 @@ func StreamScannerHandler( default: } + // 每读一行就通知主协程重置空闲超时(对齐 new-api ticker.Reset) + select { + case resetChan <- struct{}{}: + default: + } + line := scanner.Text() if len(line) < minLineLength { continue } + // 兼容两种 [DONE] 格式: + // 1. 裸 "[DONE]"(少数上游) + // 2. 标准 SSE "data: [DONE]"(主流上游如 OpenAI/DeepSeek) + // 检测到后立即退出 scanner,由 handler 层负责向客户端写入 [DONE]。 if line == "[DONE]" { - safeSendBool(stopChan) info.StreamStatus.SetEndReason(common.StreamEndReasonDone, nil) return } @@ -93,14 +125,13 @@ func StreamScannerHandler( continue } - if info.StreamStatus.GetEndReason() == "" { - info.SetFirstResponseTime() + if data == "[DONE]" { + info.StreamStatus.SetEndReason(common.StreamEndReasonDone, nil) + return } - // 通知主 goroutine 重置空闲超时 - select { - case resetChan <- struct{}{}: - default: + if info.StreamStatus.GetEndReason() == "" { + info.SetFirstResponseTime() } select { @@ -115,10 +146,8 @@ func StreamScannerHandler( if err := scanner.Err(); err != nil && err != io.EOF && ctx.Err() == nil { info.StreamStatus.SetEndReason(common.StreamEndReasonScannerErr, err) - safeSendBool(stopChan) } else if info.StreamStatus.GetEndReason() == "" { info.StreamStatus.SetEndReason(common.StreamEndReasonEOF, nil) - safeSendBool(stopChan) } }() @@ -126,14 +155,19 @@ func StreamScannerHandler( sr := NewStreamResult(info.StreamStatus) wg.Add(1) go func() { - defer wg.Done() + defer func() { + wg.Done() + if r := recover(); r != nil { + info.StreamStatus.SetEndReason(common.StreamEndReasonError, fmt.Errorf("handler panic: %v", r)) + } + safeSendBool(stopChan) + }() for data := range dataChan { sr.reset() writeMutex.Lock() handler(data, sr) writeMutex.Unlock() if sr.IsStopped() { - safeSendBool(stopChan) return } } @@ -168,8 +202,7 @@ func StreamScannerHandler( } }() - // 主 goroutine 阻塞等待 - // 使用 Timer 而非 Ticker,每次收到有效数据时 Reset,实现 per-chunk 空闲超时 + // 主 goroutine 阻塞等待(对齐 new-api:单次 select,用 Timer 实现空闲超时) timeoutTimer := time.NewTimer(streamTimeout) defer timeoutTimer.Stop() @@ -185,7 +218,7 @@ func StreamScannerHandler( info.StreamStatus.SetEndReason(common.StreamEndReasonClientGone, ctx.Err()) goto shutdown case <-resetChan: - // 收到有效数据,重置空闲超时 + // scanner 每读一行都通知重置(对齐 new-api ticker.Reset) if !timeoutTimer.Stop() { select { case <-timeoutTimer.C: @@ -197,7 +230,9 @@ func StreamScannerHandler( } shutdown: - // 等待所有协程退出 + // 通知所有协程停止(对齐 new-api:先 signal 再 wait) + safeSendBool(stopChan) + waitDone := make(chan struct{}) go func() { wg.Wait()