ituoke/internal/agent/stream_test.go

package agent

import (
	"context"
	"encoding/json"
	"errors"
	"fmt"
	"net/http"
	"net/http/httptest"
	"strings"
	"testing"
	"time"
)

func TestStreamDeliversFirstChunkBeforeCompletion(t *testing.T) {
	release := make(chan struct{})
	server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		var body map[string]any
		json.NewDecoder(r.Body).Decode(&body)
		if body["stream"] != true {
			t.Error("stream=true was not sent")
		}
		if r.Header.Get("Authorization") != "Bearer test-key" {
			t.Error("missing stream credential")
		}
		w.Header().Set("Content-Type", "text/event-stream; charset=utf-8")
		fmt.Fprint(w, ": heartbeat\r\nevent: message\r\ndata: {\"choices\":[{\"index\":0,\"delta\":{\"content\":\"你好\"}}]}\r\n\r\n")
		w.(http.Flusher).Flush()
		select {
		case <-release:
		case <-r.Context().Done():
			return
		}
		fmt.Fprint(w, "data: {\"choices\":[\ndata: {\"index\":0,\"delta\":{\"content\":\"，世界🙂\"}}]}\n\n")
		fmt.Fprint(w, "data: {\"choices\":[{\"index\":0,\"delta\":{},\"finish_reason\":\"stop\"}]}\n\ndata: [DONE]")
	}))
	defer server.Close()
	ctx, cancel := context.WithTimeout(context.Background(), 3*time.Second)
	defer cancel()
	chunks := make(chan string, 4)
	finished := make(chan error, 1)
	go func() {
		text, reason, err := completeStream(ctx, testConfig(server.URL), "test-key", nil, func(delta string) { chunks <- delta })
		if err == nil && (text != "你好，世界🙂" || reason != "stop") {
			err = fmt.Errorf("wrong completion: %q %q", text, reason)
		}
		finished <- err
	}()
	select {
	case first := <-chunks:
		if first != "你好" {
			t.Fatalf("first chunk: %q", first)
		}
	case err := <-finished:
		t.Fatalf("returned before releasing stream: %v", err)
	case <-ctx.Done():
		t.Fatal("first chunk was buffered until stream ended")
	}
	close(release)
	if err := <-finished; err != nil {
		t.Fatal(err)
	}
	if second := <-chunks; second != "，世界🙂" {
		t.Fatalf("Unicode/multiline chunk damaged: %q", second)
	}
}

func TestStreamCancellationClosesHTTPAndDoesNotCommitHistory(t *testing.T) {
	disconnected := make(chan struct{})
	server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		w.Header().Set("Content-Type", "text/event-stream")
		fmt.Fprint(w, "data: {\"choices\":[{\"index\":0,\"delta\":{\"content\":\"partial\"}}]}\n\n")
		w.(http.Flusher).Flush()
		<-r.Context().Done()
		close(disconnected)
	}))
	defer server.Close()
	s, _ := testService(t)
	if _, err := s.SaveModelConfig(ModelConfigInput{Config: testConfig(server.URL), APIKey: "key"}); err != nil {
		t.Fatal(err)
	}
	ctx, cancel := context.WithCancel(context.Background())
	defer cancel()
	s.emitDelta = func(delta ChatDelta) {
		if delta.RequestID != "request-1" || delta.Delta != "partial" {
			t.Errorf("wrong event: %+v", delta)
		}
		cancel()
	}
	_, err := s.ChatStream(ctx, "request-1", "phone", "hello")
	if !errors.Is(err, context.Canceled) {
		t.Fatalf("cancellation lost: %v", err)
	}
	if len(s.history["phone"]) != 0 {
		t.Fatal("cancelled answer committed to model context")
	}
	select {
	case <-disconnected:
	case <-time.After(2 * time.Second):
		t.Fatal("cancel did not close HTTP request")
	}
}

func TestStreamCompletionAndFailureMarkers(t *testing.T) {
	cases := []struct {
		name, input, text, reason string
		wantError                 bool
	}{
		{"done", "data: {\"choices\":[{\"delta\":{\"content\":\"ok\"}}]}\n\ndata: [DONE]\n\n", "ok", "", false},
		{"length", "data: {\"choices\":[{\"delta\":{\"content\":\"part\"},\"finish_reason\":\"length\"}]}\n\n", "part", "length", false},
		{"interrupted", "data: {\"choices\":[{\"delta\":{\"content\":\"part\"}}]}\n\n", "part", "", true},
		{"empty", "data: {\"choices\":[{\"delta\":{\"role\":\"assistant\"}}]}\n\ndata: [DONE]\n\n", "", "", true},
		{"malformed", "data: {oops}\n\n", "", "", true},
		{"provider-error", "data: {\"error\":{\"message\":\"rejected secret-key\"}}\n\n", "", "", true},
		{"filtered", "data: {\"choices\":[{\"delta\":{\"content\":\"part\"},\"finish_reason\":\"content_filter\"}]}\n\n", "part", "content_filter", true},
	}
	for _, tc := range cases {
		t.Run(tc.name, func(t *testing.T) {
			var received strings.Builder
			text, reason, err := readCompletionStream(strings.NewReader(tc.input), "secret-key", func(delta string) { received.WriteString(delta) })
			if (err != nil) != tc.wantError || text != tc.text || reason != tc.reason || received.String() != text {
				t.Fatalf("got %q %q %v, callback %q", text, reason, err, received.String())
			}
			if err != nil && strings.Contains(err.Error(), "secret-key") {
				t.Fatal("provider error exposed credential")
			}
		})
	}
}

func TestSuccessfulStreamCommitsHistory(t *testing.T) {
	server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		w.Header().Set("Content-Type", "text/event-stream")
		fmt.Fprint(w, "data: {\"choices\":[{\"delta\":{\"content\":\"hello\"}}]}\n\ndata: {\"choices\":[],\"usage\":{}}\n\ndata: [DONE]\n\n")
	}))
	defer server.Close()
	s, _ := testService(t)
	if _, err := s.SaveModelConfig(ModelConfigInput{Config: testConfig(server.URL), APIKey: "key"}); err != nil {
		t.Fatal(err)
	}
	result, err := s.ChatStream(context.Background(), "request-1", "phone", "question")
	if err != nil || result["content"] != "hello" {
		t.Fatalf("completion: %v %v", result, err)
	}
	if len(s.history["phone"]) != 2 || s.history["phone"][1].Content != "hello" {
		t.Fatal("completed stream did not enter history")
	}
}