refactor(gateway): collapse chat.go four-way duplication

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

@@ -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) {

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)
+	}
+}