From d465730fe408ca8adbf4fc47bc73f0b1357705ff Mon Sep 17 00:00:00 2001 From: Danielle Mustillo Date: Mon, 19 Oct 2026 03:04:59 -0400 Subject: [PATCH] test(e2e): pass mock LLM handlers to startModelRecorder The egress and mounted-skills mocks put a second server with an unbuffered httputil.ReverseProxy behind startModelRecorder. Once the streamed response starts, that server closes the request body. The proxy transport then reads the body once more, gets ErrBodyReadAfterClose, and closes the upstream connection mid-stream. The agent sees unexpected EOF, and the egress test times out. startModelRecorder now takes an http.Handler and calls it in process with the buffered body. The two mocks lose their second server, and the other callers wrap the mock URL with mockLLMProxy. Signed-off-by: Danielle Mustillo --- go/core/test/e2e/compaction_test.go | 2 +- go/core/test/e2e/egress_test.go | 7 +++-- go/core/test/e2e/genai_baggage_test.go | 2 +- go/core/test/e2e/interaction_test.go | 33 ++++++++++++++--------- go/core/test/e2e/model_recorder_test.go | 36 +++++++++++++++++++++++++ go/core/test/e2e/skills_test.go | 8 +++--- go/core/test/e2e/streaming_test.go | 4 +-- 7 files changed, 67 insertions(+), 25 deletions(-) create mode 100644 go/core/test/e2e/model_recorder_test.go diff --git a/go/core/test/e2e/compaction_test.go b/go/core/test/e2e/compaction_test.go index 5a60f53144..93ed6a1087 100644 --- a/go/core/test/e2e/compaction_test.go +++ b/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"}) diff --git a/go/core/test/e2e/egress_test.go b/go/core/test/e2e/egress_test.go index 075ecf00f5..4ef932ef46 100644 --- a/go/core/test/e2e/egress_test.go +++ b/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") diff --git a/go/core/test/e2e/genai_baggage_test.go b/go/core/test/e2e/genai_baggage_test.go index 359ab39eb5..38a2858b5b 100644 --- a/go/core/test/e2e/genai_baggage_test.go +++ b/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) diff --git a/go/core/test/e2e/interaction_test.go b/go/core/test/e2e/interaction_test.go index e0169dd08d..3017f6d591 100644 --- a/go/core/test/e2e/interaction_test.go +++ b/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 { diff --git a/go/core/test/e2e/model_recorder_test.go b/go/core/test/e2e/model_recorder_test.go new file mode 100644 index 0000000000..6659a1214d --- /dev/null +++ b/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)) +} diff --git a/go/core/test/e2e/skills_test.go b/go/core/test/e2e/skills_test.go index d88920d4bc..741b86053c 100644 --- a/go/core/test/e2e/skills_test.go +++ b/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]) } } diff --git a/go/core/test/e2e/streaming_test.go b/go/core/test/e2e/streaming_test.go index 267697c3e8..27dd8a90eb 100644 --- a/go/core/test/e2e/streaming_test.go +++ b/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: