From 26b5e8745782abe1fa41a2ebfc0b95e7e7d7e2b3 Mon Sep 17 00:00:00 2001 From: docooler Date: Fri, 17 Jul 2026 03:35:26 +0000 Subject: [PATCH 01/21] fix(apicompat): emit content_part events and full output in Responses stream MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The Anthropic→Responses streaming converter omits two things the OpenAI Responses wire format requires, which breaks clients that reconstruct the response from the event stream (rather than just reading deltas). 1. response.content_part.added is never emitted. A message item is opened with content: [], and the OpenAI SDK's accumulating stream helper (client.responses.stream) only appends a content part when it sees content_part.added. Without it, the next output_text.delta indexes output.content[content_index] and raises: File "openai/lib/streaming/responses/_responses.py", line 352, in accumulate_event content = output.content[event.content_index] IndexError: list index out of range Raw iteration (responses.create(stream=True)) does not accumulate and is unaffected, which is why this went unnoticed. 2. response.completed carries Output: []ResponsesOutput{}. get_final_response() and tracing integrations parse the terminal event's response directly, so callers see an empty output_text even though the deltas streamed correctly. This one is invisible when only watching the stream render. Also carries the full text on output_text.done / content_part.done (deltas carry increments only, done events carry the whole part) and fills in content/arguments/summary on output_item.done, which had the same empty-payload issue. Reproduced against a live Anthropic-platform group with the openai Python SDK 2.46.0; Arize Phoenix's playground hits the same path. Verified before (IndexError) and after (full text via get_final_response()). Adds regression tests covering event ordering, done-event payloads, and the terminal event's output for both text and tool calls. --- .../anthropic_to_responses_response.go | 95 ++++++++- .../anthropic_to_responses_stream_test.go | 187 ++++++++++++++++++ 2 files changed, 273 insertions(+), 9 deletions(-) create mode 100644 backend/internal/pkg/apicompat/anthropic_to_responses_stream_test.go diff --git a/backend/internal/pkg/apicompat/anthropic_to_responses_response.go b/backend/internal/pkg/apicompat/anthropic_to_responses_response.go index 9ca20e1cd..b22215836 100644 --- a/backend/internal/pkg/apicompat/anthropic_to_responses_response.go +++ b/backend/internal/pkg/apicompat/anthropic_to_responses_response.go @@ -152,11 +152,25 @@ type AnthropicEventToResponsesState struct { // For message output: accumulate text parts ContentIndex int + // TextAccum accumulates the current text part so that output_text.done and + // content_part.done can carry the full text (deltas carry increments only). + TextAccum string // For function_call: track per-output info CurrentCallID string CurrentName string + // Content of the currently open item, folded into Outputs when it closes. + CurrentContent []ResponsesContentPart // message + CurrentArgs string // function_call + CurrentSummary string // reasoning + + // Outputs accumulates every closed output item so that response.completed + // can carry the full output list. The OpenAI SDK's get_final_response() + // parses the terminal event's response directly; without this, clients see + // an empty output_text. + Outputs []ResponsesOutput + // Usage from message_start / message_delta. InputTokens here follows // Anthropic semantics (excludes cached tokens); they are added back when // emitting the OpenAI Responses usage. @@ -293,6 +307,22 @@ func anthToResHandleContentBlockStart(evt *AnthropicStreamEvent, state *Anthropi })) } + // response.content_part.added must precede the output_text.delta events + // for that part. The message item is added with content: [], and the + // OpenAI SDK's accumulating stream helper (client.responses.stream) only + // appends a content part when it sees content_part.added. Without it the + // following output_text.delta indexes output.content[content_index] and + // raises IndexError. Raw event iteration + // (responses.create(stream=True)) does not accumulate, which is why this + // went unnoticed. + events = append(events, makeResponsesEvent(state, "response.content_part.added", &ResponsesStreamEvent{ + OutputIndex: state.OutputIndex, + ContentIndex: state.ContentIndex, + ItemID: state.CurrentItemID, + Part: &ResponsesContentPart{Type: "output_text", Text: ""}, + })) + state.TextAccum = "" + case "tool_use": // Close previous item if any events = append(events, closeCurrentResponsesItem(state)...) @@ -327,6 +357,7 @@ func anthToResHandleContentBlockDelta(evt *AnthropicStreamEvent, state *Anthropi if evt.Delta.Text == "" { return nil } + state.TextAccum += evt.Delta.Text return []ResponsesStreamEvent{makeResponsesEvent(state, "response.output_text.delta", &ResponsesStreamEvent{ OutputIndex: state.OutputIndex, ContentIndex: state.ContentIndex, @@ -338,6 +369,7 @@ func anthToResHandleContentBlockDelta(evt *AnthropicStreamEvent, state *Anthropi if evt.Delta.Thinking == "" { return nil } + state.CurrentSummary += evt.Delta.Thinking return []ResponsesStreamEvent{makeResponsesEvent(state, "response.reasoning_summary_text.delta", &ResponsesStreamEvent{ OutputIndex: state.OutputIndex, SummaryIndex: 0, @@ -349,6 +381,7 @@ func anthToResHandleContentBlockDelta(evt *AnthropicStreamEvent, state *Anthropi if evt.Delta.PartialJSON == "" { return nil } + state.CurrentArgs += evt.Delta.PartialJSON return []ResponsesStreamEvent{makeResponsesEvent(state, "response.function_call_arguments.delta", &ResponsesStreamEvent{ OutputIndex: state.OutputIndex, Delta: evt.Delta.PartialJSON, @@ -393,12 +426,24 @@ func anthToResHandleContentBlockStop(evt *AnthropicStreamEvent, state *Anthropic return events case "message": - // Emit output_text.done (text block is done, but message item stays open for potential more blocks) + // Text block is done: emit output_text.done then content_part.done (the + // order OpenAI uses), both carrying the part's full text. The message + // item itself stays open since more blocks may follow. + text := state.TextAccum + state.TextAccum = "" + state.CurrentContent = append(state.CurrentContent, ResponsesContentPart{Type: "output_text", Text: text}) return []ResponsesStreamEvent{ makeResponsesEvent(state, "response.output_text.done", &ResponsesStreamEvent{ OutputIndex: state.OutputIndex, ContentIndex: state.ContentIndex, ItemID: state.CurrentItemID, + Text: text, + }), + makeResponsesEvent(state, "response.content_part.done", &ResponsesStreamEvent{ + OutputIndex: state.OutputIndex, + ContentIndex: state.ContentIndex, + ItemID: state.CurrentItemID, + Part: &ResponsesContentPart{Type: "output_text", Text: text}, }), } } @@ -454,24 +499,48 @@ func closeCurrentResponsesItem(state *AnthropicEventToResponsesState) []Response return nil } - itemType := state.CurrentItemType - itemID := state.CurrentItemID + // Assemble the full item: both output_item.done and response.completed must + // carry its content. Emitting only {type,id,status} makes SDK-side + // accumulation produce an empty output. + item := ResponsesOutput{ + Type: state.CurrentItemType, + ID: state.CurrentItemID, + Status: "completed", + } + switch state.CurrentItemType { + case "message": + item.Role = "assistant" + item.Content = state.CurrentContent + case "function_call": + item.CallID = state.CurrentCallID + item.Name = state.CurrentName + args := state.CurrentArgs + if args == "" { + args = "{}" + } + item.Arguments = args + case "reasoning": + if state.CurrentSummary != "" { + item.Summary = []ResponsesSummary{{Type: "summary_text", Text: state.CurrentSummary}} + } + } + state.Outputs = append(state.Outputs, item) // Reset state.CurrentItemType = "" state.CurrentItemID = "" state.CurrentCallID = "" state.CurrentName = "" + state.CurrentContent = nil + state.CurrentArgs = "" + state.CurrentSummary = "" + state.TextAccum = "" state.OutputIndex++ state.ContentIndex = 0 return []ResponsesStreamEvent{makeResponsesEvent(state, "response.output_item.done", &ResponsesStreamEvent{ OutputIndex: state.OutputIndex - 1, // Use the index before increment - Item: &ResponsesOutput{ - Type: itemType, - ID: itemID, - Status: "completed", - }, + Item: &item, })} } @@ -519,6 +588,14 @@ func makeResponsesCompletedEvent( eventType = "response.incomplete" } + // Carry the output items accumulated over the stream. The SDK's + // get_final_response() reads them straight from the terminal event, so an + // empty list leaves clients with an empty result. + outputs := state.Outputs + if outputs == nil { + outputs = []ResponsesOutput{} + } + return ResponsesStreamEvent{ Type: eventType, SequenceNumber: seq, @@ -527,7 +604,7 @@ func makeResponsesCompletedEvent( Object: "response", Model: state.Model, Status: status, - Output: []ResponsesOutput{}, + Output: outputs, Usage: usage, IncompleteDetails: incompleteDetails, }, diff --git a/backend/internal/pkg/apicompat/anthropic_to_responses_stream_test.go b/backend/internal/pkg/apicompat/anthropic_to_responses_stream_test.go new file mode 100644 index 000000000..c26e28982 --- /dev/null +++ b/backend/internal/pkg/apicompat/anthropic_to_responses_stream_test.go @@ -0,0 +1,187 @@ +package apicompat + +import "testing" + +// TestAnthropicEventToResponses_TextEmitsContentPart pins that a message text +// stream emits response.content_part.added, and that it precedes the first +// output_text.delta for that part. +// +// Why: the OpenAI SDK's accumulating stream helper (client.responses.stream) +// only appends a content part to the message item when it sees +// content_part.added. The item is added with content: [], so a missing event +// makes the following output_text.delta index output.content[content_index] and +// raise IndexError. Raw event iteration does not accumulate, so a regression +// here is easy to miss. +func TestAnthropicEventToResponses_TextEmitsContentPart(t *testing.T) { + state := NewAnthropicEventToResponsesState() + state.Model = "claude-sonnet-4-5" + + var types []string + feed := func(evt *AnthropicStreamEvent) { + for _, out := range AnthropicEventToResponsesEvents(evt, state) { + types = append(types, out.Type) + } + } + + idx := 0 + feed(&AnthropicStreamEvent{Type: "message_start", Message: &AnthropicResponse{ID: "msg_1", Model: "claude-sonnet-4-5"}}) + feed(&AnthropicStreamEvent{Type: "content_block_start", Index: &idx, ContentBlock: &AnthropicContentBlock{Type: "text"}}) + feed(&AnthropicStreamEvent{Type: "content_block_delta", Index: &idx, Delta: &AnthropicDelta{Type: "text_delta", Text: "Hel"}}) + feed(&AnthropicStreamEvent{Type: "content_block_delta", Index: &idx, Delta: &AnthropicDelta{Type: "text_delta", Text: "lo"}}) + feed(&AnthropicStreamEvent{Type: "content_block_stop", Index: &idx}) + feed(&AnthropicStreamEvent{Type: "message_stop"}) + + posOf := func(target string) int { + for i, ty := range types { + if ty == target { + return i + } + } + return -1 + } + + partAdded := posOf("response.content_part.added") + firstDelta := posOf("response.output_text.delta") + + if partAdded < 0 { + t.Fatalf("response.content_part.added was not emitted; got %v", types) + } + if firstDelta < 0 { + t.Fatalf("response.output_text.delta was not emitted; got %v", types) + } + if partAdded > firstDelta { + t.Errorf("content_part.added must precede the first output_text.delta; got %v", types) + } + if posOf("response.content_part.done") < 0 { + t.Errorf("response.content_part.done was not emitted; got %v", types) + } +} + +// TestAnthropicEventToResponses_DoneEventsCarryFullText pins that done events +// carry the part's full text (deltas carry increments only). +func TestAnthropicEventToResponses_DoneEventsCarryFullText(t *testing.T) { + state := NewAnthropicEventToResponsesState() + state.Model = "claude-sonnet-4-5" + + var events []ResponsesStreamEvent + feed := func(evt *AnthropicStreamEvent) { + events = append(events, AnthropicEventToResponsesEvents(evt, state)...) + } + + idx := 0 + feed(&AnthropicStreamEvent{Type: "message_start", Message: &AnthropicResponse{ID: "msg_1"}}) + feed(&AnthropicStreamEvent{Type: "content_block_start", Index: &idx, ContentBlock: &AnthropicContentBlock{Type: "text"}}) + feed(&AnthropicStreamEvent{Type: "content_block_delta", Index: &idx, Delta: &AnthropicDelta{Type: "text_delta", Text: "Hello "}}) + feed(&AnthropicStreamEvent{Type: "content_block_delta", Index: &idx, Delta: &AnthropicDelta{Type: "text_delta", Text: "world"}}) + feed(&AnthropicStreamEvent{Type: "content_block_stop", Index: &idx}) + + const want = "Hello world" + var sawTextDone, sawPartDone bool + for _, e := range events { + switch e.Type { + case "response.output_text.done": + sawTextDone = true + if e.Text != want { + t.Errorf("output_text.done text = %q, want %q", e.Text, want) + } + case "response.content_part.done": + sawPartDone = true + if e.Part == nil || e.Part.Text != want { + t.Errorf("content_part.done part = %+v, want text %q", e.Part, want) + } + } + } + if !sawTextDone || !sawPartDone { + t.Errorf("missing done events: output_text.done=%v content_part.done=%v", sawTextDone, sawPartDone) + } +} + +// TestAnthropicEventToResponses_CompletedCarriesOutput pins that +// response.completed carries the full output list. The SDK's +// get_final_response() and tracing integrations parse the terminal event's +// response directly; an empty output leaves them with nothing (the text still +// renders from deltas, which is why this is invisible when only watching the +// stream). +func TestAnthropicEventToResponses_CompletedCarriesOutput(t *testing.T) { + state := NewAnthropicEventToResponsesState() + state.Model = "claude-sonnet-4-5" + + var events []ResponsesStreamEvent + feed := func(evt *AnthropicStreamEvent) { + events = append(events, AnthropicEventToResponsesEvents(evt, state)...) + } + + idx := 0 + feed(&AnthropicStreamEvent{Type: "message_start", Message: &AnthropicResponse{ID: "msg_1"}}) + feed(&AnthropicStreamEvent{Type: "content_block_start", Index: &idx, ContentBlock: &AnthropicContentBlock{Type: "text"}}) + feed(&AnthropicStreamEvent{Type: "content_block_delta", Index: &idx, Delta: &AnthropicDelta{Type: "text_delta", Text: "4826"}}) + feed(&AnthropicStreamEvent{Type: "content_block_stop", Index: &idx}) + feed(&AnthropicStreamEvent{Type: "message_stop"}) + + var completed *ResponsesStreamEvent + for i := range events { + if events[i].Type == "response.completed" { + completed = &events[i] + } + } + if completed == nil || completed.Response == nil { + t.Fatalf("response.completed was not emitted") + } + if len(completed.Response.Output) == 0 { + t.Fatalf("response.completed carries an empty output; clients would see no result") + } + msg := completed.Response.Output[0] + if msg.Type != "message" || len(msg.Content) == 0 { + t.Fatalf("output[0] = %+v, want a message with content", msg) + } + if msg.Content[0].Text != "4826" { + t.Errorf("output[0].content[0].text = %q, want %q", msg.Content[0].Text, "4826") + } +} + +// TestAnthropicEventToResponses_ToolCallCompletedCarriesArguments pins that a +// function call's accumulated arguments survive into output_item.done and +// response.completed. +func TestAnthropicEventToResponses_ToolCallCompletedCarriesArguments(t *testing.T) { + state := NewAnthropicEventToResponsesState() + state.Model = "claude-sonnet-4-5" + + var events []ResponsesStreamEvent + feed := func(evt *AnthropicStreamEvent) { + events = append(events, AnthropicEventToResponsesEvents(evt, state)...) + } + + idx := 0 + feed(&AnthropicStreamEvent{Type: "message_start", Message: &AnthropicResponse{ID: "msg_1"}}) + feed(&AnthropicStreamEvent{Type: "content_block_start", Index: &idx, ContentBlock: &AnthropicContentBlock{ + Type: "tool_use", ID: "toolu_1", Name: "get_weather", + }}) + feed(&AnthropicStreamEvent{Type: "content_block_delta", Index: &idx, Delta: &AnthropicDelta{ + Type: "input_json_delta", PartialJSON: `{"city":`, + }}) + feed(&AnthropicStreamEvent{Type: "content_block_delta", Index: &idx, Delta: &AnthropicDelta{ + Type: "input_json_delta", PartialJSON: `"SH"}`, + }}) + feed(&AnthropicStreamEvent{Type: "content_block_stop", Index: &idx}) + feed(&AnthropicStreamEvent{Type: "message_stop"}) + + var completed *ResponsesStreamEvent + for i := range events { + if events[i].Type == "response.completed" { + completed = &events[i] + } + } + if completed == nil || completed.Response == nil || len(completed.Response.Output) == 0 { + t.Fatalf("response.completed carries no output") + } + fc := completed.Response.Output[0] + if fc.Type != "function_call" { + t.Fatalf("output[0].type = %q, want function_call", fc.Type) + } + if fc.Arguments != `{"city":"SH"}` { + t.Errorf("arguments = %q, want %q", fc.Arguments, `{"city":"SH"}`) + } + if fc.Name != "get_weather" { + t.Errorf("name = %q, want get_weather", fc.Name) + } +} From 7043982ae3b950fb165d15421f43fe1a37586b93 Mon Sep 17 00:00:00 2001 From: li Date: Fri, 17 Jul 2026 16:35:43 +0800 Subject: [PATCH 02/21] =?UTF-8?q?fix(grok):=20=E7=BA=AF=E5=AE=A2=E6=88=B7?= =?UTF-8?q?=E7=AB=AF=E5=87=BD=E6=95=B0=E5=B7=A5=E5=85=B7=E4=B8=8D=E5=86=8D?= =?UTF-8?q?=E6=B3=A8=E5=85=A5=E5=8E=9F=E7=94=9F=E6=90=9C=E7=B4=A2=E5=B7=A5?= =?UTF-8?q?=E5=85=B7?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Fixes #4486 --- .../service/openai_gateway_grok_cache.go | 7 ++ .../service/openai_gateway_grok_cache_test.go | 40 +++++---- .../openai_gateway_grok_cache_tool_test.go | 87 +++++++++++++++++++ .../service/openai_gateway_grok_test.go | 2 +- 4 files changed, 120 insertions(+), 16 deletions(-) create mode 100644 backend/internal/service/openai_gateway_grok_cache_tool_test.go diff --git a/backend/internal/service/openai_gateway_grok_cache.go b/backend/internal/service/openai_gateway_grok_cache.go index 9b6436e43..d9e34cbb5 100644 --- a/backend/internal/service/openai_gateway_grok_cache.go +++ b/backend/internal/service/openai_gateway_grok_cache.go @@ -296,6 +296,13 @@ func appendMissingGrokFreeCacheNativeTools(body []byte) ([]byte, error) { if !hasFunction { return body, nil } + // Only complement missing native search tools when the request already contains + // at least one search tool (native or function-form). Pure client function tools + // (e.g. view_image) must not trigger injection to avoid biasing model tool + // selection (#4486). + if !present["web_search"] && !present["x_search"] { + return body, nil + } for _, toolType := range []string{"web_search", "x_search"} { if present[toolType] { continue diff --git a/backend/internal/service/openai_gateway_grok_cache_test.go b/backend/internal/service/openai_gateway_grok_cache_test.go index c31ee7bc3..73b1ea5ec 100644 --- a/backend/internal/service/openai_gateway_grok_cache_test.go +++ b/backend/internal/service/openai_gateway_grok_cache_test.go @@ -211,6 +211,7 @@ func TestApplyGrokCacheIdentityAppendsNativeToolsToResponseFunctions(t *testing. for _, tt := range tests { t.Run(tt.name, func(t *testing.T) { + // Pure client function tools without search → no native injection (#4486). intentBody := []byte(`{"model":"grok","tools":[{"type":"function","name":"lookup","description":"look up a value","parameters":{"type":"object"}},{"type":"function","name":"save","parameters":{"type":"object"}}]` + tt.toolChoiceJSON + `}`) body, err := applyGrokResponsesCacheIdentity(intentBody, intentBody, "isolated-id", true) require.NoError(t, err) @@ -219,28 +220,35 @@ func TestApplyGrokCacheIdentityAppendsNativeToolsToResponseFunctions(t *testing. require.NoError(t, err) require.Equal(t, "isolated-id", gjson.GetBytes(body, "prompt_cache_key").String()) tools := gjson.GetBytes(body, "tools").Array() - require.Len(t, tools, 4) + require.Len(t, tools, 2, "pure client functions should not get native search injected") require.Equal(t, "function", tools[0].Get("type").String()) require.Equal(t, "lookup", tools[0].Get("name").String()) require.Equal(t, "function", tools[1].Get("type").String()) require.Equal(t, "save", tools[1].Get("name").String()) - require.Equal(t, "web_search", tools[2].Get("type").String()) - require.Equal(t, "x_search", tools[3].Get("type").String()) require.Equal(t, tt.wantChoice, gjson.GetBytes(body, "tool_choice").Exists()) - if tt.wantChoice { - require.Equal(t, "auto", gjson.GetBytes(body, "tool_choice").String()) - } - - second, err := applyGrokResponsesCacheIdentity(body, intentBody, "isolated-id", true) - require.NoError(t, err) - second, err = applyGrokFreeMessagesFunctionToolCacheRoute(second, intentBody, account, "isolated-id") - require.NoError(t, err) - require.JSONEq(t, string(body), string(second), "native tools must not be duplicated") - require.Len(t, gjson.GetBytes(second, "tools").Array(), 4) }) } } +func TestApplyGrokCacheIdentityAppendsNativeToolsWhenSearchPresent(t *testing.T) { + account := healthyGrokOAuthGatewayTestAccount(901, "access-token") + account.Credentials["subscription_tier"] = " FREE " + + // Function tools INCLUDING web_search → convert + complement with x_search. + intentBody := []byte(`{"model":"grok","tools":[{"type":"function","name":"lookup","description":"look up a value","parameters":{"type":"object"}},{"type":"function","name":"web_search","description":"search","parameters":{"type":"object"}}]}`) + body, err := applyGrokResponsesCacheIdentity(intentBody, intentBody, "isolated-id", true) + require.NoError(t, err) + body, err = applyGrokFreeMessagesFunctionToolCacheRoute(body, intentBody, account, "isolated-id") + require.NoError(t, err) + + tools := gjson.GetBytes(body, "tools").Array() + require.Len(t, tools, 3, "lookup(function) + web_search(native) + x_search(native)") + require.Equal(t, "function", tools[0].Get("type").String()) + require.Equal(t, "lookup", tools[0].Get("name").String()) + require.Equal(t, "web_search", tools[1].Get("type").String()) + require.Equal(t, "x_search", tools[2].Get("type").String()) +} + func TestApplyGrokCacheIdentityRequiresPatchedFunctionTools(t *testing.T) { account := healthyGrokOAuthGatewayTestAccount(902, "access-token") account.Credentials["subscription_tier"] = "free" @@ -272,7 +280,9 @@ func TestApplyGrokCacheIdentityRequiresPatchedFunctionTools(t *testing.T) { } func TestGrokFreeMessagesFunctionToolCacheRouteRequiresKnownFreeTier(t *testing.T) { - intentBody := []byte(`{"model":"grok","tools":[{"type":"function","name":"lookup"}],"tool_choice":"auto"}`) + // Include web_search as function to trigger native tool injection (pure client + // functions no longer trigger injection after #4486 fix). + intentBody := []byte(`{"model":"grok","tools":[{"type":"function","name":"lookup"},{"type":"function","name":"web_search"}],"tool_choice":"auto"}`) tests := []struct { name string account *Account @@ -382,7 +392,7 @@ func TestGrokFreeMessagesFunctionToolCacheRouteRequiresKnownFreeTier(t *testing. require.Equal(t, "x_search", tools[2].Get("type").String()) return } - require.Len(t, tools, 1) + require.Len(t, tools, 2, "non-free accounts should not get native search injected") }) } } diff --git a/backend/internal/service/openai_gateway_grok_cache_tool_test.go b/backend/internal/service/openai_gateway_grok_cache_tool_test.go new file mode 100644 index 000000000..77dba26ce --- /dev/null +++ b/backend/internal/service/openai_gateway_grok_cache_tool_test.go @@ -0,0 +1,87 @@ +package service + +import ( + "testing" + + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" + "github.com/tidwall/gjson" +) + +func TestAppendMissingGrokFreeCacheNativeTools_PureClientFunctionNoInject(t *testing.T) { + body := []byte(`{ + "model": "grok-4.5", + "tools": [ + {"type":"function","name":"view_image","description":"View image","parameters":{"type":"object","properties":{"path":{"type":"string"}},"required":["path"]}} + ], + "tool_choice": "auto" + }`) + + result, err := appendMissingGrokFreeCacheNativeTools(body) + require.NoError(t, err) + + tools := gjson.GetBytes(result, "tools").Array() + for _, tool := range tools { + toolType := tool.Get("type").String() + assert.NotEqual(t, "web_search", toolType, "should not inject web_search for pure client functions") + assert.NotEqual(t, "x_search", toolType, "should not inject x_search for pure client functions") + } +} + +func TestAppendMissingGrokFreeCacheNativeTools_FunctionPlusWebSearchInjects(t *testing.T) { + body := []byte(`{ + "model": "grok-4.5", + "tools": [ + {"type":"function","name":"view_image","description":"View","parameters":{"type":"object","properties":{"path":{"type":"string"}},"required":["path"]}}, + {"type":"function","name":"web_search","description":"Search","parameters":{"type":"object","properties":{"query":{"type":"string"}},"required":["query"]}} + ] + }`) + + result, err := appendMissingGrokFreeCacheNativeTools(body) + require.NoError(t, err) + + tools := gjson.GetBytes(result, "tools").Array() + types := make(map[string]bool) + for _, tool := range tools { + types[tool.Get("type").String()] = true + } + assert.True(t, types["web_search"], "web_search should be present (converted from function)") + assert.True(t, types["x_search"], "x_search should be injected when web_search is present alongside client functions") +} + +func TestAppendMissingGrokFreeCacheNativeTools_NativeSearchAlreadyPresent(t *testing.T) { + body := []byte(`{ + "model": "grok-4.5", + "tools": [ + {"type":"function","name":"view_image","description":"View","parameters":{"type":"object","properties":{"path":{"type":"string"}},"required":["path"]}}, + {"type":"web_search"} + ] + }`) + + result, err := appendMissingGrokFreeCacheNativeTools(body) + require.NoError(t, err) + + tools := gjson.GetBytes(result, "tools").Array() + types := make(map[string]bool) + for _, tool := range tools { + types[tool.Get("type").String()] = true + } + assert.True(t, types["web_search"]) + assert.True(t, types["x_search"], "x_search should be injected when web_search is already present") +} + +func TestAppendMissingGrokFreeCacheNativeTools_MultipleFunctionsNoSearch(t *testing.T) { + body := []byte(`{ + "model": "grok-4.5", + "tools": [ + {"type":"function","name":"view_image","description":"View","parameters":{"type":"object","properties":{"path":{"type":"string"}},"required":["path"]}}, + {"type":"function","name":"read_file","description":"Read","parameters":{"type":"object","properties":{"path":{"type":"string"}},"required":["path"]}} + ] + }`) + + result, err := appendMissingGrokFreeCacheNativeTools(body) + require.NoError(t, err) + + tools := gjson.GetBytes(result, "tools").Array() + require.Len(t, tools, 2, "no tools should be injected for pure client functions") +} diff --git a/backend/internal/service/openai_gateway_grok_test.go b/backend/internal/service/openai_gateway_grok_test.go index 67ad71214..b8ca3a217 100644 --- a/backend/internal/service/openai_gateway_grok_test.go +++ b/backend/internal/service/openai_gateway_grok_test.go @@ -1604,7 +1604,7 @@ func TestForwardAsAnthropicForGrokFunctionToolUsesCacheCapableMixedRoute(t *test body := []byte(`{ "model":"grok","max_tokens":32,"stream":false, "messages":[{"role":"user","content":"look up alpha"}], - "tools":[{"name":"lookup","description":"look up a key","input_schema":{"type":"object","properties":{"key":{"type":"string"}},"required":["key"]}}], + "tools":[{"name":"lookup","description":"look up a key","input_schema":{"type":"object","properties":{"key":{"type":"string"}},"required":["key"]}},{"name":"web_search","description":"search the web","input_schema":{"type":"object","properties":{"query":{"type":"string"}},"required":["query"]}}], "tool_choice":{"type":"auto"} }`) c.Request = httptest.NewRequest(http.MethodPost, "/v1/messages", bytes.NewReader(body)) From 1f7b7b91322e42abaa2a2ccb6ef7e0f3d3a1ed1a Mon Sep 17 00:00:00 2001 From: haruka <1628615876@qq.com> Date: Fri, 17 Jul 2026 18:12:33 +0800 Subject: [PATCH 03/21] fix: honor temp unschedulable rules in pool mode --- backend/internal/service/error_policy_test.go | 112 +++++++++++++++++- ...nai_account_runtime_block_fastpath_test.go | 52 ++++++++ .../internal/service/openai_alpha_search.go | 10 +- backend/internal/service/openai_embeddings.go | 4 +- .../service/openai_gateway_cc_pipeline.go | 5 +- .../service/openai_gateway_forward.go | 4 +- .../service/openai_gateway_passthrough.go | 4 +- .../service/openai_gateway_upstream_errors.go | 11 +- backend/internal/service/openai_images.go | 4 +- .../service/openai_images_responses.go | 10 +- backend/internal/service/ratelimit_service.go | 11 +- 11 files changed, 200 insertions(+), 27 deletions(-) diff --git a/backend/internal/service/error_policy_test.go b/backend/internal/service/error_policy_test.go index 2aa7a421d..178cec925 100644 --- a/backend/internal/service/error_policy_test.go +++ b/backend/internal/service/error_policy_test.go @@ -209,6 +209,50 @@ func TestCheckErrorPolicy(t *testing.T) { body: []byte(`unauthorized`), expected: ErrorPolicySkipped, }, + { + name: "pool_mode_temp_unschedulable_hit_returns_temp_unscheduled", + account: &Account{ + ID: 9, + Type: AccountTypeAPIKey, + Platform: PlatformOpenAI, + Credentials: map[string]any{ + "pool_mode": true, + "temp_unschedulable_enabled": true, + "temp_unschedulable_rules": []any{ + map[string]any{ + "error_code": float64(http.StatusServiceUnavailable), + "keywords": []any{"unavailable"}, + "duration_minutes": float64(30), + }, + }, + }, + }, + statusCode: http.StatusServiceUnavailable, + body: []byte(`Service temporarily unavailable`), + expected: ErrorPolicyTempUnscheduled, + }, + { + name: "pool_mode_temp_unschedulable_miss_returns_skipped", + account: &Account{ + ID: 10, + Type: AccountTypeAPIKey, + Platform: PlatformOpenAI, + Credentials: map[string]any{ + "pool_mode": true, + "temp_unschedulable_enabled": true, + "temp_unschedulable_rules": []any{ + map[string]any{ + "error_code": float64(http.StatusServiceUnavailable), + "keywords": []any{"maintenance"}, + "duration_minutes": float64(30), + }, + }, + }, + }, + statusCode: http.StatusServiceUnavailable, + body: []byte(`Service temporarily unavailable`), + expected: ErrorPolicySkipped, + }, } for _, tt := range tests { @@ -222,7 +266,7 @@ func TestCheckErrorPolicy(t *testing.T) { } } -func TestHandleUpstreamError_PoolModeCustomErrorCodesOverride(t *testing.T) { +func TestHandleUpstreamError_PoolModePolicies(t *testing.T) { t.Run("pool_mode_without_custom_error_codes_still_skips", func(t *testing.T) { repo := &errorPolicyRepoStub{} svc := NewRateLimitService(repo, nil, &config.Config{}, nil, nil) @@ -262,6 +306,72 @@ func TestHandleUpstreamError_PoolModeCustomErrorCodesOverride(t *testing.T) { require.Equal(t, 1, repo.setErrCalls) require.Equal(t, 0, repo.tempCalls) }) + + t.Run("pool_mode_explicit_temp_rule_stops_scheduling", func(t *testing.T) { + repo := &errorPolicyRepoStub{} + svc := NewRateLimitService(repo, nil, &config.Config{}, nil, nil) + account := &Account{ + ID: 32, + Type: AccountTypeAPIKey, + Platform: PlatformOpenAI, + Credentials: map[string]any{ + "pool_mode": true, + "temp_unschedulable_enabled": true, + "temp_unschedulable_rules": []any{ + map[string]any{ + "error_code": float64(http.StatusServiceUnavailable), + "keywords": []any{"unavailable"}, + "duration_minutes": float64(30), + }, + }, + }, + } + + shouldDisable := svc.HandleUpstreamError( + context.Background(), + account, + http.StatusServiceUnavailable, + http.Header{}, + []byte("Service temporarily unavailable"), + ) + + require.True(t, shouldDisable) + require.Equal(t, 0, repo.setErrCalls) + require.Equal(t, 1, repo.tempCalls) + }) + + t.Run("pool_mode_temp_rule_miss_still_skips", func(t *testing.T) { + repo := &errorPolicyRepoStub{} + svc := NewRateLimitService(repo, nil, &config.Config{}, nil, nil) + account := &Account{ + ID: 33, + Type: AccountTypeAPIKey, + Platform: PlatformOpenAI, + Credentials: map[string]any{ + "pool_mode": true, + "temp_unschedulable_enabled": true, + "temp_unschedulable_rules": []any{ + map[string]any{ + "error_code": float64(http.StatusServiceUnavailable), + "keywords": []any{"maintenance"}, + "duration_minutes": float64(30), + }, + }, + }, + } + + shouldDisable := svc.HandleUpstreamError( + context.Background(), + account, + http.StatusServiceUnavailable, + http.Header{}, + []byte("Service temporarily unavailable"), + ) + + require.False(t, shouldDisable) + require.Equal(t, 0, repo.setErrCalls) + require.Equal(t, 0, repo.tempCalls) + }) } // --------------------------------------------------------------------------- diff --git a/backend/internal/service/openai_account_runtime_block_fastpath_test.go b/backend/internal/service/openai_account_runtime_block_fastpath_test.go index 7d7cfbda5..6a1758419 100644 --- a/backend/internal/service/openai_account_runtime_block_fastpath_test.go +++ b/backend/internal/service/openai_account_runtime_block_fastpath_test.go @@ -83,6 +83,58 @@ func TestOpenAIRuntimeBlocker_IgnoresNonOpenAIFromRateLimitService(t *testing.T) require.False(t, gateway.isOpenAIAccountRuntimeBlocked(account)) } +func TestOpenAIPoolModeTempRule_StopsSameAccountRetryAndBlocksAcrossModels(t *testing.T) { + repo := &errorPolicyRepoStub{} + rateLimitService := NewRateLimitService(repo, nil, &config.Config{}, nil, nil) + gateway := &OpenAIGatewayService{ + cfg: &config.Config{}, + rateLimitService: rateLimitService, + } + rateLimitService.SetAccountRuntimeBlocker(gateway) + account := &Account{ + ID: 46, + Platform: PlatformOpenAI, + Type: AccountTypeAPIKey, + Status: StatusActive, + Schedulable: true, + Credentials: map[string]any{ + "pool_mode": true, + "pool_mode_retry_status_codes": []any{float64(http.StatusServiceUnavailable)}, + "temp_unschedulable_enabled": true, + "temp_unschedulable_rules": []any{ + map[string]any{ + "error_code": float64(http.StatusServiceUnavailable), + "keywords": []any{"unavailable"}, + "duration_minutes": float64(30), + }, + }, + }, + } + body := []byte(`{"error":{"message":"Service temporarily unavailable"}}`) + resp := &http.Response{ + StatusCode: http.StatusServiceUnavailable, + Header: http.Header{}, + } + + failoverErr := gateway.failoverOpenAIUpstreamHTTPError( + context.Background(), + nil, + account, + resp, + body, + "Service temporarily unavailable", + "gpt-5.4", + ) + + require.NotNil(t, failoverErr) + require.False(t, failoverErr.RetryableOnSameAccount) + require.Equal(t, 1, repo.tempCalls) + require.Equal(t, 0, repo.setErrCalls) + require.Equal(t, StatusActive, account.Status) + require.True(t, gateway.isOpenAIAccountRequestRuntimeBlocked(account, "gpt-5.4")) + require.True(t, gateway.isOpenAIAccountRequestRuntimeBlocked(account, "gpt-5.5")) +} + func TestOpenAIModelNotFound_DoesNotRuntimeBlockWholeAccount(t *testing.T) { repo := &modelNotFoundAccountRepoStub{} svc := &OpenAIGatewayService{ diff --git a/backend/internal/service/openai_alpha_search.go b/backend/internal/service/openai_alpha_search.go index f11a21958..488db038e 100644 --- a/backend/internal/service/openai_alpha_search.go +++ b/backend/internal/service/openai_alpha_search.go @@ -95,13 +95,14 @@ func (s *OpenAIGatewayService) ForwardAlphaSearch(ctx context.Context, c *gin.Co // 被永久标记为 error；历史导入且缺少 auth_mode 标记的 at- token 也会 // 漏过 PAT 类型判断。这里仍允许本次请求换号，但不修改任何账号状态； // 真正的凭据失效由普通 Responses 请求或 whoami 校验判定。 + shouldDisable := false if shouldApplyOpenAIAlphaSearchAccountErrorSideEffects(resp.StatusCode) { - s.handleFailoverSideEffects(ctx, resp, account, respBody, upstreamModel) + shouldDisable = s.handleFailoverSideEffects(ctx, resp, account, respBody, upstreamModel) } return nil, &UpstreamFailoverError{ StatusCode: resp.StatusCode, ResponseBody: respBody, - RetryableOnSameAccount: account.IsPoolMode() && account.IsPoolModeRetryableStatus(resp.StatusCode), + RetryableOnSameAccount: !shouldDisable && account.IsPoolMode() && account.IsPoolModeRetryableStatus(resp.StatusCode), } } } @@ -169,13 +170,14 @@ func (s *OpenAIGatewayService) forwardAlphaSearchViaResponsesWebSearch( if s.shouldFailoverOpenAIUpstreamResponse(resp.StatusCode, upstreamMessage, respBody) { resp.Body = io.NopCloser(bytes.NewReader(respBody)) // 仍按 alpha/search 工具请求处理：PAT 的工具链路失败不能直接永久置错。 + shouldDisable := false if shouldApplyOpenAIAlphaSearchAccountErrorSideEffects(resp.StatusCode) { - s.handleFailoverSideEffects(ctx, resp, account, respBody, upstreamModel) + shouldDisable = s.handleFailoverSideEffects(ctx, resp, account, respBody, upstreamModel) } return nil, &UpstreamFailoverError{ StatusCode: resp.StatusCode, ResponseBody: respBody, - RetryableOnSameAccount: account.IsPoolMode() && account.IsPoolModeRetryableStatus(resp.StatusCode), + RetryableOnSameAccount: !shouldDisable && account.IsPoolMode() && account.IsPoolModeRetryableStatus(resp.StatusCode), } } } diff --git a/backend/internal/service/openai_embeddings.go b/backend/internal/service/openai_embeddings.go index fa821d8f3..c0c590029 100644 --- a/backend/internal/service/openai_embeddings.go +++ b/backend/internal/service/openai_embeddings.go @@ -132,11 +132,11 @@ func (s *OpenAIGatewayService) ForwardEmbeddings( Message: upstreamMsg, Detail: upstreamDetail, }) - s.handleOpenAIAccountUpstreamError(ctx, account, resp.StatusCode, resp.Header, respBody, upstreamModel) + shouldDisable := s.handleOpenAIAccountUpstreamError(ctx, account, resp.StatusCode, resp.Header, respBody, upstreamModel) return nil, &UpstreamFailoverError{ StatusCode: resp.StatusCode, ResponseBody: respBody, - RetryableOnSameAccount: account.IsPoolMode() && account.IsPoolModeRetryableStatus(resp.StatusCode), + RetryableOnSameAccount: !shouldDisable && account.IsPoolMode() && account.IsPoolModeRetryableStatus(resp.StatusCode), } } writeOpenAIEmbeddingsUpstreamResponse(c, resp, respBody, s.responseHeaderFilter) diff --git a/backend/internal/service/openai_gateway_cc_pipeline.go b/backend/internal/service/openai_gateway_cc_pipeline.go index 52658b272..533f2ad38 100644 --- a/backend/internal/service/openai_gateway_cc_pipeline.go +++ b/backend/internal/service/openai_gateway_cc_pipeline.go @@ -112,15 +112,16 @@ func (s *OpenAIGatewayService) failoverOpenAIUpstreamHTTPError( Message: upstreamMsg, Detail: upstreamDetail, }) + shouldDisable := false if account.Platform != PlatformGrok { - s.handleOpenAIAccountUpstreamError(ctx, account, resp.StatusCode, resp.Header, respBody, upstreamModel) + shouldDisable = s.handleOpenAIAccountUpstreamError(ctx, account, resp.StatusCode, resp.Header, respBody, upstreamModel) } return newOpenAIUpstreamFailoverError( resp.StatusCode, resp.Header, respBody, upstreamMsg, - account.IsPoolMode() && (account.IsPoolModeRetryableStatus(resp.StatusCode) || isOpenAITransientProcessingError(resp.StatusCode, upstreamMsg, respBody)), + !shouldDisable && account.IsPoolMode() && (account.IsPoolModeRetryableStatus(resp.StatusCode) || isOpenAITransientProcessingError(resp.StatusCode, upstreamMsg, respBody)), ) } diff --git a/backend/internal/service/openai_gateway_forward.go b/backend/internal/service/openai_gateway_forward.go index 16c034f19..93697be77 100644 --- a/backend/internal/service/openai_gateway_forward.go +++ b/backend/internal/service/openai_gateway_forward.go @@ -877,13 +877,13 @@ func (s *OpenAIGatewayService) Forward(ctx context.Context, c *gin.Context, acco Detail: upstreamDetail, }) - s.handleFailoverSideEffects(ctx, resp, account, respBody, upstreamModel) + shouldDisable := s.handleFailoverSideEffects(ctx, resp, account, respBody, upstreamModel) return nil, newOpenAIUpstreamFailoverError( resp.StatusCode, resp.Header, respBody, upstreamMsg, - account.IsPoolMode() && (account.IsPoolModeRetryableStatus(resp.StatusCode) || isOpenAITransientProcessingError(resp.StatusCode, upstreamMsg, respBody)), + !shouldDisable && account.IsPoolMode() && (account.IsPoolModeRetryableStatus(resp.StatusCode) || isOpenAITransientProcessingError(resp.StatusCode, upstreamMsg, respBody)), ) } return s.handleErrorResponse(ctx, resp, c, account, body, billingModel) diff --git a/backend/internal/service/openai_gateway_passthrough.go b/backend/internal/service/openai_gateway_passthrough.go index a17d050af..1fb281b4a 100644 --- a/backend/internal/service/openai_gateway_passthrough.go +++ b/backend/internal/service/openai_gateway_passthrough.go @@ -580,7 +580,7 @@ func (s *OpenAIGatewayService) handleFailoverErrorResponsePassthrough( logOpenAIInstructionsRequiredDebug(ctx, c, account, resp.StatusCode, upstreamMsg, requestBody, body) reqModel, _, _ := extractOpenAIRequestMetaFromBody(requestBody) canonicalModel := canonicalOpenAIAccountSchedulingModel(account, reqModel) - _ = s.handleOpenAIAccountUpstreamError(ctx, account, resp.StatusCode, resp.Header, body, canonicalModel) + shouldDisable := s.handleOpenAIAccountUpstreamError(ctx, account, resp.StatusCode, resp.Header, body, canonicalModel) appendOpsUpstreamError(c, OpsUpstreamErrorEvent{ Platform: account.Platform, AccountID: account.ID, @@ -598,7 +598,7 @@ func (s *OpenAIGatewayService) handleFailoverErrorResponsePassthrough( resp.Header, body, upstreamMsg, - account.IsPoolMode() && account.IsPoolModeRetryableStatus(resp.StatusCode), + !shouldDisable && account.IsPoolMode() && account.IsPoolModeRetryableStatus(resp.StatusCode), ) } diff --git a/backend/internal/service/openai_gateway_upstream_errors.go b/backend/internal/service/openai_gateway_upstream_errors.go index a21775308..ff97e43fc 100644 --- a/backend/internal/service/openai_gateway_upstream_errors.go +++ b/backend/internal/service/openai_gateway_upstream_errors.go @@ -305,12 +305,11 @@ func (s *OpenAIGatewayService) readUpstreamErrorBody(resp *http.Response) []byte return body } -func (s *OpenAIGatewayService) handleFailoverSideEffects(ctx context.Context, resp *http.Response, account *Account, responseBody []byte, canonicalModel ...string) { +func (s *OpenAIGatewayService) handleFailoverSideEffects(ctx context.Context, resp *http.Response, account *Account, responseBody []byte, canonicalModel ...string) bool { if len(canonicalModel) > 0 { - s.handleOpenAIAccountUpstreamError(ctx, account, resp.StatusCode, resp.Header, responseBody, canonicalModel[0]) - return + return s.handleOpenAIAccountUpstreamError(ctx, account, resp.StatusCode, resp.Header, responseBody, canonicalModel[0]) } - s.handleOpenAIAccountUpstreamError(ctx, account, resp.StatusCode, resp.Header, responseBody) + return s.handleOpenAIAccountUpstreamError(ctx, account, resp.StatusCode, resp.Header, responseBody) } func (s *OpenAIGatewayService) handleErrorResponse( @@ -470,7 +469,7 @@ func (s *OpenAIGatewayService) handleErrorResponse( return nil, &UpstreamFailoverError{ StatusCode: resp.StatusCode, ResponseBody: body, - RetryableOnSameAccount: account.IsPoolMode() && account.IsPoolModeRetryableStatus(resp.StatusCode), + RetryableOnSameAccount: false, } } @@ -640,7 +639,7 @@ func (s *OpenAIGatewayService) handleCompatErrorResponse( return nil, &UpstreamFailoverError{ StatusCode: resp.StatusCode, ResponseBody: body, - RetryableOnSameAccount: account.IsPoolMode() && account.IsPoolModeRetryableStatus(resp.StatusCode), + RetryableOnSameAccount: false, } } diff --git a/backend/internal/service/openai_images.go b/backend/internal/service/openai_images.go index fa0b7f8b6..f32fc5615 100644 --- a/backend/internal/service/openai_images.go +++ b/backend/internal/service/openai_images.go @@ -652,11 +652,11 @@ func (s *OpenAIGatewayService) forwardOpenAIImagesAPIKey( Kind: "failover", Message: upstreamMsg, }) - s.handleFailoverSideEffects(upstreamCtx, resp, account, respBody, upstreamModel) + shouldDisable := s.handleFailoverSideEffects(upstreamCtx, resp, account, respBody, upstreamModel) return nil, &UpstreamFailoverError{ StatusCode: resp.StatusCode, ResponseBody: respBody, - RetryableOnSameAccount: account.IsPoolMode() && account.IsPoolModeRetryableStatus(resp.StatusCode), + RetryableOnSameAccount: !shouldDisable && account.IsPoolMode() && account.IsPoolModeRetryableStatus(resp.StatusCode), } } return s.handleOpenAIImagesErrorResponse(upstreamCtx, resp, c, account, upstreamModel) diff --git a/backend/internal/service/openai_images_responses.go b/backend/internal/service/openai_images_responses.go index f5b99ac9c..7924de4fc 100644 --- a/backend/internal/service/openai_images_responses.go +++ b/backend/internal/service/openai_images_responses.go @@ -920,7 +920,7 @@ func (s *OpenAIGatewayService) handleOpenAIImagesErrorResponse( return nil, &UpstreamFailoverError{ StatusCode: resp.StatusCode, ResponseBody: body, - RetryableOnSameAccount: account.IsPoolMode() && account.IsPoolModeRetryableStatus(resp.StatusCode), + RetryableOnSameAccount: false, } } @@ -1752,11 +1752,11 @@ func (s *OpenAIGatewayService) forwardOpenAIImagesOAuth( Kind: "failover", Message: upstreamMsg, }) - s.handleFailoverSideEffects(upstreamCtx, resp, account, respBody, requestModel) + shouldDisable := s.handleFailoverSideEffects(upstreamCtx, resp, account, respBody, requestModel) return nil, &UpstreamFailoverError{ StatusCode: resp.StatusCode, ResponseBody: respBody, - RetryableOnSameAccount: account.IsPoolMode() && account.IsPoolModeRetryableStatus(resp.StatusCode), + RetryableOnSameAccount: !shouldDisable && account.IsPoolMode() && account.IsPoolModeRetryableStatus(resp.StatusCode), } } return s.handleOpenAIImagesErrorResponse(upstreamCtx, resp, c, account, requestModel) @@ -1885,11 +1885,11 @@ func (s *OpenAIGatewayService) handleOpenAIImagesOAuthResponseError( } responseBody := openAIImagesUpstreamErrorResponseBody(upstreamErr) - s.handleOpenAIAccountUpstreamError(ctx, account, upstreamErr.StatusCode, headers, responseBody, requestedModel) + shouldDisable := s.handleOpenAIAccountUpstreamError(ctx, account, upstreamErr.StatusCode, headers, responseBody, requestedModel) return &UpstreamFailoverError{ StatusCode: upstreamErr.StatusCode, ResponseBody: responseBody, ResponseHeaders: headers, - RetryableOnSameAccount: account.IsPoolMode() && account.IsPoolModeRetryableStatus(upstreamErr.StatusCode), + RetryableOnSameAccount: !shouldDisable && account.IsPoolMode() && account.IsPoolModeRetryableStatus(upstreamErr.StatusCode), } } diff --git a/backend/internal/service/ratelimit_service.go b/backend/internal/service/ratelimit_service.go index 507e02d87..559ee2e05 100644 --- a/backend/internal/service/ratelimit_service.go +++ b/backend/internal/service/ratelimit_service.go @@ -159,6 +159,11 @@ func (s *RateLimitService) CheckErrorPolicy(ctx context.Context, account *Accoun return ErrorPolicySkipped } if account.IsPoolMode() { + // 池模式只跳过默认账号状态处理；管理员显式配置的临时不可调度规则仍应生效。 + // 401 保留现有认证错误语义，避免改变重复 401 的升级行为。 + if statusCode != http.StatusUnauthorized && s.tryTempUnschedulable(ctx, account, statusCode, responseBody) { + return ErrorPolicyTempUnscheduled + } return ErrorPolicySkipped } if s.tryTempUnschedulable(ctx, account, statusCode, responseBody) { @@ -172,8 +177,12 @@ func (s *RateLimitService) CheckErrorPolicy(ctx context.Context, account *Accoun func (s *RateLimitService) HandleUpstreamError(ctx context.Context, account *Account, statusCode int, headers http.Header, responseBody []byte, requestedModel ...string) (shouldDisable bool) { customErrorCodesEnabled := account.IsCustomErrorCodesEnabled() - // 池模式默认不标记本地账号状态；仅当用户显式配置自定义错误码时按本地策略处理。 + // 池模式默认不标记本地账号状态；但管理员显式配置的临时不可调度规则优先。 + // 401 保留现有认证错误语义，不在这里改变池模式的认证处理。 if account.IsPoolMode() && !customErrorCodesEnabled { + if statusCode != http.StatusUnauthorized && s.tryTempUnschedulable(ctx, account, statusCode, responseBody) { + return true + } slog.Info("pool_mode_error_skipped", "account_id", account.ID, "status_code", statusCode) return false } From e86063155f456491c4951fd9dde06ea879a87a85 Mon Sep 17 00:00:00 2001 From: Heatherm Huang Date: Fri, 17 Jul 2026 19:15:16 +0800 Subject: [PATCH 04/21] fix(grok): gate OAuth media on paid eligibility --- README.md | 4 +- backend/cmd/server/wire_gen.go | 2 +- .../internal/handler/admin/account_handler.go | 2 +- .../handler/admin/grok_import_probe.go | 12 ++-- .../handler/admin/grok_import_probe_test.go | 4 +- .../handler/admin/grok_oauth_handler.go | 2 +- backend/internal/handler/grok_media.go | 37 ++++++++++- backend/internal/handler/grok_media_test.go | 66 +++++++++++++++++++ .../handler/openai_gateway_handler.go | 29 ++++---- backend/internal/handler/wire.go | 2 + backend/internal/service/account.go | 22 +++++-- .../account_grok_media_eligibility_test.go | 42 +++++++++++- backend/internal/service/grok_media.go | 19 +++--- .../internal/service/grok_quota_service.go | 17 +++++ .../service/grok_quota_service_test.go | 59 +++++++++++++++++ .../service/openai_gateway_grok_test.go | 45 +++++++++++-- 16 files changed, 317 insertions(+), 47 deletions(-) diff --git a/README.md b/README.md index dd415c307..9695d7d5a 100644 --- a/README.md +++ b/README.md @@ -788,9 +788,9 @@ xAI quota is passive. Sub2API does not invent subscription quota values; it reco `401` responses temporarily remove accounts with invalid credentials from scheduling. `403` responses are treated as access or entitlement failures instead of token-refresh loops. `429` responses use `Retry-After` or a short cooldown to temporarily remove the account from scheduling. -New Grok image and video generation requests use a media-specific eligibility check. An OAuth account is excluded from new media generation when its recorded weekly or monthly billing probe returns `403`; chat requests and video status lookups are not affected by this media-only quarantine. If no eligible account remains, the media endpoint returns HTTP `503` with error type `grok_media_no_eligible_account` instead of forwarding the request to a known-ineligible account. +New Grok image and video generation requests use a media-specific eligibility check. API-key accounts remain eligible. OAuth accounts require positive paid-entitlement evidence from the xAI billing probe; Free, forbidden, missing, malformed, and inconclusive billing observations are excluded from new media generation. Unobserved OAuth accounts are probed before the first media request is forwarded, and imports run the billing-first quota probe proactively. Chat requests and video status lookups are not affected by this media-only quarantine. If no eligible account remains, the media endpoint returns HTTP `503` with error type `grok_media_no_eligible_account`. -Administrators can override automatic media eligibility through the account create/update API by setting `extra.grok_media_eligible` to `false` (exclude) or `true` (force eligible). On update, set it to `null` to remove the override and return to automatic probe-based behavior; omitting the field preserves the current override. A missing billing observation does not block legacy routing, and a weekly allowance period by itself is not treated as evidence that the account is ineligible. +Administrators can override automatic media eligibility through the account create/update API by setting `extra.grok_media_eligible` to `false` (exclude) or `true` (force eligible). On update, set it to `null` to remove the override and return to automatic probe-based behavior; omitting the field preserves the current override. A weekly allowance period alone is not treated as a paid tier signal. Successful image responses must contain at least one actual image output; empty HTTP `200` responses trigger account failover instead of being counted and returned as successful generations. --- diff --git a/backend/cmd/server/wire_gen.go b/backend/cmd/server/wire_gen.go index df6659d68..17ae478dd 100644 --- a/backend/cmd/server/wire_gen.go +++ b/backend/cmd/server/wire_gen.go @@ -270,7 +270,7 @@ func initializeApplication(buildInfo handler.BuildInfo) (*Application, error) { legacyEngine := securityaudit.NewLegacyModerationAdapter(contentModerationService) coordinator := securityaudit.NewCoordinator(legacyEngine, promptService) gatewayHandler := handler.ProvideGatewayHandler(gatewayService, openAIGatewayService, geminiMessagesCompatService, antigravityGatewayService, userService, concurrencyService, billingCacheService, usageService, apiKeyService, usageRecordWorkerPool, errorPassthroughService, contentModerationService, userMessageQueueService, configConfig, settingService, coordinator) - openAIGatewayHandler := handler.ProvideOpenAIGatewayHandler(openAIGatewayService, concurrencyService, billingCacheService, apiKeyService, usageRecordWorkerPool, errorPassthroughService, contentModerationService, opsService, configConfig, coordinator) + openAIGatewayHandler := handler.ProvideOpenAIGatewayHandler(openAIGatewayService, concurrencyService, billingCacheService, apiKeyService, usageRecordWorkerPool, errorPassthroughService, contentModerationService, opsService, grokQuotaService, configConfig, coordinator) handlerSettingHandler := handler.ProvideSettingHandler(settingService, buildInfo, notificationEmailService) totpHandler := handler.NewTotpHandler(totpService) handlerPaymentHandler := handler.NewPaymentHandler(paymentService, paymentConfigService) diff --git a/backend/internal/handler/admin/account_handler.go b/backend/internal/handler/admin/account_handler.go index d9b742003..46cc43ad5 100644 --- a/backend/internal/handler/admin/account_handler.go +++ b/backend/internal/handler/admin/account_handler.go @@ -61,7 +61,7 @@ type AccountHandler struct { sessionLimitCache service.SessionLimitCache rpmCache service.RPMCache tokenCacheInvalidator service.TokenCacheInvalidator - grokImportProber grokUsageProber + grokImportProber grokImportProber upstreamBillingProbe *service.UpstreamBillingProbeService } diff --git a/backend/internal/handler/admin/grok_import_probe.go b/backend/internal/handler/admin/grok_import_probe.go index b42957ce6..9ef4ce8cd 100644 --- a/backend/internal/handler/admin/grok_import_probe.go +++ b/backend/internal/handler/admin/grok_import_probe.go @@ -15,12 +15,12 @@ const ( grokImportProbeTimeout = 25 * time.Second ) -type grokUsageProber interface { - ProbeUsage(ctx context.Context, accountID int64) (*service.GrokQuotaProbeResult, error) +type grokImportProber interface { + QueryQuota(ctx context.Context, accountID int64) (*service.GrokQuotaProbeResult, error) } type grokImportProbeTask struct { - prober grokUsageProber + prober grokImportProber accountID int64 } @@ -51,7 +51,7 @@ func newGrokImportProbeScheduler(concurrency int, timeout time.Duration) *grokIm } } -func (s *grokImportProbeScheduler) schedule(prober grokUsageProber, account *service.Account) { +func (s *grokImportProbeScheduler) schedule(prober grokImportProber, account *service.Account) { if s == nil || prober == nil || account == nil || account.ID <= 0 { return } @@ -97,7 +97,7 @@ func (s *grokImportProbeScheduler) nextTask() (grokImportProbeTask, bool) { return task, true } -func (s *grokImportProbeScheduler) run(prober grokUsageProber, accountID int64) { +func (s *grokImportProbeScheduler) run(prober grokImportProber, accountID int64) { defer func() { if recovered := recover(); recovered != nil { slog.Error( @@ -112,7 +112,7 @@ func (s *grokImportProbeScheduler) run(prober grokUsageProber, accountID int64) // while this timeout only bounds the actual upstream probe execution. ctx, cancel := context.WithTimeout(context.Background(), s.timeout) defer cancel() - result, err := prober.ProbeUsage(ctx, accountID) + result, err := prober.QueryQuota(ctx, accountID) if err != nil { slog.Warn( "grok_import_active_probe_failed", diff --git a/backend/internal/handler/admin/grok_import_probe_test.go b/backend/internal/handler/admin/grok_import_probe_test.go index 3b8fc0ca6..255f04d0d 100644 --- a/backend/internal/handler/admin/grok_import_probe_test.go +++ b/backend/internal/handler/admin/grok_import_probe_test.go @@ -36,7 +36,7 @@ func newGrokImportProbeStub(buffer int) *grokImportProbeStub { } } -func (s *grokImportProbeStub) ProbeUsage(ctx context.Context, accountID int64) (*service.GrokQuotaProbeResult, error) { +func (s *grokImportProbeStub) QueryQuota(ctx context.Context, accountID int64) (*service.GrokQuotaProbeResult, error) { _, deadlineSeen := ctx.Deadline() s.mu.Lock() s.calls[accountID]++ @@ -69,7 +69,7 @@ func (s *grokImportProbeStub) ProbeUsage(ctx context.Context, accountID int64) ( return nil, failure } return &service.GrokQuotaProbeResult{ - Source: "active_probe", + Source: "hybrid_probe", Model: "grok-4.5", StatusCode: 200, ResetSupported: false, diff --git a/backend/internal/handler/admin/grok_oauth_handler.go b/backend/internal/handler/admin/grok_oauth_handler.go index ac05fd8a6..1f679b956 100644 --- a/backend/internal/handler/admin/grok_oauth_handler.go +++ b/backend/internal/handler/admin/grok_oauth_handler.go @@ -23,7 +23,7 @@ type GrokOAuthHandler struct { grokOAuthService *service.GrokOAuthService adminService service.AdminService quotaService *service.GrokQuotaService - importProber grokUsageProber + importProber grokImportProber reconciler service.GrokOAuthReconciler } diff --git a/backend/internal/handler/grok_media.go b/backend/internal/handler/grok_media.go index 6ae712945..9b96a9a12 100644 --- a/backend/internal/handler/grok_media.go +++ b/backend/internal/handler/grok_media.go @@ -167,6 +167,7 @@ func (h *OpenAIGatewayHandler) handleGrokMedia(c *gin.Context, endpoint service. sameAccountRetryCount := make(map[int64]int) var lastFailoverErr *service.UpstreamFailoverError var oauth429FailoverState service.OpenAIOAuth429FailoverState + mediaEligibilityRejected := false switchCount := 0 maxAccountSwitches := h.maxAccountSwitches if maxAccountSwitches <= 0 { @@ -202,7 +203,8 @@ func (h *OpenAIGatewayHandler) handleGrokMedia(c *gin.Context, endpoint service. zap.Error(err), zap.Int("excluded_account_count", len(failedAccountIDs)), ) - if endpoint.IsGenerationRequest() && len(failedAccountIDs) == 0 && errors.Is(err, service.ErrNoAvailableAccounts) { + if endpoint.IsGenerationRequest() && errors.Is(err, service.ErrNoAvailableAccounts) && + (len(failedAccountIDs) == 0 || (mediaEligibilityRejected && lastFailoverErr == nil)) { markOpsRoutingCapacityLimitedIfNoAvailable(c, err) h.errorResponse(c, http.StatusServiceUnavailable, "grok_media_no_eligible_account", "No eligible Grok media accounts") return @@ -246,6 +248,25 @@ func (h *OpenAIGatewayHandler) handleGrokMedia(c *gin.Context, endpoint service. ) account := selection.Account + if endpoint.IsGenerationRequest() { + eligible, eligibilityReason, eligibilityErr := h.ensureGrokMediaAccountEligibility(requestCtx, account) + if !eligible { + mediaEligibilityRejected = true + failedAccountIDs[account.ID] = struct{}{} + reqLog.Warn("grok_media.account_eligibility_rejected", + zap.Int64("account_id", account.ID), + zap.String("reason", eligibilityReason), + zap.Bool("probe_failed", eligibilityErr != nil), + ) + if switchCount >= maxAccountSwitches { + markOpsRoutingCapacityLimited(c) + h.errorResponse(c, http.StatusServiceUnavailable, "grok_media_no_eligible_account", "No eligible Grok media accounts") + return + } + switchCount++ + continue + } + } sessionHash = ensureOpenAIPoolModeSessionHash(sessionHash, account) setOpsSelectedAccount(c, account.ID, account.Platform) @@ -365,6 +386,20 @@ func (h *OpenAIGatewayHandler) handleGrokMedia(c *gin.Context, endpoint service. } } +func (h *OpenAIGatewayHandler) ensureGrokMediaAccountEligibility(ctx context.Context, account *service.Account) (bool, string, error) { + if account == nil { + return false, "missing_account", errors.New("grok media account is required") + } + eligible, reason := account.GrokMediaGenerationEligibility() + if eligible || reason != "billing_unobserved" { + return eligible, reason, nil + } + if h == nil || h.grokMediaEligibilityProber == nil { + return false, "billing_probe_unavailable", errors.New("grok media eligibility probe is not configured") + } + return h.grokMediaEligibilityProber.ProbeMediaEligibility(ctx, account.ID) +} + func grokMediaRequiredCapability(endpoint service.GrokMediaEndpoint) service.OpenAIEndpointCapability { if endpoint.IsGenerationRequest() { return service.OpenAIEndpointCapabilityGrokMediaGeneration diff --git a/backend/internal/handler/grok_media_test.go b/backend/internal/handler/grok_media_test.go index 3c58742c7..53348958b 100644 --- a/backend/internal/handler/grok_media_test.go +++ b/backend/internal/handler/grok_media_test.go @@ -1,12 +1,26 @@ package handler import ( + "context" + "errors" "testing" "github.com/Wei-Shaw/sub2api/internal/service" "github.com/stretchr/testify/require" ) +type grokMediaEligibilityProberStub struct { + eligible bool + reason string + err error + calls int +} + +func (s *grokMediaEligibilityProberStub) ProbeMediaEligibility(context.Context, int64) (bool, string, error) { + s.calls++ + return s.eligible, s.reason, s.err +} + func TestShouldRecordGrokMediaUsage(t *testing.T) { tests := []struct { name string @@ -73,3 +87,55 @@ func TestGrokMediaRequiredCapability(t *testing.T) { }) } } + +func TestEnsureGrokMediaAccountEligibility(t *testing.T) { + t.Run("non oauth account does not probe", func(t *testing.T) { + prober := &grokMediaEligibilityProberStub{} + h := &OpenAIGatewayHandler{grokMediaEligibilityProber: prober} + account := &service.Account{Platform: service.PlatformGrok, Type: service.AccountTypeAPIKey} + + eligible, reason, err := h.ensureGrokMediaAccountEligibility(context.Background(), account) + + require.NoError(t, err) + require.True(t, eligible) + require.Equal(t, "non_oauth", reason) + require.Zero(t, prober.calls) + }) + + t.Run("unobserved oauth is probed before forwarding", func(t *testing.T) { + prober := &grokMediaEligibilityProberStub{eligible: true, reason: "eligible"} + h := &OpenAIGatewayHandler{grokMediaEligibilityProber: prober} + account := &service.Account{ID: 7, Platform: service.PlatformGrok, Type: service.AccountTypeOAuth} + + eligible, reason, err := h.ensureGrokMediaAccountEligibility(context.Background(), account) + + require.NoError(t, err) + require.True(t, eligible) + require.Equal(t, "eligible", reason) + require.Equal(t, 1, prober.calls) + }) + + t.Run("missing prober fails closed", func(t *testing.T) { + h := &OpenAIGatewayHandler{} + account := &service.Account{ID: 8, Platform: service.PlatformGrok, Type: service.AccountTypeOAuth} + + eligible, reason, err := h.ensureGrokMediaAccountEligibility(context.Background(), account) + + require.Error(t, err) + require.False(t, eligible) + require.Equal(t, "billing_probe_unavailable", reason) + }) + + t.Run("probe failure fails closed", func(t *testing.T) { + probeErr := errors.New("probe failed") + prober := &grokMediaEligibilityProberStub{reason: "billing_unobserved", err: probeErr} + h := &OpenAIGatewayHandler{grokMediaEligibilityProber: prober} + account := &service.Account{ID: 9, Platform: service.PlatformGrok, Type: service.AccountTypeOAuth} + + eligible, reason, err := h.ensureGrokMediaAccountEligibility(context.Background(), account) + + require.ErrorIs(t, err, probeErr) + require.False(t, eligible) + require.Equal(t, "billing_unobserved", reason) + }) +} diff --git a/backend/internal/handler/openai_gateway_handler.go b/backend/internal/handler/openai_gateway_handler.go index ec4e06dd8..3dcc170cc 100644 --- a/backend/internal/handler/openai_gateway_handler.go +++ b/backend/internal/handler/openai_gateway_handler.go @@ -28,18 +28,23 @@ import ( // OpenAIGatewayHandler handles OpenAI API gateway requests type OpenAIGatewayHandler struct { - gatewayService *service.OpenAIGatewayService - billingCacheService *service.BillingCacheService - apiKeyService *service.APIKeyService - usageRecordWorkerPool *service.UsageRecordWorkerPool - errorPassthroughService *service.ErrorPassthroughService - contentModerationService *service.ContentModerationService - securityAuditCoordinator *securityaudit.Coordinator - opsService *service.OpsService - concurrencyHelper *ConcurrencyHelper - imageLimiter *imageConcurrencyLimiter - maxAccountSwitches int - cfg *config.Config + gatewayService *service.OpenAIGatewayService + billingCacheService *service.BillingCacheService + apiKeyService *service.APIKeyService + usageRecordWorkerPool *service.UsageRecordWorkerPool + errorPassthroughService *service.ErrorPassthroughService + contentModerationService *service.ContentModerationService + securityAuditCoordinator *securityaudit.Coordinator + grokMediaEligibilityProber grokMediaEligibilityProber + opsService *service.OpsService + concurrencyHelper *ConcurrencyHelper + imageLimiter *imageConcurrencyLimiter + maxAccountSwitches int + cfg *config.Config +} + +type grokMediaEligibilityProber interface { + ProbeMediaEligibility(ctx context.Context, accountID int64) (bool, string, error) } const maxOpenAIFirstOutputTimeoutSwitches = 1 diff --git a/backend/internal/handler/wire.go b/backend/internal/handler/wire.go index 57bf64d02..888df38b1 100644 --- a/backend/internal/handler/wire.go +++ b/backend/internal/handler/wire.go @@ -120,12 +120,14 @@ func ProvideOpenAIGatewayHandler( errorPassthroughService *service.ErrorPassthroughService, contentModerationService *service.ContentModerationService, opsService *service.OpsService, + grokQuotaService *service.GrokQuotaService, cfg *config.Config, coordinator *securityaudit.Coordinator, ) *OpenAIGatewayHandler { h := NewOpenAIGatewayHandler(gatewayService, concurrencyService, billingCacheService, apiKeyService, usageRecordWorkerPool, errorPassthroughService, contentModerationService, opsService, cfg) h.securityAuditCoordinator = coordinator + h.grokMediaEligibilityProber = grokQuotaService return h } diff --git a/backend/internal/service/account.go b/backend/internal/service/account.go index 724925631..5b3707d2a 100644 --- a/backend/internal/service/account.go +++ b/backend/internal/service/account.go @@ -1423,8 +1423,12 @@ func (a *Account) SupportsOpenAIEndpointCapability(capability OpenAIEndpointCapa case OpenAIEndpointCapabilityChatCompletions: return true case OpenAIEndpointCapabilityGrokMediaGeneration: - eligible, _ := a.GrokMediaGenerationEligibility() - return eligible + eligible, reason := a.GrokMediaGenerationEligibility() + // Unobserved OAuth accounts remain scheduler candidates only so the + // request path can run the billing probe before forwarding. The + // forwarding gate itself fails closed if that probe is unavailable or + // cannot produce positive paid-entitlement evidence. + return eligible || reason == "billing_unobserved" default: return false } @@ -1469,9 +1473,9 @@ func (a *Account) SupportsOpenAIEndpointCapability(capability OpenAIEndpointCapa } // GrokMediaGenerationEligibility reports whether a Grok account may receive -// new image/video generation requests. Missing observations preserve legacy -// routing; operators can fail closed for known-bad accounts with the explicit -// override. A successful override takes precedence over stale probe data. +// new image/video generation requests. OAuth media fails closed unless billing +// observations provide positive paid-entitlement evidence. An explicit +// operator override takes precedence over probe data. func (a *Account) GrokMediaGenerationEligibility() (bool, string) { if a == nil || !a.IsGrok() { return false, "not_grok" @@ -1488,11 +1492,17 @@ func (a *Account) GrokMediaGenerationEligibility() (bool, string) { billing, err := grokBillingSnapshotFromExtra(a.Extra) if err != nil || billing == nil { - return true, "billing_unobserved" + return false, "billing_unobserved" } if billing.StatusCode == 403 || billing.WeeklyStatusCode == 403 || billing.MonthlyStatusCode == 403 { return false, "billing_forbidden" } + if isKnownGrokFreeAccount(a) { + return false, "billing_free_tier" + } + if !grokBillingHasAuthoritativeQuota(billing) { + return false, "billing_inconclusive" + } return true, "eligible" } diff --git a/backend/internal/service/account_grok_media_eligibility_test.go b/backend/internal/service/account_grok_media_eligibility_test.go index be7b4759a..a79f9c02e 100644 --- a/backend/internal/service/account_grok_media_eligibility_test.go +++ b/backend/internal/service/account_grok_media_eligibility_test.go @@ -13,6 +13,7 @@ import ( ) func TestGrokMediaGenerationEligibility(t *testing.T) { + weeklyUsagePercent := 12.5 forbiddenBilling := &xai.BillingSummary{ StatusCode: http.StatusForbidden, WeeklyStatusCode: http.StatusForbidden, @@ -20,9 +21,24 @@ func TestGrokMediaGenerationEligibility(t *testing.T) { } weeklyAllowance := &xai.BillingSummary{ PeriodType: "weekly", + UsagePercent: &weeklyUsagePercent, StatusCode: http.StatusOK, WeeklyStatusCode: http.StatusOK, } + freeBilling := &xai.BillingSummary{ + PeriodType: "monthly", + StatusCode: http.StatusOK, + WeeklyStatusCode: http.StatusOK, + MonthlyStatusCode: http.StatusOK, + MonthlyUpdatedAt: "2026-07-17T00:00:00Z", + } + inconclusiveBilling := &xai.BillingSummary{ + StatusCode: http.StatusOK, + WeeklyStatusCode: http.StatusOK, + MonthlyStatusCode: http.StatusBadGateway, + Partial: true, + FailedWindows: []string{"monthly"}, + } weeklyForbidden := &xai.BillingSummary{ StatusCode: http.StatusOK, WeeklyStatusCode: http.StatusForbidden, @@ -43,12 +59,14 @@ func TestGrokMediaGenerationEligibility(t *testing.T) { {name: "nil account", account: nil, want: false, wantReason: "not_grok"}, {name: "non grok account", account: &Account{Platform: PlatformOpenAI}, want: false, wantReason: "not_grok"}, {name: "non oauth grok account stays eligible", account: &Account{Platform: PlatformGrok, Type: AccountTypeAPIKey}, want: true, wantReason: "non_oauth"}, - {name: "unobserved oauth preserves legacy routing", account: &Account{Platform: PlatformGrok, Type: AccountTypeOAuth}, want: true, wantReason: "billing_unobserved"}, - {name: "weekly allowance is not treated as weekly subscription", account: &Account{Platform: PlatformGrok, Type: AccountTypeOAuth, Extra: map[string]any{grokBillingExtraKey: weeklyAllowance}}, want: true, wantReason: "eligible"}, + {name: "unobserved oauth fails closed", account: &Account{Platform: PlatformGrok, Type: AccountTypeOAuth}, want: false, wantReason: "billing_unobserved"}, + {name: "weekly paid usage is eligible without inferring from period type", account: &Account{Platform: PlatformGrok, Type: AccountTypeOAuth, Extra: map[string]any{grokBillingExtraKey: weeklyAllowance}}, want: true, wantReason: "eligible"}, + {name: "observed free account is rejected", account: &Account{Platform: PlatformGrok, Type: AccountTypeOAuth, Extra: map[string]any{grokBillingExtraKey: freeBilling}}, want: false, wantReason: "billing_free_tier"}, + {name: "inconclusive billing fails closed", account: &Account{Platform: PlatformGrok, Type: AccountTypeOAuth, Extra: map[string]any{grokBillingExtraKey: inconclusiveBilling}}, want: false, wantReason: "billing_inconclusive"}, {name: "billing forbidden is rejected", account: &Account{Platform: PlatformGrok, Type: AccountTypeOAuth, Extra: map[string]any{grokBillingExtraKey: forbiddenBilling}}, want: false, wantReason: "billing_forbidden"}, {name: "weekly billing forbidden is rejected after partial success", account: &Account{Platform: PlatformGrok, Type: AccountTypeOAuth, Extra: map[string]any{grokBillingExtraKey: weeklyForbidden}}, want: false, wantReason: "billing_forbidden"}, {name: "monthly billing forbidden is rejected after partial success", account: &Account{Platform: PlatformGrok, Type: AccountTypeOAuth, Extra: map[string]any{grokBillingExtraKey: monthlyForbidden}}, want: false, wantReason: "billing_forbidden"}, - {name: "malformed billing observation preserves legacy routing", account: &Account{Platform: PlatformGrok, Type: AccountTypeOAuth, Extra: map[string]any{grokBillingExtraKey: make(chan int)}}, want: true, wantReason: "billing_unobserved"}, + {name: "malformed billing observation fails closed", account: &Account{Platform: PlatformGrok, Type: AccountTypeOAuth, Extra: map[string]any{grokBillingExtraKey: make(chan int)}}, want: false, wantReason: "billing_unobserved"}, {name: "malformed override falls back to observations", account: &Account{Platform: PlatformGrok, Type: AccountTypeOAuth, Extra: map[string]any{GrokMediaEligibleExtraKey: "false", grokBillingExtraKey: weeklyAllowance}}, want: true, wantReason: "eligible"}, {name: "explicit disable wins", account: &Account{Platform: PlatformGrok, Type: AccountTypeOAuth, Extra: map[string]any{GrokMediaEligibleExtraKey: false}}, want: false, wantReason: "override_disabled"}, {name: "explicit enable wins over forbidden probe", account: &Account{Platform: PlatformGrok, Type: AccountTypeOAuth, Extra: map[string]any{GrokMediaEligibleExtraKey: true, grokBillingExtraKey: forbiddenBilling}}, want: true, wantReason: "override_enabled"}, @@ -63,6 +81,24 @@ func TestGrokMediaGenerationEligibility(t *testing.T) { } } +func TestGrokMediaCapabilityKeepsOnlyUnobservedOAuthAsProbeCandidate(t *testing.T) { + unobserved := &Account{Platform: PlatformGrok, Type: AccountTypeOAuth} + eligible, reason := unobserved.GrokMediaGenerationEligibility() + require.False(t, eligible) + require.Equal(t, "billing_unobserved", reason) + require.True(t, unobserved.SupportsOpenAIEndpointCapability(OpenAIEndpointCapabilityGrokMediaGeneration)) + + inconclusive := &Account{ + Platform: PlatformGrok, + Type: AccountTypeOAuth, + Extra: map[string]any{grokBillingExtraKey: &xai.BillingSummary{ + StatusCode: http.StatusOK, + Partial: true, + }}, + } + require.False(t, inconclusive.SupportsOpenAIEndpointCapability(OpenAIEndpointCapabilityGrokMediaGeneration)) +} + func TestGrokMediaCapabilityFiltersOnlyGeneration(t *testing.T) { account := &Account{ ID: 1, diff --git a/backend/internal/service/grok_media.go b/backend/internal/service/grok_media.go index c75bd271f..abfe0df26 100644 --- a/backend/internal/service/grok_media.go +++ b/backend/internal/service/grok_media.go @@ -364,6 +364,16 @@ func (s *OpenAIGatewayService) ForwardGrokMedia( if err != nil { return nil, err } + if endpoint == GrokMediaEndpointImagesGenerations || endpoint == GrokMediaEndpointImagesEdits { + if countOpenAIResponseImageOutputsFromJSONBytes(respBody) <= 0 { + setOpsUpstreamError(c, http.StatusBadGateway, "xAI upstream returned no image output", truncateString(string(respBody), 512)) + return nil, &UpstreamFailoverError{ + StatusCode: http.StatusBadGateway, + ResponseBody: respBody, + ResponseHeaders: resp.Header.Clone(), + } + } + } writeGrokMediaResponse(c, resp, respBody, s.responseHeaderFilter) usage := grokMediaUsageFromResponse(endpoint, requestInfo, respBody) return &OpenAIForwardResult{ @@ -584,14 +594,7 @@ func grokMediaUsageFromResponse(endpoint GrokMediaEndpoint, requestInfo GrokMedi meta := grokMediaUsageMetadata{Usage: usage} switch endpoint { case GrokMediaEndpointImagesGenerations, GrokMediaEndpointImagesEdits: - imageCount := countOpenAIResponseImageOutputsFromJSONBytes(responseBody) - if imageCount <= 0 { - imageCount = requestInfo.N - } - if imageCount <= 0 { - imageCount = 1 - } - meta.ImageCount = imageCount + meta.ImageCount = countOpenAIResponseImageOutputsFromJSONBytes(responseBody) meta.ImageSize = requestInfo.SizeTier meta.ImageInputSize = requestInfo.Size meta.ImageOutputSizes = collectOpenAIResponseImageOutputSizesFromJSONBytes(responseBody) diff --git a/backend/internal/service/grok_quota_service.go b/backend/internal/service/grok_quota_service.go index 45e7be83a..db11fe6cc 100644 --- a/backend/internal/service/grok_quota_service.go +++ b/backend/internal/service/grok_quota_service.go @@ -221,6 +221,23 @@ func (s *GrokQuotaService) ProbeBilling(ctx context.Context, accountID int64) (* }) } +// ProbeMediaEligibility refreshes billing state and evaluates the persisted +// account snapshot used by media scheduling. Probe failures remain fail-closed; +// deterministic persisted states such as forbidden or Free are returned as +// normal ineligibility decisions rather than transport errors. +func (s *GrokQuotaService) ProbeMediaEligibility(ctx context.Context, accountID int64) (bool, string, error) { + _, probeErr := s.ProbeBilling(ctx, accountID) + account, err := s.loadGrokOAuthAccount(ctx, accountID) + if err != nil { + return false, "billing_probe_failed", err + } + eligible, reason := account.GrokMediaGenerationEligibility() + if reason == "billing_unobserved" && probeErr != nil { + return false, reason, probeErr + } + return eligible, reason, nil +} + func (s *GrokQuotaService) probeBilling(ctx context.Context, accountID int64) (*GrokQuotaProbeResult, error) { account, token, proxyURL, err := s.prepareProbe(ctx, accountID) if err != nil { diff --git a/backend/internal/service/grok_quota_service_test.go b/backend/internal/service/grok_quota_service_test.go index 59b19ef27..728277e62 100644 --- a/backend/internal/service/grok_quota_service_test.go +++ b/backend/internal/service/grok_quota_service_test.go @@ -44,6 +44,18 @@ func (r *grokQuotaAccountRepo) UpdateExtra(_ context.Context, id int64, updates r.updates = make(map[int64]map[string]any) } r.updates[id] = updates + if r.mockAccountRepoForPlatform != nil { + account := r.accountsByID[id] + if account == nil { + return nil + } + if account.Extra == nil { + account.Extra = make(map[string]any) + } + for key, value := range updates { + account.Extra[key] = value + } + } return nil } @@ -824,6 +836,53 @@ func TestGrokQuotaServicePartialBilling403PersistsMediaEligibilitySignal(t *test require.Equal(t, "billing_forbidden", reason) } +func TestGrokQuotaServiceProbeMediaEligibility(t *testing.T) { + t.Run("positive paid evidence enables media", func(t *testing.T) { + usagePercent := 10.0 + monthlyLimit := 15_000.0 + account := healthyGrokQuotaOAuthAccount(60) + repo := &grokQuotaAccountRepo{mockAccountRepoForPlatform: &mockAccountRepoForPlatform{ + accountsByID: map[int64]*Account{account.ID: account}, + }} + upstream := &grokHybridUpstream{weeklyUsagePercent: &usagePercent, monthlyLimitCents: &monthlyLimit} + svc := NewGrokQuotaService(repo, nil, NewGrokTokenProvider(repo, nil), upstream, nil) + + eligible, reason, err := svc.ProbeMediaEligibility(context.Background(), account.ID) + + require.NoError(t, err) + require.True(t, eligible) + require.Equal(t, "eligible", reason) + }) + + t.Run("successful empty billing identifies free account", func(t *testing.T) { + account := healthyGrokQuotaOAuthAccount(61) + repo := &grokQuotaAccountRepo{mockAccountRepoForPlatform: &mockAccountRepoForPlatform{ + accountsByID: map[int64]*Account{account.ID: account}, + }} + svc := NewGrokQuotaService(repo, nil, NewGrokTokenProvider(repo, nil), &grokHybridUpstream{}, nil) + + eligible, reason, err := svc.ProbeMediaEligibility(context.Background(), account.ID) + + require.NoError(t, err) + require.False(t, eligible) + require.Equal(t, "billing_free_tier", reason) + }) + + t.Run("forbidden billing is deterministic ineligibility", func(t *testing.T) { + account := healthyGrokQuotaOAuthAccount(62) + repo := &grokQuotaAccountRepo{mockAccountRepoForPlatform: &mockAccountRepoForPlatform{ + accountsByID: map[int64]*Account{account.ID: account}, + }} + svc := NewGrokQuotaService(repo, nil, NewGrokTokenProvider(repo, nil), &grokHybridUpstream{billingStatus: http.StatusForbidden}, nil) + + eligible, reason, err := svc.ProbeMediaEligibility(context.Background(), account.ID) + + require.NoError(t, err) + require.False(t, eligible) + require.Equal(t, "billing_forbidden", reason) + }) +} + func TestPreferBillingObservationStatus(t *testing.T) { t.Parallel() diff --git a/backend/internal/service/openai_gateway_grok_test.go b/backend/internal/service/openai_gateway_grok_test.go index 67ad71214..c238e69d3 100644 --- a/backend/internal/service/openai_gateway_grok_test.go +++ b/backend/internal/service/openai_gateway_grok_test.go @@ -551,7 +551,7 @@ func TestForwardGrokMediaImagesGenerationNormalizesImagineAlias(t *testing.T) { "Content-Type": []string{"application/json"}, "Xai-Request-Id": []string{"xai-image-req"}, }, - Body: io.NopCloser(strings.NewReader(`{"data":[]}`)), + Body: io.NopCloser(strings.NewReader(`{"data":[{"url":"https://images.test/cat.png"}]}`)), }} svc := &OpenAIGatewayService{httpUpstream: upstream} @@ -565,7 +565,7 @@ func TestForwardGrokMediaImagesGenerationNormalizesImagineAlias(t *testing.T) { require.NotEqual(t, grokUpstreamUserAgent, upstream.lastReq.Header.Get("User-Agent")) require.JSONEq(t, `{"model":"grok-imagine-image-quality","prompt":"draw a cat"}`, string(upstream.lastBody)) require.Equal(t, http.StatusOK, recorder.Code) - require.JSONEq(t, `{"data":[]}`, recorder.Body.String()) + require.JSONEq(t, `{"data":[{"url":"https://images.test/cat.png"}]}`, recorder.Body.String()) require.Equal(t, "xai-image-req", result.RequestID) require.Equal(t, "grok-imagine-image-quality", result.Model) require.Equal(t, "grok-imagine-image-quality", result.BillingModel) @@ -573,6 +573,43 @@ func TestForwardGrokMediaImagesGenerationNormalizesImagineAlias(t *testing.T) { require.Equal(t, ImageBillingSize2K, result.ImageSize) } +func TestForwardGrokMediaImagesGenerationRejectsEmptySuccessfulResponse(t *testing.T) { + t.Setenv(xai.EnvAllowUnsafeURLOverrides, "true") + gin.SetMode(gin.TestMode) + + recorder := httptest.NewRecorder() + c, _ := gin.CreateTestContext(recorder) + body := []byte(`{"model":"grok-imagine-image","prompt":"draw a cat"}`) + c.Request = httptest.NewRequest(http.MethodPost, "/v1/images/generations", bytes.NewReader(body)) + c.Request.Header.Set("Content-Type", "application/json") + + account := &Account{ + ID: 66, + Name: "grok", + Platform: PlatformGrok, + Type: AccountTypeAPIKey, + Concurrency: 1, + Credentials: map[string]any{ + "api_key": "api-key", + "base_url": "https://xai.test/v1", + }, + } + upstream := &httpUpstreamRecorder{resp: &http.Response{ + StatusCode: http.StatusOK, + Header: http.Header{"Content-Type": []string{"application/json"}}, + Body: io.NopCloser(strings.NewReader(`{"data":[]}`)), + }} + svc := &OpenAIGatewayService{httpUpstream: upstream} + + result, err := svc.ForwardGrokMedia(context.Background(), c, account, GrokMediaEndpointImagesGenerations, "", body, "application/json") + require.Nil(t, result) + var failoverErr *UpstreamFailoverError + require.ErrorAs(t, err, &failoverErr) + require.Equal(t, http.StatusBadGateway, failoverErr.StatusCode) + require.JSONEq(t, `{"data":[]}`, string(failoverErr.ResponseBody)) + require.Empty(t, recorder.Body.String()) +} + func TestForwardGrokMediaImagesGenerationStripsUnsupportedSize(t *testing.T) { t.Setenv(xai.EnvAllowUnsafeURLOverrides, "true") gin.SetMode(gin.TestMode) @@ -599,7 +636,7 @@ func TestForwardGrokMediaImagesGenerationStripsUnsupportedSize(t *testing.T) { Header: http.Header{ "Content-Type": []string{"application/json"}, }, - Body: io.NopCloser(strings.NewReader(`{"data":[]}`)), + Body: io.NopCloser(strings.NewReader(`{"data":[{"url":"https://images.test/cat.png"}]}`)), }} svc := &OpenAIGatewayService{httpUpstream: upstream} @@ -648,7 +685,7 @@ func TestForwardGrokMediaImagesEditMultipartConvertsToJSON(t *testing.T) { Header: http.Header{ "Content-Type": []string{"application/json"}, }, - Body: io.NopCloser(strings.NewReader(`{"data":[]}`)), + Body: io.NopCloser(strings.NewReader(`{"data":[{"url":"https://images.test/edited.png"}]}`)), }} svc := &OpenAIGatewayService{httpUpstream: upstream} From d2585097f38f8e495a0060f9b8ebf395e2139603 Mon Sep 17 00:00:00 2001 From: Tian Lee <498756723@qq.com> Date: Fri, 17 Jul 2026 14:15:05 +0800 Subject: [PATCH 05/21] =?UTF-8?q?fix:=20=E5=AE=8C=E5=96=84=E4=B8=8A?= =?UTF-8?q?=E6=B8=B8=20Sub2API=20=E8=AE=A1=E8=B4=B9=E5=80=8D=E7=8E=87?= =?UTF-8?q?=E6=8E=A2=E6=B5=8B=E4=B8=8E=E8=B4=A6=E5=8F=B7=E5=B1=95=E7=A4=BA?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- .../internal/handler/admin/account_handler.go | 7 +- .../account_handler_mixed_channel_test.go | 19 +++ .../handler/admin/admin_service_stub_test.go | 2 + backend/internal/repository/account_repo.go | 76 ++++++++- .../account_repo_sort_integration_test.go | 105 +++++++++++++ ...repo_upstream_billing_probe_update_test.go | 40 +++++ backend/internal/service/account_service.go | 1 + backend/internal/service/admin_account.go | 41 ++++- ...min_account_upstream_billing_probe_test.go | 83 ++++++++++ backend/internal/service/admin_service.go | 2 + .../service/upstream_billing_probe.go | 25 ++- .../service/upstream_billing_probe_test.go | 22 ++- .../account/BulkEditAccountModal.vue | 52 +++++++ .../components/account/CreateAccountModal.vue | 35 ++++- .../account/UpstreamBillingRateCell.vue | 111 +++++++++---- .../__tests__/BulkEditAccountModal.spec.ts | 67 ++++++++ .../__tests__/CreateAccountModal.spec.ts | 49 +++++- .../__tests__/UpstreamBillingRateCell.spec.ts | 146 ++++++++++++++++-- frontend/src/components/common/DataTable.vue | 60 +++---- .../common/__tests__/DataTable.spec.ts | 4 + .../src/i18n/locales/en/admin/accounts.ts | 20 ++- .../src/i18n/locales/en/admin/settings.ts | 10 ++ .../src/i18n/locales/zh/admin/accounts.ts | 20 ++- .../src/i18n/locales/zh/admin/settings.ts | 10 ++ frontend/src/types/index.ts | 1 + frontend/src/views/admin/AccountsView.vue | 141 ++++++++--------- frontend/src/views/admin/SettingsView.vue | 129 ++++++++++++++++ .../__tests__/AccountsView.bulkEdit.spec.ts | 146 +++++++++++++++++- .../AccountsView.usageWindowsHint.spec.ts | 2 + .../admin/__tests__/SettingsView.spec.ts | 68 ++++++++ 30 files changed, 1311 insertions(+), 183 deletions(-) diff --git a/backend/internal/handler/admin/account_handler.go b/backend/internal/handler/admin/account_handler.go index d9b742003..98adcf514 100644 --- a/backend/internal/handler/admin/account_handler.go +++ b/backend/internal/handler/admin/account_handler.go @@ -121,6 +121,7 @@ type CreateAccountRequest struct { GroupIDs []int64 `json:"group_ids"` ExpiresAt *int64 `json:"expires_at"` AutoPauseOnExpired *bool `json:"auto_pause_on_expired"` + ProbeEnabled *bool `json:"upstream_billing_probe_enabled"` ConfirmMixedChannelRisk *bool `json:"confirm_mixed_channel_risk"` // 用户确认混合渠道风险 } @@ -159,6 +160,7 @@ type BulkUpdateAccountsRequest struct { GroupIDs *[]int64 `json:"group_ids"` Credentials map[string]any `json:"credentials"` Extra map[string]any `json:"extra"` + ProbeEnabled *bool `json:"upstream_billing_probe_enabled"` ConfirmMixedChannelRisk *bool `json:"confirm_mixed_channel_risk"` // 用户确认混合渠道风险 } @@ -828,6 +830,7 @@ func (h *AccountHandler) Create(c *gin.Context) { GroupIDs: req.GroupIDs, ExpiresAt: req.ExpiresAt, AutoPauseOnExpired: req.AutoPauseOnExpired, + ProbeEnabled: req.ProbeEnabled, SkipMixedChannelCheck: skipCheck, }) if execErr != nil { @@ -1907,7 +1910,8 @@ func (h *AccountHandler) BulkUpdate(c *gin.Context) { req.Schedulable != nil || req.GroupIDs != nil || len(req.Credentials) > 0 || - len(req.Extra) > 0 + len(req.Extra) > 0 || + req.ProbeEnabled != nil if !hasUpdates { response.BadRequest(c, "No updates provided") @@ -1928,6 +1932,7 @@ func (h *AccountHandler) BulkUpdate(c *gin.Context) { GroupIDs: req.GroupIDs, Credentials: req.Credentials, Extra: req.Extra, + ProbeEnabled: req.ProbeEnabled, SkipMixedChannelCheck: skipCheck, }) if err != nil { diff --git a/backend/internal/handler/admin/account_handler_mixed_channel_test.go b/backend/internal/handler/admin/account_handler_mixed_channel_test.go index a1dc25124..bc4a3cf49 100644 --- a/backend/internal/handler/admin/account_handler_mixed_channel_test.go +++ b/backend/internal/handler/admin/account_handler_mixed_channel_test.go @@ -222,3 +222,22 @@ func TestBulkUpdateAcceptsFilterTargetRequest(t *testing.T) { require.NoError(t, json.Unmarshal(rec.Body.Bytes(), &resp)) require.Equal(t, float64(0), resp["code"]) } + +func TestBulkUpdateAcceptsDedicatedUpstreamBillingProbeSetting(t *testing.T) { + adminSvc := newStubAdminService() + router := setupAccountMixedChannelRouter(adminSvc) + + body, _ := json.Marshal(map[string]any{ + "account_ids": []int64{1, 2}, + "upstream_billing_probe_enabled": false, + }) + rec := httptest.NewRecorder() + req := httptest.NewRequest(http.MethodPost, "/api/v1/admin/accounts/bulk-update", bytes.NewReader(body)) + req.Header.Set("Content-Type", "application/json") + router.ServeHTTP(rec, req) + + require.Equal(t, http.StatusOK, rec.Code) + require.NotNil(t, adminSvc.lastBulkUpdateAccountInput) + require.NotNil(t, adminSvc.lastBulkUpdateAccountInput.ProbeEnabled) + require.False(t, *adminSvc.lastBulkUpdateAccountInput.ProbeEnabled) +} diff --git a/backend/internal/handler/admin/admin_service_stub_test.go b/backend/internal/handler/admin/admin_service_stub_test.go index f5c238a61..035e171c9 100644 --- a/backend/internal/handler/admin/admin_service_stub_test.go +++ b/backend/internal/handler/admin/admin_service_stub_test.go @@ -33,6 +33,7 @@ type stubAdminService struct { createSparkShadowErr error updateAccountErr error bulkUpdateAccountErr error + lastBulkUpdateAccountInput *service.BulkUpdateAccountsInput getAccountResult *service.Account updateAccountCalls int updateAccountExtraCalls int @@ -478,6 +479,7 @@ func (s *stubAdminService) SetAccountSchedulable(ctx context.Context, id int64, } func (s *stubAdminService) BulkUpdateAccounts(ctx context.Context, input *service.BulkUpdateAccountsInput) (*service.BulkUpdateAccountsResult, error) { + s.lastBulkUpdateAccountInput = input if s.bulkUpdateAccountErr != nil { return nil, s.bulkUpdateAccountErr } diff --git a/backend/internal/repository/account_repo.go b/backend/internal/repository/account_repo.go index f014155bb..52c19e143 100644 --- a/backend/internal/repository/account_repo.go +++ b/backend/internal/repository/account_repo.go @@ -893,6 +893,20 @@ func (r *accountRepository) ListOpsAccountsForStats(ctx context.Context, platfor func accountListOrder(params pagination.PaginationParams) []func(*entsql.Selector) { sortBy := strings.ToLower(strings.TrimSpace(params.SortBy)) sortOrder := params.NormalizedSortOrder(pagination.SortOrderAsc) + if sortBy == "upstream_billing_rate" { + direction := "ASC" + tieOrder := entsql.Asc + if sortOrder == pagination.SortOrderDesc { + direction = "DESC" + tieOrder = entsql.Desc + } + return []func(*entsql.Selector){func(s *entsql.Selector) { + extra := s.C(dbaccount.FieldExtra) + expression := upstreamBillingRateSortExpression(extra) + s.OrderExpr(entsql.Expr(expression + " " + direction + " NULLS LAST")) + s.OrderBy(tieOrder(s.C(dbaccount.FieldID))) + }} + } field := dbaccount.FieldName defaultOrder := true @@ -934,6 +948,40 @@ func accountListOrder(params pagination.PaginationParams) []func(*entsql.Selecto return []func(*entsql.Selector){dbent.Asc(field), dbent.Asc(dbaccount.FieldID)} } +func upstreamBillingRateSortExpression(extra string) string { + status := extra + " #>> '{upstream_billing_probe,status}'" + effectiveJSON := extra + " #> '{upstream_billing_probe,data,effective_rate_multiplier}'" + effective := extra + " #>> '{upstream_billing_probe,data,effective_rate_multiplier}'" + resolvedJSON := extra + " #> '{upstream_billing_probe,data,resolved_rate_multiplier}'" + resolved := extra + " #>> '{upstream_billing_probe,data,resolved_rate_multiplier}'" + peakEnabledJSON := extra + " #> '{upstream_billing_probe,data,peak_rate_enabled}'" + peakEnabled := extra + " #>> '{upstream_billing_probe,data,peak_rate_enabled}'" + peakStart := extra + " #>> '{upstream_billing_probe,data,peak_start}'" + peakEnd := extra + " #>> '{upstream_billing_probe,data,peak_end}'" + peakMultiplierJSON := extra + " #> '{upstream_billing_probe,data,peak_rate_multiplier}'" + peakMultiplier := extra + " #>> '{upstream_billing_probe,data,peak_rate_multiplier}'" + peakMultiplierValue := "(CASE WHEN jsonb_typeof(" + peakMultiplierJSON + ") = 'number' THEN (" + peakMultiplier + ")::numeric END)" + billingScope := extra + " #>> '{upstream_billing_probe,data,billing_scope}'" + timezone := extra + " #>> '{upstream_billing_probe,data,timezone}'" + validClock := "'^([01][0-9]|2[0-3]):[0-5][0-9]$'" + startMinute := "(CASE WHEN " + peakStart + " ~ " + validClock + " THEN split_part(" + peakStart + ", ':', 1)::numeric * 60 + split_part(" + peakStart + ", ':', 2)::numeric END)" + endMinute := "(CASE WHEN " + peakEnd + " ~ " + validClock + " THEN split_part(" + peakEnd + ", ':', 1)::numeric * 60 + split_part(" + peakEnd + ", ':', 2)::numeric END)" + localMinute := "(EXTRACT(HOUR FROM (CURRENT_TIMESTAMP AT TIME ZONE (" + timezone + "))) * 60 + EXTRACT(MINUTE FROM (CURRENT_TIMESTAMP AT TIME ZONE (" + timezone + "))))" + validPeakWindow := peakStart + " ~ " + validClock + " AND " + + peakEnd + " ~ " + validClock + " AND " + + startMinute + " < " + endMinute + validPeakConfig := validPeakWindow + " AND " + peakMultiplierValue + " >= 0 AND " + + "EXISTS (SELECT 1 FROM pg_timezone_names WHERE name = " + timezone + ")" + dynamicRate := "CASE WHEN " + peakEnabled + " = 'false' THEN (" + resolved + ")::numeric WHEN " + peakEnabled + " = 'true' AND " + validPeakConfig + + " THEN (" + resolved + ")::numeric * CASE WHEN " + localMinute + " >= " + startMinute + " AND " + localMinute + " < " + endMinute + + " THEN " + peakMultiplierValue + " ELSE 1 END ELSE NULL END" + legacySnapshot := "jsonb_typeof(" + resolvedJSON + ") IS NULL AND jsonb_typeof(" + peakEnabledJSON + ") IS NULL" + + return "CASE WHEN " + status + " IN ('ok', 'failed') AND (jsonb_typeof(" + resolvedJSON + ") = 'number' OR jsonb_typeof(" + effectiveJSON + ") = 'number') THEN CASE WHEN jsonb_typeof(" + + resolvedJSON + ") = 'number' AND jsonb_typeof(" + peakEnabledJSON + ") = 'boolean' THEN CASE WHEN " + billingScope + " = 'token' THEN " + dynamicRate + " ELSE NULL END WHEN " + legacySnapshot + + " AND jsonb_typeof(" + effectiveJSON + ") = 'number' THEN (" + effective + ")::numeric END END" +} + func (r *accountRepository) ListByGroup(ctx context.Context, groupID int64) ([]service.Account, error) { accounts, err := r.queryAccountsByGroup(ctx, groupID, accountGroupQueryOptions{ status: service.StatusActive, @@ -2575,6 +2623,12 @@ func (r *accountRepository) BulkUpdate(ctx context.Context, ids []int64, updates args = append(args, *updates.Schedulable) idx++ } + if updates.ProbeEnabled != nil { + if updates.Extra == nil { + updates.Extra = make(map[string]any) + } + updates.Extra[service.UpstreamBillingProbeEnabledExtraKey] = *updates.ProbeEnabled + } // JSONB 需要合并而非覆盖，使用 raw SQL 保持旧行为。 if len(updates.Credentials) > 0 { payload, err := json.Marshal(updates.Credentials) @@ -2605,8 +2659,14 @@ func (r *accountRepository) BulkUpdate(ctx context.Context, ids []int64, updates setClauses = append(setClauses, "updated_at = NOW()") - query := "UPDATE accounts SET " + joinClauses(setClauses, ", ") + " WHERE id = ANY($" + itoa(idx) + ") AND deleted_at IS NULL" + whereClause := " WHERE id = ANY($" + itoa(idx) + ") AND deleted_at IS NULL" args = append(args, pq.Array(ids)) + idx++ + if updates.ProbeEnabled != nil { + whereClause += " AND platform = $" + itoa(idx) + " AND type = $" + itoa(idx+1) + args = append(args, service.PlatformOpenAI, service.AccountTypeAPIKey) + } + query := "UPDATE accounts SET " + joinClauses(setClauses, ", ") + whereClause baseCtx := ctx contextTx := dbent.TxFromContext(ctx) @@ -2635,6 +2695,20 @@ func (r *accountRepository) BulkUpdate(ctx context.Context, ids []int64, updates if err != nil { return 0, err } + if updates.ProbeEnabled != nil { + expectedRows := int64(0) + seenIDs := make(map[int64]struct{}, len(ids)) + for _, id := range ids { + if _, seen := seenIDs[id]; seen { + continue + } + seenIDs[id] = struct{}{} + expectedRows++ + } + if rows != expectedRows { + return 0, service.ErrUpstreamBillingProbeAccountInvalid + } + } if rows > 0 { payload := map[string]any{"account_ids": ids} if err := enqueueSchedulerOutbox(ctx, exec, service.SchedulerOutboxEventAccountBulkChanged, nil, nil, payload); err != nil { diff --git a/backend/internal/repository/account_repo_sort_integration_test.go b/backend/internal/repository/account_repo_sort_integration_test.go index 098dde7bb..19c974f37 100644 --- a/backend/internal/repository/account_repo_sort_integration_test.go +++ b/backend/internal/repository/account_repo_sort_integration_test.go @@ -3,6 +3,9 @@ package repository import ( + "fmt" + "time" + "github.com/Wei-Shaw/sub2api/internal/pkg/pagination" "github.com/Wei-Shaw/sub2api/internal/service" ) @@ -33,3 +36,105 @@ func (s *AccountRepoSuite) TestListWithFilters_SortByPriorityDesc() { s.Require().Equal("high-priority", accounts[0].Name) s.Require().Equal("low-priority", accounts[1].Name) } + +func (s *AccountRepoSuite) TestListWithFilters_SortByUpstreamBillingRateWithMissingLast() { + makeAccount := func(name, status string, rate any) { + extra := map[string]any{} + if rate != nil { + extra[service.UpstreamBillingProbeExtraKey] = map[string]any{ + "status": status, + "data": map[string]any{"effective_rate_multiplier": rate}, + } + } + mustCreateAccount(s.T(), s.client, &service.Account{Name: name, Extra: extra}) + } + makeAccount("high-rate", service.UpstreamBillingProbeStatusOK, 0.8) + makeAccount("low-rate", service.UpstreamBillingProbeStatusOK, 0.03) + makeAccount("missing-rate", "", nil) + makeAccount("unsupported-with-retained-rate", service.UpstreamBillingProbeStatusUnsupported, 0.01) + + for _, tc := range []struct { + order string + want []string + }{ + {order: "asc", want: []string{"low-rate", "high-rate", "missing-rate", "unsupported-with-retained-rate"}}, + {order: "desc", want: []string{"high-rate", "low-rate", "unsupported-with-retained-rate", "missing-rate"}}, + } { + accounts, _, err := s.repo.ListWithFilters(s.ctx, pagination.PaginationParams{ + Page: 1, PageSize: 10, SortBy: "upstream_billing_rate", SortOrder: tc.order, + }, "", "", "", "", 0, "") + s.Require().NoError(err) + s.Require().Len(accounts, 4) + for i, name := range tc.want { + s.Require().Equal(name, accounts[i].Name) + } + } +} + +func (s *AccountRepoSuite) TestListWithFilters_SortByCurrentUpstreamBillingRateDuringPeak() { + now := time.Now() + locations := []string{"UTC", "Asia/Shanghai", "America/New_York", "Europe/London"} + var timezone string + var minute int + for _, name := range locations { + location, err := time.LoadLocation(name) + s.Require().NoError(err) + local := now.In(location) + candidate := local.Hour()*60 + local.Minute() + if candidate >= 2 && candidate <= 1436 { + timezone = name + minute = candidate + break + } + } + s.Require().NotEmpty(timezone) + + peakStart := fmt.Sprintf("%02d:%02d", (minute-2)/60, (minute-2)%60) + peakEnd := fmt.Sprintf("%02d:%02d", (minute+3)/60, (minute+3)%60) + mustCreateAccount(s.T(), s.client, &service.Account{ + Name: "current-peak-rate", + Extra: map[string]any{ + service.UpstreamBillingProbeExtraKey: map[string]any{ + "status": service.UpstreamBillingProbeStatusOK, + "data": map[string]any{ + "billing_scope": "token", + "resolved_rate_multiplier": 1.0, + "effective_rate_multiplier": 1.0, + "peak_rate_enabled": true, + "peak_start": peakStart, + "peak_end": peakEnd, + "peak_rate_multiplier": 10.0, + "timezone": timezone, + }, + }, + }, + }) + mustCreateAccount(s.T(), s.client, &service.Account{ + Name: "current-off-peak-rate", + Extra: map[string]any{ + service.UpstreamBillingProbeExtraKey: map[string]any{ + "status": service.UpstreamBillingProbeStatusOK, + "data": map[string]any{ + "effective_rate_multiplier": 5.0, + }, + }, + }, + }) + + for _, tc := range []struct { + order string + want []string + }{ + {order: "asc", want: []string{"current-off-peak-rate", "current-peak-rate"}}, + {order: "desc", want: []string{"current-peak-rate", "current-off-peak-rate"}}, + } { + accounts, _, err := s.repo.ListWithFilters(s.ctx, pagination.PaginationParams{ + Page: 1, PageSize: 10, SortBy: "upstream_billing_rate", SortOrder: tc.order, + }, "", "", "", "", 0, "") + s.Require().NoError(err) + s.Require().Len(accounts, 2) + for i, name := range tc.want { + s.Require().Equal(name, accounts[i].Name) + } + } +} diff --git a/backend/internal/repository/account_repo_upstream_billing_probe_update_test.go b/backend/internal/repository/account_repo_upstream_billing_probe_update_test.go index 26fbb74dc..2a78f38c1 100644 --- a/backend/internal/repository/account_repo_upstream_billing_probe_update_test.go +++ b/backend/internal/repository/account_repo_upstream_billing_probe_update_test.go @@ -163,6 +163,46 @@ func TestBulkUpdateNilProbeRemovesKeyInsteadOfWritingJSONNull(t *testing.T) { require.Contains(t, normalizeSQLWhitespace(exec.execQueries[0]), "- 'upstream_billing_probe'") } +func TestBulkUpdateDisablingProbeRemovesSnapshot(t *testing.T) { + exec := &recordingSQLExecutor{result: rowsAffectedResult(1)} + repo := newAccountRepositoryWithSQL(nil, exec, nil) + + _, err := repo.BulkUpdate(context.Background(), []int64{27}, service.AccountBulkUpdate{ + Extra: map[string]any{service.UpstreamBillingProbeEnabledExtraKey: false}, + }) + + require.NoError(t, err) + require.NotEmpty(t, exec.execQueries) + require.Contains(t, normalizeSQLWhitespace(exec.execQueries[0]), "- 'upstream_billing_probe'") + payload, ok := exec.execArgs[0][0].([]byte) + require.True(t, ok) + require.Equal(t, `{"upstream_billing_probe_enabled":false}`, string(payload)) +} + +func TestBulkUpdateProbeEligibilityMismatchRollsBack(t *testing.T) { + db, mock, err := sqlmock.New() + require.NoError(t, err) + t.Cleanup(func() { _ = db.Close() }) + client := dbent.NewClient(dbent.Driver(entsql.OpenDB(dialect.Postgres, db))) + t.Cleanup(func() { _ = client.Close() }) + + enabled := true + mock.ExpectBegin() + mock.ExpectExec(`(?s)UPDATE accounts SET extra = .* WHERE id = ANY\(\$2\) AND deleted_at IS NULL AND platform = \$3 AND type = \$4`). + WithArgs(sqlmock.AnyArg(), `{27,28}`, service.PlatformOpenAI, service.AccountTypeAPIKey). + WillReturnResult(sqlmock.NewResult(0, 1)) + mock.ExpectRollback() + + repo := newAccountRepositoryWithSQL(client, db, nil) + rows, err := repo.BulkUpdate(context.Background(), []int64{27, 28}, service.AccountBulkUpdate{ + ProbeEnabled: &enabled, + }) + + require.ErrorIs(t, err, service.ErrUpstreamBillingProbeAccountInvalid) + require.Zero(t, rows) + require.NoError(t, mock.ExpectationsWereMet()) +} + func TestUpdateCredentialsAtomicallyClearsProbeForOpenAIAPIKeyIdentityChange(t *testing.T) { db, mock, err := sqlmock.New() require.NoError(t, err) diff --git a/backend/internal/service/account_service.go b/backend/internal/service/account_service.go index b7fe0a68c..d1a57e6d2 100644 --- a/backend/internal/service/account_service.go +++ b/backend/internal/service/account_service.go @@ -144,6 +144,7 @@ type AccountBulkUpdate struct { Schedulable *bool Credentials map[string]any Extra map[string]any + ProbeEnabled *bool } // CreateAccountRequest 创建账号请求 diff --git a/backend/internal/service/admin_account.go b/backend/internal/service/admin_account.go index 899679221..ae3a9da36 100644 --- a/backend/internal/service/admin_account.go +++ b/backend/internal/service/admin_account.go @@ -469,6 +469,15 @@ func buildAccountForCreate(input *CreateAccountInput, accountExtra map[string]an Status: StatusActive, Schedulable: true, } + if input.ProbeEnabled != nil && *input.ProbeEnabled { + if !isUpstreamBillingProbeAccount(account) { + return nil, ErrUpstreamBillingProbeAccountInvalid + } + if account.Extra == nil { + account.Extra = make(map[string]any) + } + account.Extra[UpstreamBillingProbeEnabledExtraKey] = true + } // 预计算固定时间重置的下次重置时间 if account.Extra != nil { if err := ValidateQuotaResetConfig(account.Extra); err != nil { @@ -842,7 +851,7 @@ func (s *adminServiceImpl) UpdateAccountExtra(ctx context.Context, id int64, upd // BulkUpdateAccounts updates multiple accounts in one request. // It merges credentials/extra keys instead of overwriting the whole object. func (s *adminServiceImpl) BulkUpdateAccounts(ctx context.Context, input *BulkUpdateAccountsInput) (*BulkUpdateAccountsResult, error) { - // Probe state is updated only through its dedicated endpoints. + // Managed probe state may only enter through the dedicated typed field below. delete(input.Extra, UpstreamBillingProbeEnabledExtraKey) delete(input.Extra, UpstreamBillingProbeExtraKey) @@ -874,13 +883,30 @@ func (s *adminServiceImpl) BulkUpdateAccounts(ctx context.Context, input *BulkUp // 预取所有目标账号，供凭据守卫/代理守卫/混合渠道检查共用，避免多次 DB 查询。 var cachedTargets []*Account - if len(input.Credentials) > 0 || input.ProxyID != nil || needMixedChannelCheck || hasLongContextBillingUpdate { + if len(input.Credentials) > 0 || input.ProxyID != nil || needMixedChannelCheck || hasLongContextBillingUpdate || input.ProbeEnabled != nil { loaded, err := s.accountRepo.GetByIDs(ctx, input.AccountIDs) if err != nil { return nil, err } cachedTargets = loaded } + if input.ProbeEnabled != nil { + targetsByID := make(map[int64]*Account, len(cachedTargets)) + for _, account := range cachedTargets { + if account != nil { + targetsByID[account.ID] = account + } + } + for _, accountID := range input.AccountIDs { + account, ok := targetsByID[accountID] + if !ok { + return nil, ErrAccountNotFound + } + if !isUpstreamBillingProbeAccount(account) { + return nil, ErrUpstreamBillingProbeAccountInvalid + } + } + } if hasLongContextBillingUpdate { for _, account := range cachedTargets { if account == nil || account.Platform != PlatformOpenAI { @@ -952,8 +978,15 @@ func (s *adminServiceImpl) BulkUpdateAccounts(ctx context.Context, input *BulkUp // Prepare bulk updates for columns and JSONB fields. repoUpdates := AccountBulkUpdate{ - Credentials: input.Credentials, - Extra: input.Extra, + Credentials: input.Credentials, + Extra: input.Extra, + ProbeEnabled: input.ProbeEnabled, + } + if input.ProbeEnabled != nil { + if repoUpdates.Extra == nil { + repoUpdates.Extra = make(map[string]any) + } + repoUpdates.Extra[UpstreamBillingProbeEnabledExtraKey] = *input.ProbeEnabled } if updatesUpstreamBillingProbeIdentity(input.Credentials) || input.ProxyID != nil { if repoUpdates.Extra == nil { diff --git a/backend/internal/service/admin_account_upstream_billing_probe_test.go b/backend/internal/service/admin_account_upstream_billing_probe_test.go index 727a66cf1..88db65386 100644 --- a/backend/internal/service/admin_account_upstream_billing_probe_test.go +++ b/backend/internal/service/admin_account_upstream_billing_probe_test.go @@ -38,6 +38,32 @@ func TestCreateAccountDropsManagedUpstreamBillingProbeState(t *testing.T) { require.NotContains(t, created.Extra, UpstreamBillingProbeExtraKey) } +func TestCreateAccountAcceptsDedicatedUpstreamBillingProbeSetting(t *testing.T) { + enabled := true + repo := &upstreamBillingProbeAccountRepo{} + created, err := (&adminServiceImpl{accountRepo: repo}).CreateAccount(context.Background(), &CreateAccountInput{ + Name: "upstream", + Platform: PlatformOpenAI, + Type: AccountTypeAPIKey, + Credentials: map[string]any{"api_key": "sk-test"}, + ProbeEnabled: &enabled, + SkipDefaultGroupBind: true, + }) + + require.NoError(t, err) + require.Equal(t, true, created.Extra[UpstreamBillingProbeEnabledExtraKey]) + + _, err = (&adminServiceImpl{accountRepo: repo}).CreateAccount(context.Background(), &CreateAccountInput{ + Name: "oauth", + Platform: PlatformOpenAI, + Type: AccountTypeOAuth, + Credentials: map[string]any{"access_token": "token"}, + ProbeEnabled: &enabled, + SkipDefaultGroupBind: true, + }) + require.ErrorIs(t, err, ErrUpstreamBillingProbeAccountInvalid) +} + func TestUpdateAccountPreservesManagedUpstreamBillingProbeStateForUnrelatedEdit(t *testing.T) { accountID := int64(110) repo := &upstreamBillingProbeAccountRepo{accounts: map[int64]*Account{ @@ -360,6 +386,63 @@ func TestBulkUpdateAccountsDropsManagedUpstreamBillingProbeState(t *testing.T) { require.NotContains(t, repo.bulkUpdates[0].Extra, UpstreamBillingProbeExtraKey) } +func TestBulkUpdateAccountsAcceptsDedicatedUpstreamBillingProbeSetting(t *testing.T) { + for _, enabled := range []bool{true, false} { + t.Run(map[bool]string{true: "enable", false: "disable"}[enabled], func(t *testing.T) { + repo := &upstreamBillingProbeAccountRepo{accounts: map[int64]*Account{ + 1: {ID: 1, Platform: PlatformOpenAI, Type: AccountTypeAPIKey}, + 2: {ID: 2, Platform: PlatformOpenAI, Type: AccountTypeAPIKey}, + }} + + result, err := (&adminServiceImpl{accountRepo: repo}).BulkUpdateAccounts(context.Background(), &BulkUpdateAccountsInput{ + AccountIDs: []int64{1, 2}, + ProbeEnabled: &enabled, + }) + + require.NoError(t, err) + require.Equal(t, 2, result.Success) + require.Len(t, repo.bulkUpdates, 1) + require.Equal(t, enabled, repo.bulkUpdates[0].Extra[UpstreamBillingProbeEnabledExtraKey]) + require.NotNil(t, repo.bulkUpdates[0].ProbeEnabled) + require.Equal(t, enabled, *repo.bulkUpdates[0].ProbeEnabled) + }) + } +} + +func TestBulkUpdateAccountsRejectsProbeSettingForIneligibleTargetBeforeWrite(t *testing.T) { + for _, enabled := range []bool{true, false} { + t.Run(map[bool]string{true: "enable", false: "disable"}[enabled], func(t *testing.T) { + repo := &upstreamBillingProbeAccountRepo{accounts: map[int64]*Account{ + 1: {ID: 1, Platform: PlatformOpenAI, Type: AccountTypeAPIKey}, + 2: {ID: 2, Platform: PlatformOpenAI, Type: AccountTypeOAuth}, + }} + + _, err := (&adminServiceImpl{accountRepo: repo}).BulkUpdateAccounts(context.Background(), &BulkUpdateAccountsInput{ + AccountIDs: []int64{1, 2}, + ProbeEnabled: &enabled, + }) + + require.ErrorIs(t, err, ErrUpstreamBillingProbeAccountInvalid) + require.Empty(t, repo.bulkUpdates) + }) + } +} + +func TestBulkUpdateAccountsRejectsProbeSettingWhenTargetIsMissing(t *testing.T) { + enabled := true + repo := &upstreamBillingProbeAccountRepo{accounts: map[int64]*Account{ + 1: {ID: 1, Platform: PlatformOpenAI, Type: AccountTypeAPIKey}, + }} + + _, err := (&adminServiceImpl{accountRepo: repo}).BulkUpdateAccounts(context.Background(), &BulkUpdateAccountsInput{ + AccountIDs: []int64{1, 2}, + ProbeEnabled: &enabled, + }) + + require.ErrorIs(t, err, ErrAccountNotFound) + require.Empty(t, repo.bulkUpdates) +} + func TestBulkUpdateAccountsInvalidatesProbeSnapshotForIdentityCredentials(t *testing.T) { repo := &upstreamBillingProbeAccountRepo{} input := &BulkUpdateAccountsInput{ diff --git a/backend/internal/service/admin_service.go b/backend/internal/service/admin_service.go index c167b6be8..013cb0404 100644 --- a/backend/internal/service/admin_service.go +++ b/backend/internal/service/admin_service.go @@ -329,6 +329,7 @@ type CreateAccountInput struct { GroupIDs []int64 ExpiresAt *int64 AutoPauseOnExpired *bool + ProbeEnabled *bool // SkipDefaultGroupBind prevents auto-binding to platform default group when GroupIDs is empty. SkipDefaultGroupBind bool // SkipMixedChannelCheck skips the mixed channel risk check when binding groups. @@ -378,6 +379,7 @@ type BulkUpdateAccountsInput struct { GroupIDs *[]int64 Credentials map[string]any Extra map[string]any + ProbeEnabled *bool // SkipMixedChannelCheck skips the mixed channel risk check when binding groups. // This should only be set when the caller has explicitly confirmed the risk. SkipMixedChannelCheck bool diff --git a/backend/internal/service/upstream_billing_probe.go b/backend/internal/service/upstream_billing_probe.go index 70dbf8f8f..d5c28e68b 100644 --- a/backend/internal/service/upstream_billing_probe.go +++ b/backend/internal/service/upstream_billing_probe.go @@ -38,7 +38,7 @@ const ( upstreamBillingProbeMaxBodyBytes = 64 * 1024 upstreamBillingProbeMaxPerCycle = 20 upstreamBillingProbeConcurrency = 4 - upstreamBillingProbeMaxBackoff = 24 * time.Hour + upstreamBillingProbeMaxDelay = 24 * time.Hour upstreamBillingProbeLeaderLockKey = "upstream:billing:probe:leader" upstreamBillingProbeLeaderLockTTL = 2 * time.Minute ) @@ -620,7 +620,7 @@ func (s *UpstreamBillingProbeService) probeLoadedAccount(ctx context.Context, ac ReceivedAt: probeTimePtr(now), FreshUntil: probeTimePtr(now.Add(2 * time.Duration(intervalMinutes) * time.Minute)), LastAttemptAt: now, - NextProbeAt: now.Add(nextProbeDelay(intervalMinutes, 0, 0)), + NextProbeAt: now.Add(nextProbeDelay(intervalMinutes, 0)), HTTPStatus: resp.StatusCode, } if err := s.updateSnapshot(ctx, account, snapshot); err != nil { @@ -650,7 +650,7 @@ func (s *UpstreamBillingProbeService) persistProbeFailure( snapshot := &UpstreamBillingProbeSnapshot{ Status: status, LastAttemptAt: now, - NextProbeAt: now.Add(nextProbeDelay(intervalMinutes, failureCount, retryAfterDuration)), + NextProbeAt: now.Add(nextProbeDelay(intervalMinutes, retryAfterDuration)), FailureCount: failureCount, HTTPStatus: statusCode, LastError: reason, @@ -859,20 +859,13 @@ func (s *UpstreamBillingProbeService) currentTime() time.Time { return time.Now() } -func nextProbeDelay(intervalMinutes, failureCount int, retryAfterDuration time.Duration) time.Duration { +func nextProbeDelay(intervalMinutes int, retryAfterDuration time.Duration) time.Duration { interval := time.Duration(intervalMinutes) * time.Minute if interval < upstreamBillingProbeMinIntervalMinutes*time.Minute { interval = upstreamBillingProbeMinIntervalMinutes * time.Minute } - if failureCount > 0 { - shift := failureCount - if shift > 5 { - shift = 5 - } - interval *= time.Duration(1 << shift) - } - if interval > upstreamBillingProbeMaxBackoff { - interval = upstreamBillingProbeMaxBackoff + if interval > upstreamBillingProbeMaxDelay { + interval = upstreamBillingProbeMaxDelay } jitterRange := interval / 5 if jitterRange > 5*time.Minute { @@ -883,11 +876,11 @@ func nextProbeDelay(intervalMinutes, failureCount int, retryAfterDuration time.D } if retryAfterDuration > interval { // Retry-After is an explicit upstream instruction; do not shorten it - // with the local exponential-backoff ceiling. + // with the local maximum delay. return retryAfterDuration } - if interval > upstreamBillingProbeMaxBackoff { - return upstreamBillingProbeMaxBackoff + if interval > upstreamBillingProbeMaxDelay { + return upstreamBillingProbeMaxDelay } return interval } diff --git a/backend/internal/service/upstream_billing_probe_test.go b/backend/internal/service/upstream_billing_probe_test.go index 0ffcc6006..eecd866a9 100644 --- a/backend/internal/service/upstream_billing_probe_test.go +++ b/backend/internal/service/upstream_billing_probe_test.go @@ -495,7 +495,7 @@ func TestUpstreamBillingProbeFailurePreservesLastSuccessAndRetryAfter(t *testing } func TestUpstreamBillingProbeRetryAfterIsNotShortened(t *testing.T) { - delay := nextProbeDelay(30, 1, 48*time.Hour) + delay := nextProbeDelay(30, 48*time.Hour) require.Equal(t, 48*time.Hour, delay) } @@ -510,11 +510,21 @@ func TestUpstreamBillingProbeEmptyResponseIsPersistedAsFailure(t *testing.T) { } repo := &upstreamBillingProbeAccountRepo{accounts: map[int64]*Account{account.ID: account}} svc := newUpstreamBillingProbeTestService(repo, &httpUpstreamRecorder{}, &upstreamBillingProbeSettingRepo{}) + fixedNow := time.Date(2026, time.July, 13, 2, 0, 0, 0, time.UTC) + svc.now = func() time.Time { return fixedNow } snapshot, err := svc.ProbeAccount(context.Background(), account.ID) require.NoError(t, err) require.Equal(t, UpstreamBillingProbeStatusFailed, snapshot.Status) require.Equal(t, "empty_response", snapshot.LastError) + require.False(t, snapshot.NextProbeAt.Before(fixedNow.Add(24*time.Minute))) + require.False(t, snapshot.NextProbeAt.After(fixedNow.Add(36*time.Minute))) + + snapshot, err = svc.ProbeAccount(context.Background(), account.ID) + require.NoError(t, err) + require.Equal(t, 2, snapshot.FailureCount) + require.False(t, snapshot.NextProbeAt.Before(fixedNow.Add(24*time.Minute))) + require.False(t, snapshot.NextProbeAt.After(fixedNow.Add(36*time.Minute))) } func TestUpstreamBillingProbeUnsupportedAndAccountToggle(t *testing.T) { @@ -533,6 +543,8 @@ func TestUpstreamBillingProbeUnsupportedAndAccountToggle(t *testing.T) { Body: io.NopCloser(strings.NewReader("not found")), }} svc := newUpstreamBillingProbeTestService(repo, upstream, &upstreamBillingProbeSettingRepo{}) + fixedNow := time.Date(2026, time.July, 13, 2, 0, 0, 0, time.UTC) + svc.now = func() time.Time { return fixedNow } require.NoError(t, svc.SetAccountEnabled(context.Background(), account.ID, true)) require.Equal(t, true, account.Extra[UpstreamBillingProbeEnabledExtraKey]) @@ -540,6 +552,14 @@ func TestUpstreamBillingProbeUnsupportedAndAccountToggle(t *testing.T) { require.NoError(t, err) require.Equal(t, UpstreamBillingProbeStatusUnsupported, snapshot.Status) require.Equal(t, "unsupported", snapshot.LastError) + require.False(t, snapshot.NextProbeAt.Before(fixedNow.Add(24*time.Minute))) + require.False(t, snapshot.NextProbeAt.After(fixedNow.Add(36*time.Minute))) + + snapshot, err = svc.ProbeAccount(context.Background(), account.ID) + require.NoError(t, err) + require.Equal(t, 2, snapshot.FailureCount) + require.False(t, snapshot.NextProbeAt.Before(fixedNow.Add(24*time.Minute))) + require.False(t, snapshot.NextProbeAt.After(fixedNow.Add(36*time.Minute))) invalid := &Account{ID: 20, Platform: PlatformOpenAI, Type: AccountTypeOAuth} repo.accounts[invalid.ID] = invalid diff --git a/frontend/src/components/account/BulkEditAccountModal.vue b/frontend/src/components/account/BulkEditAccountModal.vue index 5f744c0fc..a4268fb92 100644 --- a/frontend/src/components/account/BulkEditAccountModal.vue +++ b/frontend/src/components/account/BulkEditAccountModal.vue @@ -854,6 +854,45 @@ + +
+
+
+ +

+ {{ t('admin.accounts.upstreamBilling.autoProbeHint') }} +

+
+ +
+
+
-
+