AutoZero/llm/streamguard_test.go

package llm

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

// stallServer 吐一个增量之后彻底静默——正是"网关半开连接/服务端卡在生成阶段"的形态：
// TCP 不断、HTTP 头已回、之后再无字节。没有看门狗时 bufio.Scanner 会永久阻塞在 Read。
func stallServer(t *testing.T, hold chan struct{}) *httptest.Server {
	t.Helper()
	return httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		if r.Method == http.MethodPost {
			w.Header().Set("Content-Type", "text/event-stream")
			w.Write([]byte("data: {\"choices\":[{\"delta\":{\"content\":\"partial\"}}]}\n\n"))
			if f, ok := w.(http.Flusher); ok {
				f.Flush()
			}
		}
		<-hold // 永不返回
	}))
}

func TestStreamIdleWatchdogUnblocks(t *testing.T) {
	old := streamIdleTimeout
	streamIdleTimeout = 150 * time.Millisecond
	defer func() { streamIdleTimeout = old }()

	hold := make(chan struct{})
	srv := stallServer(t, hold)
	defer func() {
		close(hold)
		srv.Close()
	}()

	p := &OpenAIProvider{BaseURL: srv.URL, APIKey: "k", Model: "m"}
	done := make(chan error, 1)
	go func() {
		_, err := p.Stream(context.Background(), "sys",
			[]Message{{Role: RoleUser, Content: "hi"}}, nil, StreamOptions{}, nil)
		done <- err
	}()

	select {
	case err := <-done:
		var se *StalledError
		if !errors.As(err, &se) {
			t.Fatalf("要的是 StalledError，got %v", err)
		}
		if !strings.Contains(err.Error(), "stalled") {
			t.Errorf("错误必须说清是静默掐断，got=%q", err)
		}
	case <-time.After(10 * time.Second):
		t.Fatal("看门狗没生效：Stream 仍被静默的连接卡着")
	}
}

// TestStreamIdleWatchdogSparesHealthyStream：不能错杀慢但活的流——
// 边睡边吐字节的流必须完整读完。
func TestStreamIdleWatchdogSparesHealthyStream(t *testing.T) {
	old := streamIdleTimeout
	streamIdleTimeout = 400 * time.Millisecond
	defer func() { streamIdleTimeout = old }()

	srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		w.Header().Set("Content-Type", "text/event-stream")
		f, _ := w.(http.Flusher)
		// 三段，每段间隔 250ms：单段静默都低于 400ms 阈值。
		for _, s := range []string{"he", "ll", "o"} {
			w.Write([]byte("data: {\"choices\":[{\"delta\":{\"content\":\"" + s + "\"}}]}\n\n"))
			if f != nil {
				f.Flush()
			}
			time.Sleep(250 * time.Millisecond)
		}
		w.Write([]byte("data: [DONE]\n\n"))
		if f != nil {
			f.Flush()
		}
	}))
	defer srv.Close()

	p := &OpenAIProvider{BaseURL: srv.URL, APIKey: "k", Model: "m"}
	msg, err := p.Stream(context.Background(), "sys",
		[]Message{{Role: RoleUser, Content: "hi"}}, nil, StreamOptions{}, nil)
	if err != nil {
		t.Fatalf("健康流被误杀：%v", err)
	}
	if msg.Content != "hello" {
		t.Errorf("Content = %q, want hello", msg.Content)
	}
}

// trickle 是一个可控的"滴流"响应体：按给定节奏每次吐一个字节，
// 滴完之后彻底静默（阻塞在读上），用来精确复现"服务端连上不吐字节"。
// 和真响应体一样：Close 会让阻塞中的 Read 立刻返回错误。
type trickle struct {
	gaps   []time.Duration
	i      int
	closed chan struct{}
}

func (r *trickle) Read(p []byte) (int, error) {
	if r.i >= len(r.gaps) {
		<-r.closed
		return 0, io.ErrClosedPipe
	}
	time.Sleep(r.gaps[r.i])
	r.i++
	p[0] = 'x'
	return 1, nil
}

func (r *trickle) Close() error { close(r.closed); return nil }

func (r *trickle) closedNow() bool {
	select {
	case <-r.closed:
		return true
	default:
		return false
	}
}

// TestIdleGuardMarksStall：看护的两条基本契约——有字节就重置计时，
// 真静默了才掐断，并且能把原因报出来。
func TestIdleGuardMarksStall(t *testing.T) {
	const idle = 300 * time.Millisecond
	src := &trickle{gaps: []time.Duration{60 * time.Millisecond, 60, 60, 60}, closed: make(chan struct{})}
	g := guardStream(src, idle)
	defer g.Stop()

	buf := make([]byte, 1)
	for i := 0; i < len(src.gaps); i++ {
		if _, err := g.Read(buf); err != nil {
			t.Fatalf("Read 透传失败：%v", err)
		}
		if err := g.Err(); err != nil {
			t.Fatalf("活跃流被误判为静默：%v", err)
		}
	}

	// 此后彻底没有数据：Read 阻塞在源上，只有看门狗关掉它能解阻塞。
	readErr := make(chan error, 1)
	go func() { _, err := g.Read(buf); readErr <- err }()

	deadline := time.Now().Add(10 * time.Second)
	for g.Err() == nil {
		if time.Now().After(deadline) {
			t.Fatal("静默没有被判定")
		}
		time.Sleep(20 * time.Millisecond)
	}
	if !strings.Contains(g.Err().Error(), idle.String()) {
		t.Errorf("错误里应带阈值，got=%q", g.Err())
	}
	select {
	case err := <-readErr:
		if err == nil {
			t.Error("body 被掐断后 Read 必须返回错误（否则调用方仍会阻塞）")
		}
	case <-time.After(2 * time.Second):
		t.Fatal("掐断后 Read 仍卡着——调用方解不了阻塞，看门狗等于没做")
	}
	if !src.closedNow() {
		t.Error("看门狗没有关闭响应体")
	}
}