Agent-System/internal/agentx/llm/sse_test.go

package llm

import (
	"context"
	"io"
	"net/http"
	"net/http/httptest"
	"strings"
	"testing"
	"time"
)

func TestParseSSETextAndUsage(t *testing.T) {
	body := strings.Join([]string{
		`data: {"choices":[{"delta":{"content":"你"},"index":0}]}`,
		`data: {"choices":[{"delta":{"content":"好"},"finish_reason":"stop"}],"usage":{"prompt_tokens":3,"completion_tokens":2,"total_tokens":5}}`,
		`data: [DONE]`,
		"",
	}, "\n")
	var deltas []string
	var end Event
	for ev := range ParseSSE(context.Background(), strings.NewReader(body)) {
		switch ev.Kind {
		case KindTextDelta:
			deltas = append(deltas, ev.Text)
		case KindMessageEnd:
			end = ev
		default:
			t.Fatalf("unexpected event %#v", ev)
		}
	}
	if strings.Join(deltas, "") != "你好" {
		t.Fatalf("deltas %v", deltas)
	}
	if end.Message.StopReason != StopEnd || end.Message.Content != "你好" {
		t.Fatalf("end %#v", end.Message)
	}
	if end.Usage.TotalTokens != 5 {
		t.Fatalf("usage %#v", end.Usage)
	}
}

func TestParseSSEToolCallIncrements(t *testing.T) {
	body := strings.Join([]string{
		`data: {"choices":[{"delta":{"tool_calls":[{"index":0,"id":"call_1","function":{"name":"read_file","arguments":""}}]}}]}`,
		`data: {"choices":[{"delta":{"tool_calls":[{"index":0,"function":{"arguments":"{\"path\":"}}]}}]}`,
		`data: {"choices":[{"delta":{"tool_calls":[{"index":0,"function":{"arguments":"\"a.md\"}"}}]},"finish_reason":"tool_calls"}]}`,
		`data: [DONE]`,
		"",
	}, "\n")
	var end Event
	deltas := 0
	for ev := range ParseSSE(context.Background(), strings.NewReader(body)) {
		if ev.Kind == KindToolCallDelta {
			deltas++
		}
		if ev.Kind == KindMessageEnd {
			end = ev
		}
	}
	if deltas != 3 {
		t.Fatalf("tool deltas %d", deltas)
	}
	if end.Message.StopReason != StopToolCalls || len(end.Message.ToolCalls) != 1 {
		t.Fatalf("end %#v", end.Message)
	}
	tc := end.Message.ToolCalls[0]
	if tc.ID != "call_1" || tc.Name != "read_file" || tc.Arguments != `{"path":"a.md"}` {
		t.Fatalf("tool call %#v", tc)
	}
}

func TestParseSSEMalformedIsError(t *testing.T) {
	body := "data: {not-json}\n\n"
	var got Event
	for ev := range ParseSSE(context.Background(), strings.NewReader(body)) {
		got = ev
	}
	if got.Kind != KindError || got.Message.StopReason != StopError || got.Retryable {
		t.Fatalf("event %#v", got)
	}
}

func TestParseSSECancel(t *testing.T) {
	pr, pw := io.Pipe()
	ctx, cancel := context.WithCancel(context.Background())
	ch := ParseSSE(ctx, pr)
	go func() {
		_, _ = io.WriteString(pw, "data: {\"choices\":[{\"delta\":{\"content\":\"x\"}}]}\n\n")
		cancel()
		time.Sleep(20 * time.Millisecond)
		_ = pw.Close()
	}()
	sawAbort := false
	for ev := range ch {
		if ev.Kind == KindError && ev.Message.StopReason == StopAborted {
			sawAbort = true
		}
	}
	if !sawAbort {
		t.Fatal("expected aborted event")
	}
}

func TestClientHTTPErrorIsRetryable(t *testing.T) {
	srv := httpServer(t, func(w http.ResponseWriter, r *http.Request) {
		if r.Header.Get("Authorization") != "Bearer secret" {
			t.Errorf("auth header %q", r.Header.Get("Authorization"))
		}
		http.Error(w, "busy", http.StatusTooManyRequests)
	})
	c := &Client{BaseURL: srv, APIKey: "secret", HTTPClient: srvClient()}
	var got Event
	for ev := range c.Stream(context.Background(), "MiniMax-M2.7-highspeed", []Message{{Role: RoleUser, Content: "hi"}}, nil, StreamOptions{}) {
		got = ev
	}
	if got.Kind != KindError || !got.Retryable || got.Message.StopReason != StopError {
		t.Fatalf("event %#v", got)
	}
}

func TestClientStreamsCompletion(t *testing.T) {
	srv := httpServer(t, func(w http.ResponseWriter, r *http.Request) {
		w.Header().Set("Content-Type", "text/event-stream")
		_, _ = io.WriteString(w, "data: {\"choices\":[{\"delta\":{\"content\":\"ok\"},\"finish_reason\":\"stop\"}]}\n\ndata: [DONE]\n")
	})
	c := &Client{BaseURL: srv, APIKey: "k", HTTPClient: srvClient()}
	var end Event
	for ev := range c.Stream(context.Background(), "m", []Message{{Role: RoleUser, Content: "hi"}}, []ToolDecl{{Name: "read_file", Parameters: []byte(`{"type":"object"}`)}}, StreamOptions{}) {
		if ev.Kind == KindMessageEnd {
			end = ev
		}
	}
	if end.Message.Content != "ok" || end.Message.StopReason != StopEnd {
		t.Fatalf("end %#v", end.Message)
	}
}

func TestClientCancel(t *testing.T) {
	started := make(chan struct{})
	srv := httpServer(t, func(w http.ResponseWriter, r *http.Request) {
		close(started)
		w.Header().Set("Content-Type", "text/event-stream")
		fl, _ := w.(http.Flusher)
		_, _ = io.WriteString(w, "data: {\"choices\":[{\"delta\":{\"content\":\"a\"}}]}\n\n")
		if fl != nil {
			fl.Flush()
		}
		<-r.Context().Done()
	})
	c := &Client{BaseURL: srv, APIKey: "k", HTTPClient: &http.Client{}}
	ctx, cancel := context.WithCancel(context.Background())
	ch := c.Stream(ctx, "m", []Message{{Role: RoleUser, Content: "hi"}}, nil, StreamOptions{})
	<-started
	cancel()
	saw := false
	deadline := time.After(2 * time.Second)
	for {
		select {
		case ev, ok := <-ch:
			if !ok {
				if !saw {
					t.Fatal("channel closed without abort")
				}
				return
			}
			if ev.Kind == KindError && ev.Message.StopReason == StopAborted {
				saw = true
			}
		case <-deadline:
			t.Fatal("timed out waiting for cancel")
		}
	}
}

func httpServer(t *testing.T, h http.HandlerFunc) string {
	t.Helper()
	srv := httptest.NewServer(h)
	t.Cleanup(srv.Close)
	return srv.URL
}

func srvClient() *http.Client {
	return &http.Client{Timeout: 2 * time.Second}
}

func TestStreamErrorString(t *testing.T) {
	if (&StreamError{Message: "boom"}).Error() != "boom" {
		t.Fatal("error string")
	}
}