AutoZero/agent/stream_tools_test.go

package agent

import (
	"context"
	"errors"
	"sync/atomic"
	"testing"
	"time"

	"git.gopher.ink/mirrors2/AutoZero/llm"
	"git.gopher.ink/mirrors2/AutoZero/tools"
)

// streamToolProvider 在返回最终消息之前先通过 onDelta 通知一个工具调用已完整，
// 然后等只读工具真的启动——用来证明它是"边流边跑"，而不是消息结束才跑。
type streamToolProvider struct {
	read       llm.ToolCall
	side       llm.ToolCall
	readStart  <-chan struct{}
	sideStart  <-chan struct{}
	failStream bool
	calls      int
}

func (p *streamToolProvider) Stream(_ context.Context, _ string, _ []llm.Message, _ []llm.ToolDef,
	_ llm.StreamOptions, onDelta func(llm.StreamDelta)) (llm.Message, error) {
	p.calls++
	if p.calls > 1 {
		return llm.Message{Role: llm.RoleAssistant, Content: "done"}, nil
	}
	onDelta(llm.StreamDelta{ToolCall: &p.read})
	if p.side.Name != "" {
		onDelta(llm.StreamDelta{ToolCall: &p.side})
	}
	select {
	case <-p.readStart:
	case <-time.After(2 * time.Second):
		return llm.Message{}, errors.New("read tool did not start during the stream")
	}
	// 给潜在的（错误的）提前启动留出调度窗口，再确认副作用工具没被提前跑。
	time.Sleep(50 * time.Millisecond)
	if p.sideStart != nil {
		select {
		case <-p.sideStart:
			return llm.Message{}, errors.New("side-effect tool started during the stream")
		default:
		}
	}
	if p.failStream {
		return llm.Message{}, errors.New("stream died")
	}
	calls := []llm.ToolCall{p.read}
	if p.side.Name != "" {
		calls = append(calls, p.side)
	}
	return llm.Message{Role: llm.RoleAssistant, ToolCalls: calls}, nil
}

func streamTools(t *testing.T) (read, side *tools.Tool, readStart, sideStart chan struct{}, readRuns, sideRuns *int32) {
	t.Helper()
	readStart = make(chan struct{})
	sideStart = make(chan struct{})
	var rr, sr int32
	read = &tools.Tool{Name: "rotool", ReadOnly: true, ConcurrencySafe: true, Run: func(context.Context, string) (tools.Result, error) {
		atomic.AddInt32(&rr, 1)
		close(readStart)
		return tools.TextResult("read-ok"), nil
	}}
	side = &tools.Tool{Name: "sidetool", SideEffect: true, Kind: tools.KindExec, Run: func(context.Context, string) (tools.Result, error) {
		atomic.AddInt32(&sr, 1)
		close(sideStart)
		return tools.TextResult("side-ok"), nil
	}}
	return read, side, readStart, sideStart, &rr, &sr
}

// TestReadOnlyToolStartsDuringStream：只读工具在 provider 还没返回时就已开跑；
// 消息落盘后按原顺序回填，且不重复执行。
func TestReadOnlyToolStartsDuringStream(t *testing.T) {
	read, side, readStart, sideStart, readRuns, sideRuns := streamTools(t)
	prov := &streamToolProvider{
		read:      llm.ToolCall{ID: "r1", Name: "rotool", Input: "{}"},
		side:      llm.ToolCall{ID: "s1", Name: "sidetool", Input: "{}"},
		readStart: readStart,
		sideStart: sideStart,
	}
	a, sess := newTestAgent(t, read, side)
	a.Provider = prov

	if err := a.Run(context.Background(), "go"); err != nil {
		t.Fatal(err)
	}
	if got := atomic.LoadInt32(readRuns); got != 1 {
		t.Fatalf("read tool ran %d times, want 1", got)
	}
	if got := atomic.LoadInt32(sideRuns); got != 1 {
		t.Fatalf("side tool ran %d times, want 1", got)
	}
	// 结果按调用顺序回填（r1 在 s1 前），且 assistant 消息在它们之前。
	var toolIDs []string
	for _, m := range sess.Messages {
		if m.Role == llm.RoleTool {
			toolIDs = append(toolIDs, m.ToolCallID)
		}
	}
	if len(toolIDs) != 2 || toolIDs[0] != "r1" || toolIDs[1] != "s1" {
		t.Fatalf("tool result order = %v", toolIDs)
	}
}

// TestStreamFailureDiscardsEarlyResults：流在提前执行之后失败，结果不能落盘
// （没有对应的 tool_use），重试/下一轮从干净历史开始；只读工具没有副作用。
func TestStreamFailureDiscardsEarlyResults(t *testing.T) {
	read, _, readStart, _, readRuns, _ := streamTools(t)
	prov := &streamToolProvider{
		read:       llm.ToolCall{ID: "r1", Name: "rotool", Input: "{}"},
		readStart:  readStart,
		failStream: true,
	}
	a, sess := newTestAgent(t, read)
	a.Provider = prov
	var events []ToolEvent
	a.Opt.OnTool = func(ev ToolEvent) { events = append(events, ev) }

	if err := a.Run(context.Background(), "go"); err == nil {
		t.Fatal("want stream error")
	}
	if got := atomic.LoadInt32(readRuns); got != 1 {
		t.Fatalf("read tool ran %d times, want 1", got)
	}
	if len(events) != 0 {
		t.Fatalf("failed attempt leaked tool events into the UI: %+v", events)
	}
	for _, m := range sess.Messages {
		if m.Role == llm.RoleTool || (m.Role == llm.RoleAssistant && len(m.ToolCalls) > 0) {
			t.Fatalf("failed attempt leaked tool data into the transcript: %+v", sess.Messages)
		}
	}
}

// TestStartEarlyToolEligibility：只有只读、非副作用、并发安全的工具能提前跑。
func TestStartEarlyToolEligibility(t *testing.T) {
	read, side, readStart, sideStart, _, _ := streamTools(t)
	a, _ := newTestAgent(t, read, side)
	early := map[string]chan earlyOutcome{}

	a.startEarlyTool(context.Background(), llm.ToolCall{ID: "s1", Name: "sidetool", Input: "{}"}, early)
	if _, ok := early["s1"]; ok {
		t.Fatal("side-effect tool must not start early")
	}
	select {
	case <-sideStart:
		t.Fatal("side-effect tool actually ran")
	default:
	}

	a.startEarlyTool(context.Background(), llm.ToolCall{ID: "r1", Name: "rotool", Input: "{}"}, early)
	if _, ok := early["r1"]; !ok {
		t.Fatal("read-only tool should start early")
	}
	select {
	case <-readStart:
	case <-time.After(2 * time.Second):
		t.Fatal("read-only tool did not start")
	}
}