v2.3: fix the queue=false release race in the given test; add TestLoadIsHeldForTheWholeStream; run notes

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

docs/plans/v2.3/_files/internal/proxy/affinity_test.go

@@ -6,7 +6,9 @@ package proxy_test
 // the server's queue is the one that must show it.
 
 import (
+	"bufio"
 	"net/http"
+	"strings"
 	"sync"
 	"testing"
 	"time"
@@ -118,9 +120,8 @@ func TestQueueFalseNeitherHoldsNorRefuses(t *testing.T) {
 	if d := time.Since(start); d > 1100*time.Millisecond {
 		t.Errorf("three queue = false requests took %v; they were held", d)
 	}
-	if n := r.lim.InFlight(host, "shared"); n != 0 {
-		t.Errorf("in flight after = %d, want 0", n)
-	}
+	// The slot is given back just after the answer is sent (a deferred release), so wait for it.
+	waitUntil(t, func() bool { return r.lim.InFlight(host, "shared") == 0 })
 	// Accounting is unchanged: each chat is still a row.
 	waitUntil(t, func() bool { return r.rows("bm") == 3 })
 }
@@ -151,3 +152,29 @@ func TestConversationAffinityStillSpreads(t *testing.T) {
 		t.Errorf("two concurrent conversations on route r used %v, want both hosts", hosts)
 	}
 }
+
+// A request counts against its host for as long as its answer is streaming, not only until the
+// first byte: a slot (queueing route) or a tracked place (queue = false) is given back when the
+// stream ends.
+func TestLoadIsHeldForTheWholeStream(t *testing.T) {
+	for _, route := range []string{"r", "bm"} {
+		t.Run(route, func(t *testing.T) {
+			alpha, beta := newUpstream(t, "alpha"), newUpstream(t, "beta")
+			r := newRig(t, affinityHosts, alpha, beta)
+			body := strings.Replace(conversation(1, 1), `"stream":false`, `"stream":true`, 1)
+			resp := r.do(http.MethodPost, "/"+route+"/v1/chat/completions", body)
+			defer resp.Body.Close()
+			host := resp.Header.Get("X-Crossbar-Host")
+			line, err := bufio.NewReader(resp.Body).ReadString('\n')
+			if err != nil || !strings.HasPrefix(line, "data:") {
+				t.Fatalf("first line %q, err %v", line, err)
+			}
+			// The first chunk is here; the upstream sends more for another ~30 ms.
+			if n := r.lim.InFlight(host, "shared"); n != 1 {
+				t.Errorf("in flight on %s after the first chunk = %d, want 1 (released before the stream ended)", host, n)
+			}
+			drain(resp)
+			waitUntil(t, func() bool { return r.lim.InFlight(host, "shared") == 0 })
+		})
+	}
+}