From fad2cc7b652eedbe5df999ec221051a21aec86ad Mon Sep 17 00:00:00 2001
From: Badgerbees <rapabelias@gmail.com>
Date: Tue, 14 Apr 2026 20:03:10 +0700
Subject: [PATCH] OpenAI compat: request stream usage

---
 pkg/providers/openai_compat/provider.go      | 20 +++--
 pkg/providers/openai_compat/provider_test.go | 81 ++++++++++++++++++++
 2 files changed, 96 insertions(+), 5 deletions(-)

diff --git a/pkg/providers/openai_compat/provider.go b/pkg/providers/openai_compat/provider.go
index 98a70cfd2..2958937c1 100644
--- a/pkg/providers/openai_compat/provider.go
+++ b/pkg/providers/openai_compat/provider.go
@@ -257,6 +257,9 @@ func (p *Provider) ChatStream(
 
 	requestBody := p.buildRequestBody(messages, tools, model, options)
 	requestBody["stream"] = true
+	if supportsStreamingUsage(p.apiBase) {
+		requestBody["stream_options"] = map[string]any{"include_usage": true}
+	}
 
 	jsonData, err := json.Marshal(requestBody)
 	if err != nil {
@@ -479,11 +482,10 @@ func isNativeSearchHost(apiBase string) bool {
 	return host == "api.openai.com" || strings.HasSuffix(host, ".openai.azure.com")
 }
 
-// supportsPromptCacheKey reports whether the given API base is known to
-// support the prompt_cache_key request field. Currently only OpenAI's own
-// API and Azure OpenAI support this. All other OpenAI-compatible providers
-// (Mistral, Gemini, DeepSeek, Groq, etc.) reject unknown fields with 422 errors.
-func supportsPromptCacheKey(apiBase string) bool {
+// isOpenAINativeBaseURL reports whether the given API base is the OpenAI or
+// Azure OpenAI native endpoint family. We reuse it for request fields that are
+// known to work only on those native hosts.
+func isOpenAINativeBaseURL(apiBase string) bool {
 	u, err := url.Parse(apiBase)
 	if err != nil {
 		return false
@@ -491,3 +493,11 @@ func supportsPromptCacheKey(apiBase string) bool {
 	host := u.Hostname()
 	return host == "api.openai.com" || strings.HasSuffix(host, ".openai.azure.com")
 }
+
+func supportsPromptCacheKey(apiBase string) bool {
+	return isOpenAINativeBaseURL(apiBase)
+}
+
+func supportsStreamingUsage(apiBase string) bool {
+	return isOpenAINativeBaseURL(apiBase)
+}
diff --git a/pkg/providers/openai_compat/provider_test.go b/pkg/providers/openai_compat/provider_test.go
index d140d63d6..d4597751f 100644
--- a/pkg/providers/openai_compat/provider_test.go
+++ b/pkg/providers/openai_compat/provider_test.go
@@ -991,6 +991,54 @@ func chatWithCacheKey(t *testing.T, apiBase string) map[string]any {
 	return requestBody
 }
 
+func chatStreamWithRequestBody(t *testing.T, apiBase string) map[string]any {
+	t.Helper()
+	var requestBody map[string]any
+
+	server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		if err := json.NewDecoder(r.Body).Decode(&requestBody); err != nil {
+			http.Error(w, err.Error(), http.StatusBadRequest)
+			return
+		}
+
+		w.Header().Set("Content-Type", "text/event-stream")
+		_, _ = io.WriteString(
+			w,
+			"data: {\"choices\":[{\"delta\":{\"content\":\"ok\"},\"finish_reason\":\"stop\"}]}\n\n",
+		)
+		if flusher, ok := w.(http.Flusher); ok {
+			flusher.Flush()
+		}
+		_, _ = io.WriteString(w, "data: [DONE]\n\n")
+		if flusher, ok := w.(http.Flusher); ok {
+			flusher.Flush()
+		}
+	}))
+	defer server.Close()
+
+	p := NewProvider("key", server.URL, "")
+	p.apiBase = apiBase
+	p.httpClient = &http.Client{
+		Transport: roundTripperFunc(func(r *http.Request) (*http.Response, error) {
+			r.URL, _ = url.Parse(server.URL + r.URL.Path)
+			return http.DefaultTransport.RoundTrip(r)
+		}),
+	}
+
+	_, err := p.ChatStream(
+		t.Context(),
+		[]Message{{Role: "user", Content: "hi"}},
+		nil,
+		"test-model",
+		nil,
+		nil,
+	)
+	if err != nil {
+		t.Fatalf("ChatStream() error = %v", err)
+	}
+	return requestBody
+}
+
 func TestProviderChat_PromptCacheKeySentToOpenAI(t *testing.T) {
 	body := chatWithCacheKey(t, "https://api.openai.com/v1")
 	if body["prompt_cache_key"] != "agent-main" {
@@ -998,6 +1046,39 @@ func TestProviderChat_PromptCacheKeySentToOpenAI(t *testing.T) {
 	}
 }
 
+func TestProviderChatStream_IncludesUsageForOpenAI(t *testing.T) {
+	body := chatStreamWithRequestBody(t, "https://api.openai.com/v1")
+	streamOptions, ok := body["stream_options"].(map[string]any)
+	if !ok {
+		t.Fatalf("stream_options = %T, want map[string]any", body["stream_options"])
+	}
+	if got := streamOptions["include_usage"]; got != true {
+		t.Fatalf("stream_options.include_usage = %v, want true", got)
+	}
+}
+
+func TestProviderChatStream_OmitsUsageForNonOpenAI(t *testing.T) {
+	tests := []struct {
+		name    string
+		apiBase string
+	}{
+		{"mistral", "https://api.mistral.ai/v1"},
+		{"gemini", "https://generativelanguage.googleapis.com/v1beta"},
+		{"deepseek", "https://api.deepseek.com/v1"},
+		{"groq", "https://api.groq.com/openai/v1"},
+		{"minimax", "https://api.minimaxi.com/v1"},
+	}
+
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			body := chatStreamWithRequestBody(t, tt.apiBase)
+			if _, exists := body["stream_options"]; exists {
+				t.Fatalf("stream_options should NOT be sent to %s, but was included in request", tt.name)
+			}
+		})
+	}
+}
+
 func TestProviderChat_PromptCacheKeyOmittedForNonOpenAI(t *testing.T) {
 	tests := []struct {
 		name    string