diff --git a/relay/channel/claude/relay-claude.go b/relay/channel/claude/relay-claude.go index 3818af3..f34d9a7 100644 --- a/relay/channel/claude/relay-claude.go +++ b/relay/channel/claude/relay-claude.go @@ -8,6 +8,7 @@ import ( "io" "net/http" "one-api/common" + "one-api/constant" "one-api/dto" relaycommon "one-api/relay/common" "one-api/service" @@ -267,8 +268,8 @@ func claudeStreamHandler(c *gin.Context, resp *http.Response, info *relaycommon. } return 0, nil, nil }) - dataChan := make(chan string) - stopChan := make(chan bool) + dataChan := make(chan string, 5) + stopChan := make(chan bool, 2) go func() { for scanner.Scan() { data := scanner.Text() @@ -276,7 +277,11 @@ func claudeStreamHandler(c *gin.Context, resp *http.Response, info *relaycommon. continue } data = strings.TrimPrefix(data, "data: ") - dataChan <- data + if !common.SafeSendStringTimeout(dataChan, data, constant.StreamingTimeout) { + // send data timeout, stop the stream + common.LogError(c, "send data timeout, stop the stream") + break + } } stopChan <- true }() diff --git a/relay/channel/gemini/relay-gemini.go b/relay/channel/gemini/relay-gemini.go index 2ba5cef..8af08c5 100644 --- a/relay/channel/gemini/relay-gemini.go +++ b/relay/channel/gemini/relay-gemini.go @@ -7,6 +7,7 @@ import ( "io" "net/http" "one-api/common" + "one-api/constant" "one-api/dto" relaycommon "one-api/relay/common" "one-api/service" @@ -163,8 +164,8 @@ func streamResponseGeminiChat2OpenAI(geminiResponse *GeminiChatResponse) *dto.Ch func geminiChatStreamHandler(c *gin.Context, resp *http.Response, info *relaycommon.RelayInfo) (*dto.OpenAIErrorWithStatusCode, string) { responseText := "" - dataChan := make(chan string) - stopChan := make(chan bool) + dataChan := make(chan string, 5) + stopChan := make(chan bool, 2) scanner := bufio.NewScanner(resp.Body) scanner.Split(func(data []byte, atEOF bool) (advance int, token []byte, err error) { if atEOF && len(data) == 0 { @@ -187,7 +188,11 @@ func geminiChatStreamHandler(c *gin.Context, resp *http.Response, info *relaycom } data = strings.TrimPrefix(data, "\"text\": \"") data = strings.TrimSuffix(data, "\"") - dataChan <- data + if !common.SafeSendStringTimeout(dataChan, data, constant.StreamingTimeout) { + // send data timeout, stop the stream + common.LogError(c, "send data timeout, stop the stream") + break + } } stopChan <- true }() diff --git a/relay/channel/openai/relay-openai.go b/relay/channel/openai/relay-openai.go index 5146a4f..5c2acf4 100644 --- a/relay/channel/openai/relay-openai.go +++ b/relay/channel/openai/relay-openai.go @@ -54,7 +54,7 @@ func OpenaiStreamHandler(c *gin.Context, resp *http.Response, info *relaycommon. } if !common.SafeSendStringTimeout(dataChan, data, constant.StreamingTimeout) { // send data timeout, stop the stream - common.LogInfo(c, "send data timeout, stop the stream") + common.LogError(c, "send data timeout, stop the stream") break } data = data[6:]