Update sse.go

sse.go

@@ -40,7 +40,8 @@ func streamOpenAIResponseToClaudeSSE(
 	calculatedInputTokens := calculateInputTokensFromClaudeRequest(originalClaudeRequest)
 	log.Printf("INFO: [%s] (OpenAI->Claude SSE) Bắt đầu SSE. Input tokens ước tính: %d. Model: %s", messageID, calculatedInputTokens, requestedModel)
 
-	inputTokens := calculatedInputTokens
+	// inputTokens from OpenAI stream can update this if OpenAI provides it
+	currentOpenAIInputTokens := calculatedInputTokens 
 	outputTokens := 0
 	finalUsageReceivedFromStream := false // Flag to check if usage is explicitly received from OpenAI stream
 	eventIndex := 0
@@ -120,11 +121,11 @@ func streamOpenAIResponseToClaudeSSE(
 				}
 				if chunk.Usage != nil { // Some OpenAI versions might send usage in chunks
 					if chunk.Usage.CompletionTokens > 0 {
-						outputTokens = chunk.Usage.CompletionTokens
+						outputTokens = chunk.Usage.CompletionTokens // Use OpenAI's count if available
 						finalUsageReceivedFromStream = true
 					}
 					if chunk.Usage.PromptTokens > 0 {
-						inputTokens = chunk.Usage.PromptTokens // Update input tokens if provided by OpenAI
+						currentOpenAIInputTokens = chunk.Usage.PromptTokens // Update input tokens if provided by OpenAI
 					}
 				}
 			}
@@ -180,7 +181,7 @@ func streamOpenAIResponseToClaudeSSE(
 		claudeStopReason = mapOpenAIFinishReasonToClaude(openAIFinishReason)
 	}
 
