fix: improve streaming bootstrap and forwarding

sdk/api/handlers/gemini/gemini-cli_handlers.go

@@ -182,19 +182,18 @@ func (h *GeminiCLIAPIHandler) handleInternalGenerateContent(c *gin.Context, rawJ
 }
 
 func (h *GeminiCLIAPIHandler) forwardCLIStream(c *gin.Context, flusher http.Flusher, alt string, cancel func(error), data <-chan []byte, errs <-chan *interfaces.ErrorMessage) {
-	for {
-		select {
-		case <-c.Request.Context().Done():
-			cancel(c.Request.Context().Err())
-			return
-		case chunk, ok := <-data:
-			if !ok {
-				cancel(nil)
-				return
-			}
+	var keepAliveInterval *time.Duration
+	if alt != "" {
+		disabled := time.Duration(0)
+		keepAliveInterval = &disabled
+	}
+
+	h.ForwardStream(c, flusher, cancel, data, errs, handlers.StreamForwardOptions{
+		KeepAliveInterval: keepAliveInterval,
+		WriteChunk: func(chunk []byte) {
 			if alt == "" {
 				if bytes.Equal(chunk, []byte("data: [DONE]")) || bytes.Equal(chunk, []byte("[DONE]")) {
-					continue
+					return
 				}
 
 				if !bytes.HasPrefix(chunk, []byte("data:")) {
@@ -206,22 +205,25 @@ func (h *GeminiCLIAPIHandler) forwardCLIStream(c *gin.Context, flusher http.Flus
 			} else {
 				_, _ = c.Writer.Write(chunk)
 			}
-			flusher.Flush()
-		case errMsg, ok := <-errs:
-			if !ok {
-				continue
+		},
+		WriteTerminalError: func(errMsg *interfaces.ErrorMessage) {
+			if errMsg == nil {
+				return
 			}
-			if errMsg != nil {
-				h.WriteErrorResponse(c, errMsg)
-				flusher.Flush()
+			status := http.StatusInternalServerError
+			if errMsg.StatusCode > 0 {
+				status = errMsg.StatusCode
 			}
-			var execErr error
-			if errMsg != nil {
-				execErr = errMsg.Error
+			errText := http.StatusText(status)
+			if errMsg.Error != nil && errMsg.Error.Error() != "" {
+				errText = errMsg.Error.Error()
 			}
-			cancel(execErr)
-			return
-		case <-time.After(500 * time.Millisecond):
-		}
-	}
+			body := handlers.BuildErrorResponseBody(status, errText)
+			if alt == "" {
+				_, _ = fmt.Fprintf(c.Writer, "event: error\ndata: %s\n\n", string(body))
+			} else {
+				_, _ = c.Writer.Write(body)
+			}
+		},
+	})
 }