From b6a157c19c5741aa8b7b06f14e1db2b8270a7311 Mon Sep 17 00:00:00 2001 From: SPHINX998 <51872346+SPHINX998@users.noreply.github.com> Date: Thu, 1 Oct 2026 22:30:02 +0800 Subject: [PATCH] fix(openai): account image input usage on WS passthrough The WS v2 passthrough relay only read output_tokens_details.image_tokens, so image generation over WebSocket billed image input tokens at the text input rate and ignored the hosted response.tool_usage.image_gen breakdown that the HTTP Responses path already merges (extractOpenAIUsageFromJSONBytes + mergeHostedImageGenToolUsage). The relay now reads image input tokens from input_tokens_details or prompt_tokens_details, fills missing image input/output counters from response.tool_usage.image_gen (usage values keep precedence), and carries ImageInputTokens per turn and in the session aggregate into the OpenAIForwardResult used for billing. Co-Authored-By: Claude Opus 5.5 --- .../service/openai_ws_v2/passthrough_relay.go | 23 +++- .../passthrough_relay_image_usage_test.go | 47 ++++++++ .../openai_ws_v2_passthrough_adapter.go | 2 + ...enai_ws_v2_passthrough_image_usage_test.go | 103 ++++++++++++++++++ 4 files changed, 174 insertions(+), 1 deletion(-) create mode 100644 backend/internal/service/openai_ws_v2/passthrough_relay_image_usage_test.go create mode 100644 backend/internal/service/openai_ws_v2_passthrough_image_usage_test.go diff --git a/backend/internal/service/openai_ws_v2/passthrough_relay.go b/backend/internal/service/openai_ws_v2/passthrough_relay.go index 3c6bd076dc..0fdbb90852 100644 --- a/backend/internal/service/openai_ws_v2/passthrough_relay.go +++ b/backend/internal/service/openai_ws_v2/passthrough_relay.go @@ -26,6 +26,7 @@ type FrameConn interface { type Usage struct { InputTokens int + ImageInputTokens int OutputTokens int CacheCreationInputTokens int CacheReadInputTokens int @@ -1117,6 +1118,21 @@ func parseUsageAndAccumulate( if imageTokens == 0 { imageTokens = usageResult.Get("completion_tokens_details.image_tokens").Int() } + imageInputTokens := usageResult.Get("input_tokens_details.image_tokens").Int() + if imageInputTokens <= 0 { + imageInputTokens = usageResult.Get("prompt_tokens_details.image_tokens").Int() + } + // 与 HTTP Responses 同口径：usage 缺图片计数时由托管工具 tool_usage.image_gen 回填。 + imageGen := gjson.GetBytes(message, "response.tool_usage.image_gen") + if !imageGen.Exists() { + imageGen = gjson.GetBytes(message, "tool_usage.image_gen") + } + if imageTokens <= 0 { + imageTokens = max(imageGen.Get("output_tokens_details.image_tokens").Int(), 0) + } + if imageInputTokens <= 0 { + imageInputTokens = max(imageGen.Get("input_tokens_details.image_tokens").Int(), 0) + } requireTotals := isTerminalEvent(strings.TrimSpace(eventType)) inputTokens, inputOK := parseUsageIntField(inputResult, requireTotals) @@ -1141,6 +1157,7 @@ func parseUsageAndAccumulate( } parsedUsage := Usage{ InputTokens: inputTokens, + ImageInputTokens: int(imageInputTokens), OutputTokens: outputTokens, CacheCreationInputTokens: openAICacheCreationTokensFromUsage(usageResult), CacheReadInputTokens: cachedTokens, @@ -1161,7 +1178,7 @@ func parseUsageAndAccumulate( func relayUsageHasTokens(usage Usage) bool { return usage.InputTokens > 0 || usage.OutputTokens > 0 || usage.CacheCreationInputTokens > 0 || usage.CacheReadInputTokens > 0 || - usage.ImageOutputTokens > 0 + usage.ImageInputTokens > 0 || usage.ImageOutputTokens > 0 } func mergeRelayUsageNonZero(dst *Usage, src Usage) { @@ -1180,6 +1197,9 @@ func mergeRelayUsageNonZero(dst *Usage, src Usage) { if src.CacheReadInputTokens > 0 { dst.CacheReadInputTokens = src.CacheReadInputTokens } + if src.ImageInputTokens > 0 { + dst.ImageInputTokens = src.ImageInputTokens + } if src.ImageOutputTokens > 0 { dst.ImageOutputTokens = src.ImageOutputTokens } @@ -1191,6 +1211,7 @@ func finalizeRelayTurnUsage(state *relayState) Usage { } turnUsage := state.turnUsage state.usage.InputTokens += turnUsage.InputTokens + state.usage.ImageInputTokens += turnUsage.ImageInputTokens state.usage.OutputTokens += turnUsage.OutputTokens state.usage.CacheCreationInputTokens += turnUsage.CacheCreationInputTokens state.usage.CacheReadInputTokens += turnUsage.CacheReadInputTokens diff --git a/backend/internal/service/openai_ws_v2/passthrough_relay_image_usage_test.go b/backend/internal/service/openai_ws_v2/passthrough_relay_image_usage_test.go new file mode 100644 index 0000000000..74beacd6c9 --- /dev/null +++ b/backend/internal/service/openai_ws_v2/passthrough_relay_image_usage_test.go @@ -0,0 +1,47 @@ +package openai_ws_v2 + +import ( + "context" + "testing" + "time" + + coderws "github.com/coder/websocket" + "github.com/stretchr/testify/require" +) + +// WS passthrough must account image tokens the way the HTTP Responses path +// does: image input from input/prompt_tokens_details, and the hosted +// tool_usage.image_gen breakdown filling counters the usage object omits. +func TestRelay_ImageUsageMatchesHTTPResponsesAccounting(t *testing.T) { + t.Parallel() + + clientConn := newPassthroughTestFrameConn(nil, false) + upstreamConn := newPassthroughTestFrameConn([]passthroughTestFrame{ + {msgType: coderws.MessageText, payload: []byte(`{"type":"response.completed","response":{"id":"resp_tool_image","usage":{"input_tokens":100,"output_tokens":60},"tool_usage":{"image_gen":{"input_tokens_details":{"image_tokens":80},"output_tokens_details":{"image_tokens":50}}}}}`)}, + {msgType: coderws.MessageText, payload: []byte(`{"type":"response.completed","response":{"id":"resp_explicit_image","usage":{"input_tokens":30,"output_tokens":15,"input_tokens_details":{"image_tokens":20},"output_tokens_details":{"image_tokens":10}},"tool_usage":{"image_gen":{"input_tokens_details":{"image_tokens":800},"output_tokens_details":{"image_tokens":500}}}}}`)}, + {msgType: coderws.MessageText, payload: []byte(`{"type":"response.completed","response":{"id":"resp_prompt_details","usage":{"prompt_tokens":12,"completion_tokens":2,"prompt_tokens_details":{"image_tokens":7}}}}`)}, + {msgType: coderws.MessageText, payload: []byte(`{"type":"response.completed","response":{"id":"resp_text","usage":{"input_tokens":3,"output_tokens":4}}}`)}, + }, true) + ctx, cancel := context.WithTimeout(context.Background(), 2*time.Second) + defer cancel() + + var turns []RelayTurnResult + result, relayExit := Relay(ctx, clientConn, upstreamConn, + []byte(`{"type":"response.create","model":"gpt-5.4","tools":[{"type":"image_generation"}],"input":[]}`), + RelayOptions{OnTurnComplete: func(turn RelayTurnResult) { turns = append(turns, turn) }}) + require.Nil(t, relayExit) + require.Len(t, turns, 4) + + require.Equal(t, 80, turns[0].Usage.ImageInputTokens, "tool_usage.image_gen fills a missing image input counter") + require.Equal(t, 50, turns[0].Usage.ImageOutputTokens, "tool_usage.image_gen fills a missing image output counter") + require.Equal(t, 20, turns[1].Usage.ImageInputTokens, "response usage takes precedence over the hosted tool breakdown") + require.Equal(t, 10, turns[1].Usage.ImageOutputTokens, "response usage takes precedence over the hosted tool breakdown") + require.Equal(t, 7, turns[2].Usage.ImageInputTokens) + require.Zero(t, turns[3].Usage.ImageInputTokens, "a text turn must not reuse earlier image usage") + require.Zero(t, turns[3].Usage.ImageOutputTokens, "a text turn must not reuse earlier image usage") + + require.Equal(t, 107, result.Usage.ImageInputTokens) + require.Equal(t, 60, result.Usage.ImageOutputTokens) + require.Equal(t, 145, result.Usage.InputTokens) + require.Equal(t, 81, result.Usage.OutputTokens) +} diff --git a/backend/internal/service/openai_ws_v2_passthrough_adapter.go b/backend/internal/service/openai_ws_v2_passthrough_adapter.go index 556ad06274..5b5f25e37f 100644 --- a/backend/internal/service/openai_ws_v2_passthrough_adapter.go +++ b/backend/internal/service/openai_ws_v2_passthrough_adapter.go @@ -1207,6 +1207,7 @@ func (s *OpenAIGatewayService) proxyResponsesWebSocketV2Passthrough( RequestID: turn.RequestID, Usage: OpenAIUsage{ InputTokens: turn.Usage.InputTokens, + ImageInputTokens: turn.Usage.ImageInputTokens, OutputTokens: turn.Usage.OutputTokens, CacheCreationInputTokens: turn.Usage.CacheCreationInputTokens, CacheReadInputTokens: turn.Usage.CacheReadInputTokens, @@ -1348,6 +1349,7 @@ func (s *OpenAIGatewayService) proxyResponsesWebSocketV2Passthrough( RequestID: relayResult.RequestID, Usage: OpenAIUsage{ InputTokens: relayResult.Usage.InputTokens, + ImageInputTokens: relayResult.Usage.ImageInputTokens, OutputTokens: relayResult.Usage.OutputTokens, CacheCreationInputTokens: relayResult.Usage.CacheCreationInputTokens, CacheReadInputTokens: relayResult.Usage.CacheReadInputTokens, diff --git a/backend/internal/service/openai_ws_v2_passthrough_image_usage_test.go b/backend/internal/service/openai_ws_v2_passthrough_image_usage_test.go new file mode 100644 index 0000000000..06d66df62d --- /dev/null +++ b/backend/internal/service/openai_ws_v2_passthrough_image_usage_test.go @@ -0,0 +1,103 @@ +package service + +import ( + "context" + "net/http" + "net/http/httptest" + "strings" + "testing" + "time" + + "github.com/Wei-Shaw/sub2api/internal/config" + coderws "github.com/coder/websocket" + "github.com/gin-gonic/gin" + "github.com/stretchr/testify/require" +) + +// The billed turn result of a WS passthrough image generation must carry the +// image input tokens and the tool_usage.image_gen backfill, as HTTP does. +func TestOpenAIWSPassthroughTurnResultCarriesImageUsage(t *testing.T) { + gin.SetMode(gin.TestMode) + + cfg := &config.Config{} + cfg.Security.URLAllowlist.Enabled = false + cfg.Security.URLAllowlist.AllowInsecureHTTP = true + cfg.Gateway.OpenAIWS.Enabled = true + cfg.Gateway.OpenAIWS.APIKeyEnabled = true + cfg.Gateway.OpenAIWS.ResponsesWebsocketsV2 = true + cfg.Gateway.OpenAIWS.ModeRouterV2Enabled = true + cfg.Gateway.OpenAIWS.IngressModeDefault = OpenAIWSIngressModeCtxPool + cfg.Gateway.OpenAIWS.DialTimeoutSeconds = 3 + cfg.Gateway.OpenAIWS.ReadTimeoutSeconds = 3 + cfg.Gateway.OpenAIWS.WriteTimeoutSeconds = 3 + + upstreamConn := &openAIWSCaptureConn{ + events: [][]byte{ + []byte(`{"type":"response.completed","response":{"id":"resp_ws_image","model":"gpt-5.4","output":[{"id":"ig_ws_1","type":"image_generation_call","status":"completed","result":"final-image"}],"usage":{"input_tokens":120,"output_tokens":70,"input_tokens_details":{"image_tokens":40}},"tool_usage":{"image_gen":{"input_tokens_details":{"image_tokens":900},"output_tokens_details":{"image_tokens":55}}}}}`), + }, + } + svc := &OpenAIGatewayService{ + cfg: cfg, + httpUpstream: &httpUpstreamRecorder{}, + cache: &stubGatewayCache{}, + openaiWSResolver: NewOpenAIWSProtocolResolver(cfg), + toolCorrector: NewCodexToolCorrector(), + openaiWSPassthroughDialer: &openAIWSCaptureDialer{conn: upstreamConn}, + } + account := &Account{ + ID: 453, Name: "openai-ws-image-usage", Platform: PlatformOpenAI, Type: AccountTypeAPIKey, + Status: StatusActive, Schedulable: true, Concurrency: 1, + Credentials: map[string]any{"api_key": "sk-test"}, + Extra: map[string]any{"openai_apikey_responses_websockets_v2_mode": OpenAIWSIngressModePassthrough}, + } + + resultCh := make(chan *OpenAIForwardResult, 1) + hooks := &OpenAIWSIngressHooks{ + AfterTurn: func(_ int, result *OpenAIForwardResult, turnErr error) { + if turnErr == nil && result != nil { + resultCh <- result + } + }, + } + wsServer := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + conn, err := coderws.Accept(w, r, nil) + if err != nil { + return + } + defer func() { _ = conn.CloseNow() }() + ginCtx, _ := gin.CreateTestContext(httptest.NewRecorder()) + ginCtx.Request = r.Clone(r.Context()) + readCtx, cancel := context.WithTimeout(r.Context(), 3*time.Second) + _, firstMessage, readErr := conn.Read(readCtx) + cancel() + if readErr != nil { + return + } + _ = svc.ProxyResponsesWebSocketFromClient(r.Context(), ginCtx, conn, account, "sk-test", firstMessage, hooks) + })) + defer wsServer.Close() + + dialCtx, cancelDial := context.WithTimeout(context.Background(), 3*time.Second) + clientConn, _, err := coderws.Dial(dialCtx, "ws"+strings.TrimPrefix(wsServer.URL, "http"), nil) + cancelDial() + require.NoError(t, err) + defer func() { _ = clientConn.CloseNow() }() + + writeCtx, cancelWrite := context.WithTimeout(context.Background(), 3*time.Second) + require.NoError(t, clientConn.Write(writeCtx, coderws.MessageText, []byte(`{"type":"response.create","model":"gpt-5.4","tools":[{"type":"image_generation"}],"input":"draw"}`))) + cancelWrite() + readCtx, cancelRead := context.WithTimeout(context.Background(), 3*time.Second) + _, _, readErr := clientConn.Read(readCtx) + cancelRead() + require.NoError(t, readErr) + _ = clientConn.Close(coderws.StatusNormalClosure, "done") + + select { + case result := <-resultCh: + require.Equal(t, 120, result.Usage.InputTokens) + require.Equal(t, 40, result.Usage.ImageInputTokens, "usage image input wins over the hosted tool breakdown") + require.Equal(t, 55, result.Usage.ImageOutputTokens, "tool_usage.image_gen fills the missing image output counter") + case <-time.After(5 * time.Second): + t.Fatal("passthrough turn result was not reported") + } +}