-	finalInputTokens := inputTokens
+	finalInputTokens := currentOpenAIInputTokens // Use potentially updated input tokens from OpenAI stream
 	finalOutputTokens := outputTokens
 	if !finalUsageReceivedFromStream { // If OpenAI didn't send final usage, estimate from accumulated content
 		finalOutputTokens = estimateTokens(accumulatedContent)
@@ -250,12 +251,12 @@ func proxyClaudeNativeSSE(
 	calculatedInputTokens := calculateInputTokensFromClaudeRequest(originalClaudeRequest)
 	log.Printf("INFO: [%s] (Claude Native SSE with Hack) Bắt đầu SSE. Input tokens ước tính: %d. Model: %s", messageID, calculatedInputTokens, requestedModel)
 
-	inputTokens := calculatedInputTokens // inputTokens sẽ không thay đổi trong quá trình stream này
+	// removed: inputTokens := calculatedInputTokens // This variable was unused
 	outputTokens := 0
 	eventIndex := 0
 
 	var upstreamFinalStopReason string
-	var upstreamFinalUsage *ClaudeUsage
+	var upstreamFinalUsage *ClaudeUsage // This is of type *ClaudeUsage
 
 	doneChan := make(chan struct{})
 	errChan := make(chan error, 1)
@@ -270,7 +271,6 @@ func proxyClaudeNativeSSE(
 	eventIndex++
 
 	// Gửi content_block_start cho client (tự tạo)
-	// Claude native thường gửi content_block_start với index 0 và content_block.type="text", content_block.text=""
 	contentStartBlock := ClaudeSSEContentBlock{Type: "text", Text: ""}
 	contentStartEvent := ClaudeSSEEvent{Type: "content_block_start", Index: func() *int { i := 0; return &i }(), ContentBlock: &contentStartBlock}
 	if !sendSSEEvent(c, "content_block_start", contentStartEvent, messageID, eventIndex, true, "(Claude Native SSE with Hack)") {
@@ -296,27 +296,23 @@ func proxyClaudeNativeSSE(
 				continue
 			}
 
-			// Chỉ xử lý các dòng data, bỏ qua event: lines vì chúng ta sẽ tự tạo/forward có chọn lọc
 			if strings.HasPrefix(line, "data:") {
 				dataStr := strings.TrimSpace(strings.TrimPrefix(line, "data:"))
-				var upstreamEvent ClaudeSSEEvent // Sử dụng struct ClaudeSSEEvent để parse
+				var upstreamEvent ClaudeSSEEvent 
 				if err := json.Unmarshal([]byte(dataStr), &upstreamEvent); err != nil {
 					log.Printf("WARN: [%s] (Claude Native SSE with Hack) Không thể giải mã chunk JSON Claude native: %v. Data: %s", messageID, err, dataStr)
 					continue
 				}
 
-				// Xử lý các loại event từ Claude native upstream
 				switch upstreamEvent.Type {
 				case "content_block_delta":
 					if upstreamEvent.Delta != nil && upstreamEvent.Delta.Text != nil {
 						textChunk := *upstreamEvent.Delta.Text
 						accumulatedContent += textChunk
 
-						// Forward content_block_delta đến client
-						// Đảm bảo index được set (thường là 0 cho một content block đơn giản)
 						clientContentDelta := ClaudeSSEEvent{
 							Type:  "content_block_delta",
-							Index: upstreamEvent.Index, // Sử dụng index từ upstream nếu có, nếu không thì 0
+							Index: upstreamEvent.Index, 
 							Delta: &ClaudeSSEDelta{
 								Type: "text_delta",
 								Text: &textChunk,
@@ -325,43 +321,41 @@ func proxyClaudeNativeSSE(
 						if clientContentDelta.Index == nil { clientContentDelta.Index = func() *int { i := 0; return &i }() }
 
 						if !sendSSEEvent(c, "content_block_delta", clientContentDelta, messageID, eventIndex, false, "(Claude Native SSE with Hack)") {
-							return // Client disconnected
+							return 
 						}
 						eventIndex++
 
-						// "Hack" để gửi message_delta với output_tokens cập nhật
 						currentOutputTokens := estimateTokens(accumulatedContent)
 						if currentOutputTokens != outputTokens {
 							outputTokens = currentOutputTokens
 							intermediateUsage := ClaudeSSEUsage{OutputTokens: outputTokens}
 							intermediateDeltaPayload := ClaudeSSEEvent{
 								Type:  "message_delta",
-								Delta: &ClaudeSSEDelta{}, // Delta rỗng
+								Delta: &ClaudeSSEDelta{}, 
 								Usage: &intermediateUsage,
 							}
 							if !sendSSEEvent(c, "message_delta", intermediateDeltaPayload, messageID, eventIndex, false, "(Claude Native SSE with Hack)") {
-								return // Client disconnected
+								return 
 							}
 							eventIndex++
 						}
 					}
-				case "message_delta": // Claude native có thể gửi message_delta với usage
-					if upstreamEvent.Usage != nil {
-						// Cập nhật outputTokens nếu upstream cung cấp
-						// Đây là thông tin chính xác hơn là ước tính
+				case "message_delta": 
+					if upstreamEvent.Usage != nil { // upstreamEvent.Usage is *ClaudeSSEUsage
 						if upstreamEvent.Usage.OutputTokens > outputTokens {
 							outputTokens = upstreamEvent.Usage.OutputTokens
 						}
-						// Gửi message_delta "hack" với usage từ upstream (nếu có)
-						// hoặc usage đã tính toán nếu upstream không gửi
-						currentHackUsage := ClaudeSSEUsage{OutputTokens: outputTokens}
+						currentHackUsage := ClaudeSSEUsage{OutputTokens: outputTokens} // Use updated outputTokens
+						// Preserve stop_reason/sequence from upstream's message_delta if present
+						var deltaDetails ClaudeSSEDelta
+						if upstreamEvent.Delta != nil {
+							deltaDetails.StopReason = upstreamEvent.Delta.StopReason
+							deltaDetails.StopSequence = upstreamEvent.Delta.StopSequence
+						}
+
 						hackMessageDelta := ClaudeSSEEvent{
-							Type: "message_delta",
-							Delta: &ClaudeSSEDelta{
-								// Claude native message_delta có thể có stop_reason/sequence
-								StopReason: upstreamEvent.Delta.StopReason,
-								StopSequence: upstreamEvent.Delta.StopSequence,
-							},
+							Type:  "message_delta",
+							Delta: &deltaDetails,
 							Usage: &currentHackUsage,
 						}
 						if !sendSSEEvent(c, "message_delta", hackMessageDelta, messageID, eventIndex, false, "(Claude Native SSE with Hack - from upstream delta)") {
@@ -369,42 +363,46 @@ func proxyClaudeNativeSSE(
 						}
 						eventIndex++
 					}
+					// Capture stop reason if it's in this message_delta
 					if upstreamEvent.Delta != nil && upstreamEvent.Delta.StopReason != nil {
 						upstreamFinalStopReason = *upstreamEvent.Delta.StopReason
 					}
 
 				case "message_stop":
 					log.Printf("INFO: [%s] (Claude Native SSE with Hack) Nhận message_stop từ upstream.", messageID)
-					if upstreamEvent.Message != nil && upstreamEvent.Message.StopReason != nil { // Older format
+					if upstreamEvent.Message != nil && upstreamEvent.Message.StopReason != nil { 
 						upstreamFinalStopReason = *upstreamEvent.Message.StopReason
 					}
-					// Newer Claude API sends stop_reason in message_delta, and message_stop has final usage.
-					if upstreamEvent.Usage != nil { // Final usage from upstream
-						upstreamFinalUsage = upstreamEvent.Usage
-						if upstreamEvent.Usage.OutputTokens > outputTokens { // Prefer upstream's final count
+					
+					if upstreamEvent.Usage != nil { // upstreamEvent.Usage is *ClaudeSSEUsage
+						// Convert ClaudeSSEUsage to ClaudeUsage for upstreamFinalUsage
+						tempUsage := &ClaudeUsage{
+							OutputTokens: upstreamEvent.Usage.OutputTokens,
+						}
+						if upstreamEvent.Usage.InputTokens != nil {
+							tempUsage.InputTokens = *upstreamEvent.Usage.InputTokens
+						}
+						// Assign the converted *ClaudeUsage to upstreamFinalUsage
+						upstreamFinalUsage = tempUsage 
+
+						if upstreamEvent.Usage.OutputTokens > outputTokens { 
 							outputTokens = upstreamEvent.Usage.OutputTokens
 						}
 					}
-					return // Kết thúc vòng lặp khi nhận message_stop
+					return 
 
 				case "error":
 					log.Printf("ERROR: [%s] (Claude Native SSE with Hack) Nhận lỗi từ upstream: %+v", messageID, upstreamEvent.Error)
 					streamErrorOccurred = true
 					errorDetails = upstreamEvent.Error
-					return // Kết thúc vòng lặp
+					return 
 				
-				// Các event khác như message_start, content_block_start, content_block_stop từ upstream sẽ được bỏ qua
-				// vì chúng ta đã tự tạo các event tương ứng cho client ở đầu luồng.
-				// Ping events có thể được forward nếu cần.
 				case "ping":
 					pingEvent := ClaudeSSEEvent{Type: "ping"}
 					if !sendSSEEvent(c, "ping", pingEvent, messageID, eventIndex, false, "(Claude Native SSE with Hack - ping)") {
 						return
 					}
 					eventIndex++
-
-				default:
-					// log.Printf("DEBUG: [%s] (Claude Native SSE with Hack) Bỏ qua event upstream không xác định: %s", messageID, upstreamEvent.Type)
 				}
 			}
 		}
@@ -421,17 +419,17 @@ func proxyClaudeNativeSSE(
 
 	select {
 	case <-doneChan:
-		// Normal completion of upstream read
+		// Normal completion
 	case err := <-errChan:
 		log.Printf("ERROR: [%s] (Claude Native SSE with Hack) Nhận lỗi từ goroutine đọc: %v", messageID, err)
 		streamErrorOccurred = true
-		if errorDetails == nil { // Ensure errorDetails is set
+		if errorDetails == nil { 
 			errorDetails = &ClaudeError{Type: "api_error", Message: fmt.Sprintf("Lỗi đọc phản hồi upstream: %v", err)}
 		}
 	case <-c.Request.Context().Done():
 		log.Printf("INFO: [%s] (Claude Native SSE with Hack) Client ngắt kết nối trong quá trình xử lý stream: %v", messageID, c.Request.Context().Err())
 		streamErrorOccurred = true
-		if errorDetails == nil { // Ensure errorDetails is set
+		if errorDetails == nil { 
 			errorDetails = &ClaudeError{Type: "client_disconnect", Message: "Client ngắt kết nối trong quá trình stream"}
 		}
 	}
@@ -444,51 +442,51 @@ func proxyClaudeNativeSSE(
 	} else if upstreamFinalStopReason != "" {
 		finalStopReasonClient = upstreamFinalStopReason
 	} else {
-		finalStopReasonClient = "end_turn" // Default if no specific reason from upstream
+		finalStopReasonClient = "end_turn" 
 	}
 
-	// Xác định finalOutputTokens
 	finalOutputTokens := outputTokens
-	if upstreamFinalUsage != nil && upstreamFinalUsage.OutputTokens > finalOutputTokens {
-		finalOutputTokens = upstreamFinalUsage.OutputTokens // Ưu tiên usage từ message_stop của upstream
-	} else {
-		finalOutputTokens = estimateTokens(accumulatedContent) // Nếu không có, ước tính lại
+	if upstreamFinalUsage != nil && upstreamFinalUsage.OutputTokens > finalOutputTokens { // upstreamFinalUsage is *ClaudeUsage
+		finalOutputTokens = upstreamFinalUsage.OutputTokens 
+	} else if accumulatedContent != "" { // Only estimate if content was accumulated and no definitive upstream usage
+		finalOutputTokens = estimateTokens(accumulatedContent)
 	}
 	finalOutputTokens = max(0, finalOutputTokens)
 
-
 	log.Printf("INFO: [%s] (Claude Native SSE with Hack) SSE Stream hoàn tất. Lý do dừng: %s. Input: %d, Output: %d. Thời gian: %v. Upstream Stop Reason: %s",
 		messageID, finalStopReasonClient, calculatedInputTokens, finalOutputTokens, time.Since(startTime), upstreamFinalStopReason)
 
-	// Gửi "Hack" message_delta cuối cùng với stop_reason và usage (chỉ output_tokens)
 	finalHackUsageData := ClaudeSSEUsage{OutputTokens: finalOutputTokens}
 	finalDeltaStopReason := finalStopReasonClient
 	priorityFinalDeltaPayload := ClaudeSSEEvent{
 		Type: "message_delta",
 		Delta: &ClaudeSSEDelta{
 			StopReason:   &finalDeltaStopReason,
-			StopSequence: nil, // Claude native không thường gửi stop_sequence trong delta này
+			StopSequence: nil, 
 		},
 		Usage: &finalHackUsageData,
 	}
 	_ = sendSSEEvent(c, "message_delta", priorityFinalDeltaPayload, messageID, eventIndex, false, "(Claude Native SSE with Hack)")
 	eventIndex++
 
-	// Gửi content_block_stop cho client
 	contentStopPayload := ClaudeSSEEvent{Type: "content_block_stop", Index: func() *int { i := 0; return &i }()}
 	_ = sendSSEEvent(c, "content_block_stop", contentStopPayload, messageID, eventIndex, false, "(Claude Native SSE with Hack)")
 	eventIndex++
 
-	// Gửi message_stop cho client với usage đầy đủ
-	// Input tokens là cái đã tính toán, output tokens là cái đã chốt cuối cùng
-	finalClientInputTokens := calculatedInputTokens
+	finalClientInputTokens := calculatedInputTokens // Use the initially calculated input tokens for consistency
+	// If upstream provided definitive input tokens in its final usage, we could consider using it,
+	// but calculatedInputTokens is based on the original request, which is reliable for the client's perspective.
+	if upstreamFinalUsage != nil && upstreamFinalUsage.InputTokens > 0 {
+		// Potentially use upstreamFinalUsage.InputTokens if it's considered more accurate
+		// For now, sticking to calculatedInputTokens for the client-facing message_stop.
+	}
+
 	finalStopUsageData := ClaudeSSEUsage{InputTokens: &finalClientInputTokens, OutputTokens: finalOutputTokens}
 	messageStopPayload := ClaudeSSEEvent{Type: "message_stop", Usage: &finalStopUsageData}
 	_ = sendSSEEvent(c, "message_stop", messageStopPayload, messageID, eventIndex, true, "(Claude Native SSE with Hack)")
 	eventIndex++
 
 	if streamErrorOccurred && errorDetails != nil && errorDetails.Type != "client_disconnect" {
-		// Nếu lỗi không phải do client disconnect, gửi event error cho client
 		errorPayload := ClaudeSSEEvent{Type: "error", Error: errorDetails}
 		_ = sendSSEEvent(c, "error", errorPayload, messageID, eventIndex, true, "(Claude Native SSE with Hack)")
 	}
@@ -506,21 +504,16 @@ func sendSSEEvent(c *gin.Context, eventName string, data interface{}, requestID
 		jsonData, err := json.Marshal(data)
 		if err != nil {
 			log.Printf("ERROR: [%s] %s Không thể marshal SSE event %d (%s): %v", requestID, logPrefix, eventIndex, eventName, err)
-			return true // Continue trying to send other events if this one fails to marshal
+			return true 
 		}
-		// Đảm bảo có một dòng trống sau mỗi sự kiện data:
 		_, err = fmt.Fprintf(c.Writer, "event: %s\ndata: %s\n\n", eventName, string(jsonData))
 		if err != nil {
 			if shouldLog || eventName == "message_stop" || eventName == "error" {
 				log.Printf("WARN: [%s] %s Không thể ghi SSE event %d (%s) cho client: %v. Client có thể đã ngắt kết nối.", requestID, logPrefix, eventIndex, err)
 			}
-			return false // Stop if writing to client fails
+			return false 
 		}
 		c.Writer.Flush()
-		// Vô hiệu hóa: Log TRACE này quá chi tiết
-		// if shouldLog {
-		// 	log.Printf("TRACE: [%s] %s SSE: Đã gửi Event %d (%s)", requestID, logPrefix, eventIndex, eventName)
-		// }
 		return true
 	}
 }