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: