test(e2e): pass mock LLM handlers to startModelRecorder

go/core/test/e2e/compaction_test.go

@@ -47,7 +47,7 @@ const (
 func TestSessionContextCompaction(t *testing.T) {
 	t.Parallel()
 	target := interactionTarget(t)
-	recorder := startModelRecorder(t, startMockLLMServer(t, compactionMocks, "mocks/invoke_golang_compaction.json"), nil)
+	recorder := startModelRecorder(t, mockLLMProxy(t, startMockLLMServer(t, compactionMocks, "mocks/invoke_golang_compaction.json")), nil)
 	kube := interactionKubeClient(t)
 	modelURL := reachableModelURL(t, recorder.URL)
 	agentModel := createInteractionModel(t, kube, modelURL, map[string]string{"X-Kagent-E2E-Model": "agent"})

go/core/test/e2e/egress_test.go

@@ -361,7 +361,7 @@ func startEgressSkillMock(t *testing.T, root, destination string, runNames ...st
 	t.Helper()
 	require.NotEmpty(t, runNames)
 	var active atomic.Pointer[egressProbeRun]
-	server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+	upstream := http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
 		run := active.Load()
 		if run == nil {
 			http.Error(w, "no active probe", http.StatusServiceUnavailable)
@@ -395,9 +395,8 @@ func startEgressSkillMock(t *testing.T, root, destination string, runNames ...st
 			return
 		}
 		run.proxy.ServeHTTP(w, r)
-	}))
-	t.Cleanup(server.Close)
-	recorder := startModelRecorder(t, server.URL, nil)
+	})
+	recorder := startModelRecorder(t, upstream, nil)
 	modelURL := reachableModelURL(t, recorder.URL)
 	origin := strings.TrimSuffix(modelURL, "/v1")
 	raw, err := interactionMocks.ReadFile("mocks/invoke_egress_skill.json")

go/core/test/e2e/genai_baggage_test.go

@@ -49,7 +49,7 @@ func TestRuntimeGenAIBaggage(t *testing.T) {
 			require.NoError(t, err)
 			prompt := "Add 3 and 5 using the configured MCP server."
 			require.NoError(t, json.Unmarshal([]byte(`{"role":"user","content":"`+prompt+`"}`), &cfg.OpenAI[0].Match.Message))
-			recorder := startModelRecorder(t, startMockLLMConfig(t, cfg), nil)
+			recorder := startModelRecorder(t, mockLLMProxy(t, startMockLLMConfig(t, cfg)), nil)
 			model := createInteractionModel(t, kube, reachableModelURL(t, recorder.URL), map[string]string{"X-Kagent-E2E-Model": "agent"})
 
 			harness := createPropagatorsHarness(t, kube, tc.propagators)

go/core/test/e2e/interaction_test.go

@@ -1369,7 +1369,7 @@ func startForkMemoryMock(t *testing.T) string {
 	config.OpenAI = append(config.OpenAI, continuation.OpenAI...)
 	config.OpenAIResponse = append(config.OpenAIResponse, continuation.OpenAIResponse...)
 	config.Anthropic = append(config.Anthropic, continuation.Anthropic...)
-	recorder := startModelRecorder(t, startMockLLMConfig(t, config), func(body []byte) error {
+	recorder := startModelRecorder(t, mockLLMProxy(t, startMockLLMConfig(t, config)), func(body []byte) error {
 		if !bytes.Contains(body, []byte("What was the answer before the checkpoint?")) {
 			return nil
 		}
@@ -1397,18 +1397,14 @@ type recordedModelRequest struct {
 	Body   []byte
 }
 
-// startModelRecorder puts a recording proxy in front of the mock LLM at
-// upstreamURL. An inspect function may reject a request: its error is
-// answered with 400 instead of being forwarded, which fails the agent's turn
-// visibly rather than letting the mock answer a prompt it should not see.
-func startModelRecorder(t *testing.T, upstreamURL string, inspect func(body []byte) error) *modelRecorder {
+// startModelRecorder puts a recording server in front of upstream, which is
+// usually a proxy to the mock LLM. An inspect function may reject a request:
+// its error is answered with 400 instead of being forwarded, which fails the
+// agent's turn visibly rather than letting the mock answer a prompt it should
+// not see.
+func startModelRecorder(t *testing.T, upstream http.Handler, inspect func(body []byte) error) *modelRecorder {
 	t.Helper()
-	upstream, err := url.Parse(upstreamURL)
-	if err != nil {
-		t.Fatal(err)
-	}
 	recorder := &modelRecorder{}
-	proxy := httputil.NewSingleHostReverseProxy(upstream)
 	server := httptest.NewUnstartedServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
 		body, err := io.ReadAll(r.Body)
 		if err != nil {
@@ -1425,10 +1421,13 @@ func startModelRecorder(t *testing.T, upstreamURL string, inspect func(body []by
 				return
 			}
 		}
+		// Upstream proxies read the body after the response starts, so they
+		// need a reader that the server cannot close under them.
 		r.Body = io.NopCloser(bytes.NewReader(body))
-		proxy.ServeHTTP(w, r)
+		upstream.ServeHTTP(w, r)
 	}))
 	_ = server.Listener.Close()
+	var err error
 	server.Listener, err = net.Listen("tcp", "0.0.0.0:0")
 	if err != nil {
 		t.Fatal(err)
@@ -1439,6 +1438,16 @@ func startModelRecorder(t *testing.T, upstreamURL string, inspect func(body []by
 	return recorder
 }
 
+// mockLLMProxy returns a proxy to the mock LLM at rawURL for startModelRecorder.
+func mockLLMProxy(t *testing.T, rawURL string) *httputil.ReverseProxy {
+	t.Helper()
+	upstream, err := url.Parse(rawURL)
+	if err != nil {
+		t.Fatal(err)
+	}
+	return httputil.NewSingleHostReverseProxy(upstream)
+}
+
 // Requests returns the recorded requests carrying the header value, in
 // arrival order.
 func (r *modelRecorder) Requests(header, value string) []recordedModelRequest {

go/core/test/e2e/model_recorder_test.go

@@ -0,0 +1,36 @@
+package e2e_test
+
+import (
+	"io"
+	"net/http"
+	"strings"
+	"testing"
+
+	"github.com/stretchr/testify/require"
+)
+
+func TestModelRecorderBodyOutlivesResponseStart(t *testing.T) {
+	t.Parallel()
+	reread := make(chan string, 1)
+	recorder := startModelRecorder(t, http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		w.Header().Set("Content-Type", "text/event-stream")
+		w.WriteHeader(http.StatusOK)
+		http.NewResponseController(w).Flush()
+		body, err := io.ReadAll(r.Body)
+		if err != nil {
+			reread <- err.Error()
+			return
+		}
+		reread <- string(body)
+		_, _ = io.WriteString(w, "data: done\n\n")
+	}), nil)
+
+	response, err := http.Post(recorder.URL+"/v1/chat/completions", "application/json", strings.NewReader(`{"stream":true}`))
+	require.NoError(t, err)
+	defer response.Body.Close()
+	streamed, err := io.ReadAll(response.Body)
+	require.NoError(t, err)
+
+	require.Equal(t, `{"stream":true}`, <-reread)
+	require.Equal(t, "data: done\n\n", string(streamed))
+}

go/core/test/e2e/skills_test.go

@@ -7,7 +7,6 @@ import (
 	"fmt"
 	"io"
 	"net/http"
-	"net/http/httptest"
 	"net/http/httputil"
 	"net/url"
 	"slices"
@@ -164,9 +163,8 @@ func startMountedSkillsMock(t *testing.T, root string) (*modelRecorder, func(str
 	}
 	var active atomic.Pointer[httputil.ReverseProxy]
 	active.Store(proxies["initial"])
-	server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+	upstream := http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
 		active.Load().ServeHTTP(w, r)
-	}))
-	t.Cleanup(server.Close)
-	return startModelRecorder(t, server.URL, nil), func(phase string) { active.Store(proxies[phase]) }
+	})
+	return startModelRecorder(t, upstream, nil), func(phase string) { active.Store(proxies[phase]) }
 }

go/core/test/e2e/streaming_test.go

@@ -34,7 +34,7 @@ func TestSessionModelStreamingDisabled(t *testing.T) {
 	} {
 		t.Run(string(format), func(t *testing.T) {
 			t.Parallel()
-			recorder := startModelRecorder(t, startMockLLMServer(t, interactionMocks, "mocks/invoke_agent.json"), func(body []byte) error {
+			recorder := startModelRecorder(t, mockLLMProxy(t, startMockLLMServer(t, interactionMocks, "mocks/invoke_agent.json")), func(body []byte) error {
 				var request struct {
 					Stream bool `json:"stream"`
 				}
@@ -85,7 +85,7 @@ func TestSessionCompletesAfterDisconnect(t *testing.T) {
 		started, release := make(chan struct{}), make(chan struct{})
 		var startedOnce, releaseOnce sync.Once
 		upstream := startMockLLMServer(t, interactionMocks, "mocks/invoke_agent.json")
-		recorder := startModelRecorder(t, upstream, func([]byte) error {
+		recorder := startModelRecorder(t, mockLLMProxy(t, upstream), func([]byte) error {
 			startedOnce.Do(func() { close(started) })
 			select {
 			case <-release: