Agent-Native/runtime/llm/client_test.go

package llm

import (
	"bytes"
	"compress/gzip"
	"context"
	"encoding/json"
	"errors"
	"io"
	"net/http"
	"net/http/httptest"
	"strconv"
	"strings"
	"sync/atomic"
	"testing"
)

type terminalRoundTripFunc func(*http.Request) (*http.Response, error)

func (f terminalRoundTripFunc) RoundTrip(r *http.Request) (*http.Response, error) { return f(r) }

func TestClientTerminalObservationHTTPStatus(t *testing.T) {
	server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		if r.URL.Path != "/chat/completions" {
			t.Errorf("path = %q", r.URL.Path)
		}
		w.Header().Set("Content-Type", "text/event-stream")
		_, _ = w.Write([]byte("data: {\"choices\":[{\"delta\":{\"content\":\"ok\"},\"finish_reason\":\"stop\"}],\"usage\":{\"prompt_tokens\":2,\"completion_tokens\":1,\"total_tokens\":3}}\n\ndata: [DONE]\n\n"))
	}))
	defer server.Close()
	client := &Client{BaseURL: server.URL, APIKey: "fixture", HTTPClient: server.Client()}
	var end Event
	for ev := range client.Stream(context.Background(), "fixture", []Message{{Role: RoleUser, Content: "hello"}}, nil, StreamOptions{}) {
		if ev.Kind == KindMessageEnd {
			end = ev
		}
	}
	if end.HTTPStatus != http.StatusOK || end.TerminalObservation == nil ||
		end.TerminalObservation.ReadOutcome != "done_marker" || !end.StreamComplete || end.Usage.TotalTokens != 3 {
		t.Fatalf("successful HTTP terminal lost transport evidence: %#v", end)
	}
}

func TestClientTerminalObservationProviderAndTransportFailure(t *testing.T) {
	server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		http.Error(w, "PRIVATE_PROVIDER_BODY_R05", http.StatusTooManyRequests)
	}))
	defer server.Close()
	tests := []struct {
		name    string
		client  *Client
		status  int
		outcome string
	}{
		{"provider-http", &Client{BaseURL: server.URL, APIKey: "fixture", HTTPClient: server.Client()}, http.StatusTooManyRequests, "provider_error"},
		{"transport", &Client{BaseURL: "http://test.invalid", APIKey: "fixture", HTTPClient: &http.Client{Transport: terminalRoundTripFunc(func(*http.Request) (*http.Response, error) {
			return nil, errors.New("PRIVATE_TRANSPORT_ERROR_R05")
		})}}, 0, "transport_error"},
	}
	for _, tc := range tests {
		t.Run(tc.name, func(t *testing.T) {
			var end Event
			for ev := range tc.client.Stream(context.Background(), "fixture", []Message{{Role: RoleUser, Content: "hello"}}, nil, StreamOptions{}) {
				end = ev
			}
			if end.Kind != KindError || end.HTTPStatus != tc.status || end.TerminalObservation == nil ||
				end.TerminalObservation.ReadOutcome != tc.outcome {
				t.Fatalf("provider/transport failure was not classified: %#v", end)
			}
		})
	}
}

// bodyBoundaryToolFrames is a complete tool call plus Usage without [DONE].
const bodyBoundaryToolFrames = "data: {\"choices\":[{\"delta\":{\"tool_calls\":[{\"index\":0,\"id\":\"call_PRIVATE_R06\",\"function\":{\"name\":\"read_PRIVATE_R06\",\"arguments\":\"{\\\"path\\\":\\\"PRIVATE_R06.txt\\\"}\"}}]},\"finish_reason\":\"tool_calls\"}]}\n\n" +
	"data: {\"choices\":[],\"usage\":{\"prompt_tokens\":5,\"completion_tokens\":2,\"total_tokens\":7}}\n\n"

const bodyBoundaryDone = "data: [DONE]\n\n"

func streamBodyBoundaryEnd(t *testing.T, ctx context.Context, client *Client) Event {
	t.Helper()
	var end Event
	terminals := 0
	for ev := range client.Stream(ctx, "fixture", []Message{{Role: RoleUser, Content: "hello"}}, nil, StreamOptions{}) {
		if ev.Kind == KindMessageEnd || ev.Kind == KindError {
			end = ev
			terminals++
		} else if ev.TerminalObservation != nil {
			t.Fatalf("non-terminal event carried an observation: %#v", ev)
		}
	}
	if terminals != 1 {
		t.Fatalf("terminal events = %d, want 1", terminals)
	}
	if end.TerminalObservation != nil {
		raw, err := json.Marshal(end.TerminalObservation)
		if err != nil || strings.Contains(string(raw), "PRIVATE") || strings.Contains(string(raw), "127.0.0.1") {
			t.Fatalf("unsafe terminal observation: %s err=%v", raw, err)
		}
	}
	return end
}

type bodyBoundaryWant struct {
	protocol, framing, relation string
	declared, eof, uncompressed bool
	bytes                       int64
}

func requireBodyBoundary(t *testing.T, end Event, want bodyBoundaryWant) {
	t.Helper()
	if end.TerminalObservation == nil || end.TerminalObservation.BodyBoundary == nil {
		t.Fatalf("missing body boundary observation: %#v", end)
	}
	body := end.TerminalObservation.BodyBoundary
	if body.ProtocolClass != want.protocol || body.FramingClass != want.framing ||
		body.LengthRelation != want.relation || body.DeclaredLengthKnown != want.declared ||
		body.EOFSeen != want.eof || body.Uncompressed != want.uncompressed || body.BytesSaturated {
		t.Fatalf("wrong body boundary: %+v want %+v", body, want)
	}
	if want.bytes >= 0 && body.BytesRead != want.bytes {
		t.Fatalf("bytes read = %d, want %d", body.BytesRead, want.bytes)
	}
}

func requireUnreleasedTools(t *testing.T, end Event) {
	t.Helper()
	if end.Kind != KindMessageEnd || end.Err == nil || end.StreamComplete || len(end.Message.ToolCalls) != 0 ||
		end.Usage.TotalTokens != 7 || end.HTTPStatus != http.StatusOK {
		t.Fatalf("unfinished tool stream was released or lost usage: %#v", end)
	}
}

func fixtureResponseClient(resp func() *http.Response) *Client {
	return &Client{BaseURL: "http://fixture.invalid", APIKey: "fixture", HTTPClient: &http.Client{
		Transport: terminalRoundTripFunc(func(*http.Request) (*http.Response, error) { return resp(), nil }),
	}}
}

func TestClientBodyBoundaryShortRead(t *testing.T) {
	t.Run("declared-longer-than-body", func(t *testing.T) {
		server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
			w.Header().Set("Content-Type", "text/event-stream")
			w.Header().Set("Content-Length", strconv.Itoa(len(bodyBoundaryToolFrames)+64))
			_, _ = io.WriteString(w, bodyBoundaryToolFrames)
		}))
		defer server.Close()
		end := streamBodyBoundaryEnd(t, context.Background(), &Client{BaseURL: server.URL, APIKey: "fixture", HTTPClient: server.Client()})
		requireUnreleasedTools(t, end)
		if end.TerminalObservation.ReadOutcome != "read_error" {
			t.Fatalf("short body was not a read error: %+v", end.TerminalObservation)
		}
		requireBodyBoundary(t, end, bodyBoundaryWant{protocol: "http1", framing: "known_length", relation: "unknown",
			declared: true, bytes: int64(len(bodyBoundaryToolFrames))})
	})
	t.Run("declared-length-mismatch", func(t *testing.T) {
		client := fixtureResponseClient(func() *http.Response {
			return &http.Response{StatusCode: http.StatusOK, ProtoMajor: 1, ProtoMinor: 1,
				ContentLength: int64(len(bodyBoundaryToolFrames) + 64), Header: http.Header{},
				Body: io.NopCloser(strings.NewReader(bodyBoundaryToolFrames))}
		})
		end := streamBodyBoundaryEnd(t, context.Background(), client)
		requireUnreleasedTools(t, end)
		requireBodyBoundary(t, end, bodyBoundaryWant{protocol: "http1", framing: "known_length", relation: "mismatch",
			declared: true, eof: true, bytes: int64(len(bodyBoundaryToolFrames))})
	})
}

func TestClientBodyBoundaryUncompressed(t *testing.T) {
	var compressed bytes.Buffer
	zw := gzip.NewWriter(&compressed)
	_, _ = io.WriteString(zw, bodyBoundaryToolFrames)
	if err := zw.Close(); err != nil {
		t.Fatal(err)
	}
	server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		w.Header().Set("Content-Type", "text/event-stream")
		w.Header().Set("Content-Encoding", "gzip")
		w.Header().Set("Content-Length", strconv.Itoa(compressed.Len()))
		_, _ = w.Write(compressed.Bytes())
	}))
	defer server.Close()
	end := streamBodyBoundaryEnd(t, context.Background(), &Client{BaseURL: server.URL, APIKey: "fixture", HTTPClient: server.Client()})
	requireUnreleasedTools(t, end)
	requireBodyBoundary(t, end, bodyBoundaryWant{protocol: "http1", framing: "unknown", relation: "unknown",
		eof: true, uncompressed: true, bytes: int64(len(bodyBoundaryToolFrames))})
}

// doneThenErrorBody returns the whole stream without io.EOF; any later read fails.
type doneThenErrorBody struct {
	data  *strings.Reader
	reads int
}

func (b *doneThenErrorBody) Read(p []byte) (int, error) {
	if b.data.Len() > 0 {
		return b.data.Read(p)
	}
	b.reads++
	return 0, errors.New("PRIVATE_READ_AFTER_DONE_R06")
}

func (b *doneThenErrorBody) Close() error { return nil }

func TestClientDoneDoesNotNeedBodyEOF(t *testing.T) {
	stream := bodyBoundaryToolFrames + bodyBoundaryDone
	body := &doneThenErrorBody{data: strings.NewReader(stream)}
	client := fixtureResponseClient(func() *http.Response {
		return &http.Response{StatusCode: http.StatusOK, ProtoMajor: 1, ProtoMinor: 1, ContentLength: -1,
			TransferEncoding: []string{"chunked"}, Header: http.Header{}, Body: body}
	})
	end := streamBodyBoundaryEnd(t, context.Background(), client)
	if end.Kind != KindMessageEnd || end.Err != nil || !end.StreamComplete || len(end.Message.ToolCalls) != 1 ||
		end.Message.StopReason != StopToolCalls || end.Usage.TotalTokens != 7 || end.HTTPStatus != http.StatusOK ||
		end.TerminalObservation.ReadOutcome != "done_marker" {
		t.Fatalf("legal [DONE] tool terminal was not kept: %#v", end)
	}
	requireBodyBoundary(t, end, bodyBoundaryWant{protocol: "http1", framing: "chunked", relation: "unknown",
		bytes: int64(len(stream))})
	if body.reads != 0 {
		t.Fatalf("parser read past [DONE] %d times", body.reads)
	}

	var standalone Event
	for ev := range ParseSSE(context.Background(), strings.NewReader(stream)) {
		standalone = ev
	}
	if standalone.TerminalObservation == nil || standalone.TerminalObservation.BodyBoundary != nil ||
		len(standalone.Message.ToolCalls) != 1 || !standalone.StreamComplete {
		t.Fatalf("standalone ParseSSE changed: %#v", standalone)
	}
}

func TestClientBodyBoundaryChunkedAndHTTP2(t *testing.T) {
	flushed := func(stream string) http.HandlerFunc {
		return func(w http.ResponseWriter, r *http.Request) {
			w.Header().Set("Content-Type", "text/event-stream")
			_, _ = io.WriteString(w, stream)
			w.(http.Flusher).Flush()
		}
	}
	tests := []struct {
		name     string
		server   func() *httptest.Server
		wantDone bool
		want     bodyBoundaryWant
	}{
		{"known-length", func() *httptest.Server {
			stream := bodyBoundaryToolFrames + bodyBoundaryDone
			return httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
				w.Header().Set("Content-Type", "text/event-stream")
				w.Header().Set("Content-Length", strconv.Itoa(len(stream)))
				_, _ = io.WriteString(w, stream)
			}))
		}, true, bodyBoundaryWant{protocol: "http1", framing: "known_length", relation: "match", declared: true, eof: true,
			bytes: int64(len(bodyBoundaryToolFrames + bodyBoundaryDone))}},
		{"chunked", func() *httptest.Server {
			return httptest.NewServer(flushed(bodyBoundaryToolFrames))
		}, false, bodyBoundaryWant{protocol: "http1", framing: "chunked", relation: "unknown", eof: true,
			bytes: int64(len(bodyBoundaryToolFrames))}},
		{"http2", func() *httptest.Server {
			server := httptest.NewUnstartedServer(flushed(bodyBoundaryToolFrames))
			server.EnableHTTP2 = true
			server.StartTLS()
			return server
		}, false, bodyBoundaryWant{protocol: "http2", framing: "http2_stream", relation: "unknown", eof: true,
			bytes: int64(len(bodyBoundaryToolFrames))}},
	}
	for _, tc := range tests {
		t.Run(tc.name, func(t *testing.T) {
			server := tc.server()
			defer server.Close()
			end := streamBodyBoundaryEnd(t, context.Background(), &Client{BaseURL: server.URL, APIKey: "fixture", HTTPClient: server.Client()})
			requireBodyBoundary(t, end, tc.want)
			if !tc.wantDone {
				requireUnreleasedTools(t, end)
				return
			}
			if end.Err != nil || !end.StreamComplete || len(end.Message.ToolCalls) != 1 || end.Usage.TotalTokens != 7 {
				t.Fatalf("legal [DONE] tool terminal was not kept: %#v", end)
			}
		})
	}
}

func TestClientBodyBoundaryCloseDelimitedUnknown(t *testing.T) {
	t.Run("close-delimited", func(t *testing.T) {
		server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
			w.Header().Set("Content-Type", "text/event-stream")
			w.Header().Set("Transfer-Encoding", "identity")
			w.WriteHeader(http.StatusOK)
			_, _ = io.WriteString(w, bodyBoundaryToolFrames)
			w.(http.Flusher).Flush()
		}))
		defer server.Close()
		end := streamBodyBoundaryEnd(t, context.Background(), &Client{BaseURL: server.URL, APIKey: "fixture", HTTPClient: server.Client()})
		requireUnreleasedTools(t, end)
		requireBodyBoundary(t, end, bodyBoundaryWant{protocol: "http1", framing: "close_delimited", relation: "unknown",
			eof: true, bytes: int64(len(bodyBoundaryToolFrames))})
	})
	t.Run("undeterminable", func(t *testing.T) {
		client := fixtureResponseClient(func() *http.Response {
			return &http.Response{StatusCode: http.StatusOK, ContentLength: -1, Header: http.Header{},
				Body: io.NopCloser(strings.NewReader(bodyBoundaryToolFrames))}
		})
		end := streamBodyBoundaryEnd(t, context.Background(), client)
		requireUnreleasedTools(t, end)
		requireBodyBoundary(t, end, bodyBoundaryWant{protocol: "other", framing: "unknown", relation: "unknown",
			eof: true, bytes: int64(len(bodyBoundaryToolFrames))})
	})
}

func TestClientBodyBoundaryFailuresDoNotClaimEOF(t *testing.T) {
	t.Run("cancelled", func(t *testing.T) {
		server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
			w.Header().Set("Content-Type", "text/event-stream")
			_, _ = io.WriteString(w, "data: {\"choices\":[{\"delta\":{\"content\":\"PRIVATE_R06\"}}]}\n\n")
			w.(http.Flusher).Flush()
			<-r.Context().Done()
		}))
		defer server.Close()
		ctx, cancel := context.WithCancel(context.Background())
		defer cancel()
		client := &Client{BaseURL: server.URL, APIKey: "fixture", HTTPClient: server.Client()}
		var end Event
		for ev := range client.Stream(ctx, "fixture", []Message{{Role: RoleUser, Content: "hello"}}, nil, StreamOptions{}) {
			if ev.Kind == KindTextDelta {
				cancel()
			}
			if ev.Kind == KindMessageEnd || ev.Kind == KindError {
				end = ev
			}
		}
		if end.Kind != KindError || end.Message.StopReason != StopAborted || end.TerminalObservation == nil ||
			end.TerminalObservation.ReadOutcome != "cancelled" {
			t.Fatalf("cancelled stream was not aborted: %#v", end)
		}
		if body := end.TerminalObservation.BodyBoundary; body == nil || body.EOFSeen || body.LengthRelation != "unknown" {
			t.Fatalf("cancelled stream claimed body EOF: %+v", body)
		}
	})
	t.Run("provider-4xx", func(t *testing.T) {
		server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
			http.Error(w, "PRIVATE_PROVIDER_BODY_R06", http.StatusBadRequest)
		}))
		defer server.Close()
		end := streamBodyBoundaryEnd(t, context.Background(), &Client{BaseURL: server.URL, APIKey: "fixture", HTTPClient: server.Client()})
		if end.Kind != KindError || end.HTTPStatus != http.StatusBadRequest || end.TerminalObservation == nil ||
			end.TerminalObservation.ReadOutcome != "provider_error" || end.TerminalObservation.BodyBoundary != nil {
			t.Fatalf("provider error carried a body boundary: %#v", end)
		}
	})
	t.Run("redirect-then-short-read", func(t *testing.T) {
		var finalHits atomic.Int32
		mux := http.NewServeMux()
		mux.HandleFunc("/chat/completions", func(w http.ResponseWriter, r *http.Request) {
			w.Header().Set("Location", "/final/chat/completions")
			w.WriteHeader(http.StatusTemporaryRedirect)
			_, _ = io.WriteString(w, "PRIVATE_REDIRECT_BODY_R06")
		})
		mux.HandleFunc("/final/chat/completions", func(w http.ResponseWriter, r *http.Request) {
			finalHits.Add(1)
			w.Header().Set("Content-Type", "text/event-stream")
			w.Header().Set("Content-Length", strconv.Itoa(len(bodyBoundaryToolFrames)+64))
			_, _ = io.WriteString(w, bodyBoundaryToolFrames)
		})
		server := httptest.NewServer(mux)
		defer server.Close()
		end := streamBodyBoundaryEnd(t, context.Background(), &Client{BaseURL: server.URL, APIKey: "fixture", HTTPClient: server.Client()})
		if hits := finalHits.Load(); hits != 1 {
			t.Fatalf("redirect target hits = %d, want 1", hits)
		}
		requireUnreleasedTools(t, end)
		requireBodyBoundary(t, end, bodyBoundaryWant{protocol: "http1", framing: "known_length", relation: "unknown",
			declared: true, bytes: int64(len(bodyBoundaryToolFrames))})
	})
}

func TestClientBodyBoundarySaturatesCount(t *testing.T) {
	size := int64(maxObservedBodyBytes + 4096)
	observer := newBodyObserver(&http.Response{ProtoMajor: 1, ContentLength: size,
		Body: io.NopCloser(io.LimitReader(zeroReader{}, size))})
	if n, err := io.Copy(io.Discard, observer); err != nil || n != size {
		t.Fatalf("observer changed the body: n=%d err=%v", n, err)
	}
	body := observer.snapshot()
	if body.BytesRead != maxObservedBodyBytes || !body.BytesSaturated || !body.EOFSeen ||
		body.LengthRelation != "unknown" || !body.DeclaredLengthKnown {
		t.Fatalf("saturated body boundary: %+v", body)
	}
}

type zeroReader struct{}

func (zeroReader) Read(p []byte) (int, error) {
	clear(p)
	return len(p), nil
}