From 98847adfd6aeaa3ef1eb1f93b24c174a5d671778 Mon Sep 17 00:00:00 2001 From: dev Date: Mon, 24 Aug 2026 22:38:38 +0800 Subject: [PATCH] refactor(gateway): collapse chat.go four-way duplication MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit singleChat / singleChatAuto / streamChat / streamChatAuto shared ~260 near- identical lines (diff after stripping comments was empty). Extract three shared bodies and shrink all four entry points to thin dispatchers: - failChat: error → record + writeError, shared ChainErr extraction (errors.As returns false for direct-path errors, so the two paths stay equivalent without a branch) - recordChatUsage: exact upstream numbers win, byte estimates fill gaps - writeChatCompletion: unified ChatCompletion rendering (model name is the only direct/auto difference, passed in) - pumpStream: the full SSE pump (preamble, delta loop, terminal finish + usage chunk, [DONE]) — shared by both streaming entries Behavior change (pinned by TestDirectStreamFailoverAuditSource): direct streams now pin rec.Source to the source that actually served the stream after a failover, instead of discarding it (_, usedModel). The audit row previously recorded the first candidate, which was wrong on failover. chat.go: 982 → 896 lines (−86). --- internal/gateway/chat.go | 264 +++++++++++-------------------- internal/gateway/gateway_test.go | 37 +++++ 2 files changed, 126 insertions(+), 175 deletions(-) diff --git a/internal/gateway/chat.go b/internal/gateway/chat.go index edcf270..5b7e96f 100644 --- a/internal/gateway/chat.go +++ b/internal/gateway/chat.go @@ -542,21 +542,26 @@ func clientUpstreamErr(err error) string { return types.OneLine(err.Error(), 160) } -func (g *Gateway) singleChat(w http.ResponseWriter, ctx context.Context, cands []*provider.Provider, req *types.ChatRequest, effective string, rec *Req) { - rec.LatMs = 0 - t0 := time.Now() - resp, usedSrc, usedModel, err := g.core.Scheduler().Chat(ctx, toScheduler(cands), req) - rec.LatMs = time.Since(t0).Milliseconds() - if err != nil { - rec.OK = false - rec.Status = upstreamErrStatus(err) - rec.Err = err.Error() - g.writeRec(rec) - writeError(w, rec.Status, "upstream_error", clientUpstreamErr(err)) - return +// failChat maps a scheduling failure onto the audit record and answers the +// client. A failed AUTO chain additionally pins its first failed tier onto +// the record; direct-path errors never match *scheduler.ChainErr, so the +// extraction is safely shared by all four entry points. Callers own the +// writeRec call (inline for non-stream paths, deferred for stream paths). +func (g *Gateway) failChat(w http.ResponseWriter, rec *Req, err error) { + rec.OK = false + rec.Status = upstreamErrStatus(err) + rec.Err = err.Error() + var ce *scheduler.ChainErr + if errors.As(err, &ce) && len(ce.Tiers) > 0 { + rec.Source = ce.Tiers[0].Source + rec.Model = ce.Tiers[0].Model } - rec.OK = true - rec.Status = http.StatusOK + writeError(w, rec.Status, "upstream_error", clientUpstreamErr(err)) +} + +// recordChatUsage fills token accounting for a finished non-streaming +// request: exact upstream numbers win, byte estimates fill the gaps. +func recordChatUsage(rec *Req, req *types.ChatRequest, resp *types.UnifiedResponse) { rec.Prompt = int64(resp.TokenUsage.Prompt) if rec.Prompt == 0 { rec.Prompt = estimatePromptTokens(req) @@ -565,9 +570,12 @@ func (g *Gateway) singleChat(w http.ResponseWriter, ctx context.Context, cands [ if rec.Compl == 0 { rec.Compl = estimateTextTokens(resp.Content, resp.ReasoningContent, resp.ToolCalls) } - rec.Source = usedSrc - rec.Model = usedModel - g.writeRec(rec) +} + +// writeChatCompletion renders a unified response as an OpenAI +// chat.completion object. modelName is the id clients see as the serving +// model: the requested id for direct routes, the exact slot model for AUTO. +func writeChatCompletion(w http.ResponseWriter, resp *types.UnifiedResponse, modelName string) { msg := RespMessage{Role: "assistant", Content: resp.Content} if resp.ReasoningContent != "" { msg.ReasoningContent = resp.ReasoningContent @@ -579,7 +587,7 @@ func (g *Gateway) singleChat(w http.ResponseWriter, ctx context.Context, cands [ ID: newID(), Object: "chat.completion", Created: time.Now().Unix(), - Model: effective, + Model: modelName, Choices: []ChatChoice{{Index: 0, Message: msg, FinishReason: resp.FinishReason}}, } if resp.TokenUsage.Total > 0 || resp.TokenUsage.Prompt > 0 || resp.TokenUsage.Completion > 0 { @@ -588,6 +596,27 @@ func (g *Gateway) singleChat(w http.ResponseWriter, ctx context.Context, cands [ writeJSON(w, http.StatusOK, out) } +// singleChat runs a direct (model-pinned) non-streaming request across the +// candidate list, falling back on failure. +func (g *Gateway) singleChat(w http.ResponseWriter, ctx context.Context, cands []*provider.Provider, req *types.ChatRequest, effective string, rec *Req) { + rec.LatMs = 0 + t0 := time.Now() + resp, usedSrc, usedModel, err := g.core.Scheduler().Chat(ctx, toScheduler(cands), req) + rec.LatMs = time.Since(t0).Milliseconds() + if err != nil { + g.failChat(w, rec, err) + g.writeRec(rec) + return + } + rec.OK = true + rec.Status = http.StatusOK + recordChatUsage(rec, req, resp) + rec.Source = usedSrc + rec.Model = usedModel + g.writeRec(rec) + writeChatCompletion(w, resp, effective) +} + // writeRec records a finished request (audit + aggregates). func (g *Gateway) writeRec(rec *Req) { if rec == nil { @@ -625,27 +654,12 @@ func mergeUsage(prev, cur *types.TokenUsage) *types.TokenUsage { return &out } -func (g *Gateway) streamChat(w http.ResponseWriter, ctx context.Context, cands []*provider.Provider, req *types.ChatRequest, effective string, rec *Req) { - rec.LatMs = 0 - t0 := time.Now() - rec.OK = true - rec.Status = http.StatusOK - defer func() { - rec.LatMs = time.Since(t0).Milliseconds() - g.writeRec(rec) - }() - chunks, _, usedModel, err := g.core.Scheduler().ChatStream(ctx, toScheduler(cands), req) - if err != nil { - rec.OK = false - rec.Status = upstreamErrStatus(err) - rec.Err = err.Error() - writeError(w, rec.Status, "upstream_error", clientUpstreamErr(err)) - return - } - if usedModel != "" { - rec.Model = usedModel - } - rec.Prompt = estimatePromptTokens(req) +// pumpStream writes the full SSE sequence for a started stream: role +// preamble, one chunk per unified delta, the terminating finish_reason, the +// OpenAI-standard final usage chunk (empty choices) and [DONE]. modelName +// follows writeChatCompletion's rule (requested id for direct routes, exact +// slot model for AUTO). +func (g *Gateway) pumpStream(w http.ResponseWriter, rec *Req, chunks <-chan types.UnifiedChunk, modelName string) { w.Header().Set("Content-Type", "text/event-stream") w.Header().Set("Cache-Control", "no-cache") w.Header().Set("Connection", "keep-alive") @@ -669,7 +683,7 @@ func (g *Gateway) streamChat(w http.ResponseWriter, ctx context.Context, cands [ } if !send(ChatChunk{ - ID: id, Object: "chat.completion.chunk", Created: created, Model: effective, + ID: id, Object: "chat.completion.chunk", Created: created, Model: modelName, Choices: []ChunkChoice{{Index: 0, Delta: RespMessage{Role: "assistant"}}}, }) { return @@ -681,7 +695,7 @@ func (g *Gateway) streamChat(w http.ResponseWriter, ctx context.Context, cands [ lastUsage = mergeUsage(lastUsage, ck.Usage) } chunk := ChatChunk{ - ID: id, Object: "chat.completion.chunk", Created: created, Model: effective, + ID: id, Object: "chat.completion.chunk", Created: created, Model: modelName, } delta := RespMessage{Role: "assistant", Content: ck.Content} if ck.ReasoningContent != "" { @@ -710,7 +724,7 @@ func (g *Gateway) streamChat(w http.ResponseWriter, ctx context.Context, cands [ finalFinish = "stop" } send(ChatChunk{ - ID: id, Object: "chat.completion.chunk", Created: created, Model: effective, + ID: id, Object: "chat.completion.chunk", Created: created, Model: modelName, Choices: []ChunkChoice{{Index: 0, Delta: RespMessage{}, FinishReason: &finalFinish}}, }) // Final usage chunk (OpenAI standard: empty choices + usage before [DONE]). @@ -729,7 +743,7 @@ func (g *Gateway) streamChat(w http.ResponseWriter, ctx context.Context, cands [ } if tut != nil { send(ChatChunk{ - ID: id, Object: "chat.completion.chunk", Created: created, Model: effective, + ID: id, Object: "chat.completion.chunk", Created: created, Model: modelName, Choices: []ChunkChoice{}, Usage: tut, }) @@ -740,6 +754,33 @@ func (g *Gateway) streamChat(w http.ResponseWriter, ctx context.Context, cands [ } } +// streamChat runs a direct (model-pinned) streaming request across the +// candidate list, falling back early on connect errors. +func (g *Gateway) streamChat(w http.ResponseWriter, ctx context.Context, cands []*provider.Provider, req *types.ChatRequest, effective string, rec *Req) { + rec.LatMs = 0 + t0 := time.Now() + rec.OK = true + rec.Status = http.StatusOK + defer func() { + rec.LatMs = time.Since(t0).Milliseconds() + g.writeRec(rec) + }() + chunks, usedSrc, usedModel, err := g.core.Scheduler().ChatStream(ctx, toScheduler(cands), req) + if err != nil { + g.failChat(w, rec, err) + return + } + if usedModel != "" { + rec.Model = usedModel + } + // Audit accuracy: pin the source that actually served the stream (after a + // failover it differs from the first candidate). Direct streams previously + // discarded it. + rec.Source = usedSrc + rec.Prompt = estimatePromptTokens(req) + g.pumpStream(w, rec, chunks, effective) +} + // singleChatAuto runs a non-streaming AUTO request down the chain (see // scheduler.ChainChat): tiers ascending (tier 1 = highest priority first), // per-tier round-robin ordered by preference, cooldown as the only hard skip, @@ -752,48 +793,17 @@ func (g *Gateway) singleChatAuto(w http.ResponseWriter, ctx context.Context, cha resp, usedSrc, usedModel, err := g.core.Scheduler().ChainChat(ctx, chain, req, quotaExhausted) rec.LatMs = time.Since(t0).Milliseconds() if err != nil { - rec.OK = false - rec.Status = upstreamErrStatus(err) - rec.Err = err.Error() - if ce, ok := err.(*scheduler.ChainErr); ok && len(ce.Tiers) > 0 { - rec.Source = ce.Tiers[0].Source - rec.Model = ce.Tiers[0].Model - } + g.failChat(w, rec, err) g.writeRec(rec) - writeError(w, rec.Status, "upstream_error", clientUpstreamErr(err)) return } rec.OK = true rec.Status = http.StatusOK - rec.Prompt = int64(resp.TokenUsage.Prompt) - if rec.Prompt == 0 { - rec.Prompt = estimatePromptTokens(req) - } - rec.Compl = int64(resp.TokenUsage.Completion) - if rec.Compl == 0 { - rec.Compl = estimateTextTokens(resp.Content, resp.ReasoningContent, resp.ToolCalls) - } + recordChatUsage(rec, req, resp) rec.Source = usedSrc rec.Model = usedModel g.writeRec(rec) - msg := RespMessage{Role: "assistant", Content: resp.Content} - if resp.ReasoningContent != "" { - msg.ReasoningContent = resp.ReasoningContent - } - if len(resp.ToolCalls) > 0 { - msg.ToolCalls = toolCallsWire(resp.ToolCalls) - } - out := ChatCompletion{ - ID: newID(), - Object: "chat.completion", - Created: time.Now().Unix(), - Model: usedModel, - Choices: []ChatChoice{{Index: 0, Message: msg, FinishReason: resp.FinishReason}}, - } - if resp.TokenUsage.Total > 0 || resp.TokenUsage.Prompt > 0 || resp.TokenUsage.Completion > 0 { - out.Usage = &resp.TokenUsage - } - writeJSON(w, http.StatusOK, out) + writeChatCompletion(w, resp, usedModel) } // streamChatAuto streams an AUTO request down the chain. A slot is abandoned @@ -811,14 +821,7 @@ func (g *Gateway) streamChatAuto(w http.ResponseWriter, ctx context.Context, cha }() chunks, usedSrc, usedModel, err := g.core.Scheduler().ChainChatStream(ctx, chain, req, quotaExhausted) if err != nil { - rec.OK = false - rec.Status = upstreamErrStatus(err) - rec.Err = err.Error() - if ce, ok := err.(*scheduler.ChainErr); ok && len(ce.Tiers) > 0 { - rec.Source = ce.Tiers[0].Source - rec.Model = ce.Tiers[0].Model - } - writeError(w, rec.Status, "upstream_error", clientUpstreamErr(err)) + g.failChat(w, rec, err) return } if usedModel != "" { @@ -826,96 +829,7 @@ func (g *Gateway) streamChatAuto(w http.ResponseWriter, ctx context.Context, cha } rec.Source = usedSrc rec.Prompt = estimatePromptTokens(req) - w.Header().Set("Content-Type", "text/event-stream") - w.Header().Set("Cache-Control", "no-cache") - w.Header().Set("Connection", "keep-alive") - w.WriteHeader(http.StatusOK) - flusher, _ := w.(http.Flusher) - id := newID() - created := time.Now().Unix() - send := func(obj interface{}) bool { - b, err := json.Marshal(obj) - if err != nil { - return false - } - if _, err := fmt.Fprintf(w, "data: %s\n\n", b); err != nil { - return false - } - if flusher != nil { - flusher.Flush() - } - return true - } - if !send(ChatChunk{ - ID: id, Object: "chat.completion.chunk", Created: created, Model: rec.Model, - Choices: []ChunkChoice{{Index: 0, Delta: RespMessage{Role: "assistant"}}}, - }) { - return - } - var lastUsage *types.TokenUsage - lastFinish := "" - for ck := range chunks { - if ck.Usage != nil { - lastUsage = mergeUsage(lastUsage, ck.Usage) - } - chunk := ChatChunk{ - ID: id, Object: "chat.completion.chunk", Created: created, Model: rec.Model, - } - delta := RespMessage{Role: "assistant", Content: ck.Content} - if ck.ReasoningContent != "" { - delta.ReasoningContent = ck.ReasoningContent - } - if len(ck.ToolCalls) > 0 { - delta.ToolCalls = ck.ToolCalls - } - choice := ChunkChoice{Index: 0, Delta: delta} - if ck.Done { - fin := ck.FinishReason - if fin == "" { - fin = "stop" - } - lastFinish = fin - choice.FinishReason = &fin - } - chunk.Choices = []ChunkChoice{choice} - rec.Compl += int64(len(ck.Content)+len(ck.ReasoningContent)+len(ck.ToolCalls)) / 3 - if !send(chunk) { - return - } - } - finalFinish := lastFinish - if finalFinish == "" { - finalFinish = "stop" - } - send(ChatChunk{ - ID: id, Object: "chat.completion.chunk", Created: created, Model: rec.Model, - Choices: []ChunkChoice{{Index: 0, Delta: RespMessage{}, FinishReason: &finalFinish}}, - }) - // Final usage chunk (OpenAI standard: empty choices + usage before [DONE]). - // Prefer the upstream's exact usage if the stream carried it; fall back to - // the gateway's estimate otherwise. - var tut *types.TokenUsage - if lastUsage != nil { - tut = lastUsage - } else if rec.Prompt+rec.Compl > 0 { - u := types.TokenUsage{ - Prompt: int(rec.Prompt), - Completion: int(rec.Compl), - Total: int(rec.Prompt + rec.Compl), - } - tut = &u - } - if tut != nil { - send(ChatChunk{ - ID: id, Object: "chat.completion.chunk", Created: created, Model: rec.Model, - Choices: []ChunkChoice{}, - Usage: tut, - }) - } - fmt.Fprintf(w, "data: [DONE]\n\n") - if flusher != nil { - flusher.Flush() - } + g.pumpStream(w, rec, chunks, usedModel) } func (g *Gateway) handleImage(w http.ResponseWriter, r *http.Request) { diff --git a/internal/gateway/gateway_test.go b/internal/gateway/gateway_test.go index e02f17c..8325ac5 100644 --- a/internal/gateway/gateway_test.go +++ b/internal/gateway/gateway_test.go @@ -702,3 +702,40 @@ type flushRecorder struct { } func (f *flushRecorder) Flush() { f.flushed = true } + +// TestDirectStreamFailoverAuditSource: a direct streaming request whose first +// candidate hard-fails must be served by the fallback, and the audit record's +// Source must name the source that actually served the stream (not just the +// first candidate). +func TestDirectStreamFailoverAuditSource(t *testing.T) { + aUp := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + w.WriteHeader(500) + fmt.Fprint(w, `{"error":"boom"}`) + })) + defer aUp.Close() + bUp := mockUpstream() + defer bUp.Close() + g := newTestGateway(t, + // Both sources expose the same model id so the direct scheduler has a + // real fallback candidate after a hard-fails; b is the healthy one. + config.Source{Name: "a", BaseURL: aUp.URL, Adapter: "openai", Models: []config.Model{{ID: "m"}}}, + config.Source{Name: "b", BaseURL: bUp.URL, Adapter: "openai", Models: []config.Model{{ID: "m"}}}, + ) + body := doReq(t, g, "POST", "/v1/chat/completions", + `{"model":"m","stream":true,"messages":[{"role":"user","content":"hi"}]}`) + if body.Code != 200 { + t.Fatalf("status=%d body=%s", body.Code, body.Body.String()) + } + sb := body.Body.String() + if !strings.Contains(sb, `"finish_reason":"stop"`) || !strings.Contains(sb, "[DONE]") { + t.Fatalf("stream incomplete: %s", sb) + } + recs := g.stats.AuditRecords(0, 0, "") + if len(recs) == 0 { + t.Fatal("no audit records") + } + last := recs[len(recs)-1] + if last.Source != "b" || !last.OK { + t.Fatalf("audit rec source=%q ok=%v, want source=b ok=true (actual serving source)", last.Source, last.OK) + } +}