chore: import upstream snapshot with attribution

agent/stream_test.go

@@ -0,0 +1,310 @@
+package agent
+
+import (
+	"context"
+	"errors"
+	"io"
+	"testing"
+	"time"
+
+	"go-micro.dev/v6/ai"
+	"go-micro.dev/v6/flow"
+	"go-micro.dev/v6/store"
+)
+
+func TestStreamAskEmitsToolEventsAndFinalTokens(t *testing.T) {
+	calls := 0
+	fakeGen = func(ctx context.Context, opts ai.Options, req *ai.Request) (*ai.Response, error) {
+		if opts.ToolHandler == nil {
+			t.Fatal("StreamAsk must configure a tool handler")
+		}
+		calls++
+		result := opts.ToolHandler(ctx, ai.ToolCall{ID: "call-1", Name: "echo", Input: map[string]any{"text": "hello"}})
+		return &ai.Response{
+			Reply:     "planning",
+			Answer:    "final answer",
+			ToolCalls: []ai.ToolCall{{ID: "call-1", Name: "echo", Input: map[string]any{"text": "hello"}, Result: result.Content}},
+		}, nil
+	}
+	defer func() { fakeGen = nil }()
+
+	a := newTestAgent(Name("streamer"), WithTool("echo", "echo text", nil, func(ctx context.Context, input map[string]any) (string, error) {
+		return input["text"].(string), nil
+	}))
+	stream, err := a.StreamAsk(context.Background(), "say hello")
+	if err != nil {
+		t.Fatalf("StreamAsk: %v", err)
+	}
+
+	var types []StreamEventType
+	var tokens string
+	var done *Response
+	for {
+		event, err := stream.Recv()
+		if errors.Is(err, io.EOF) {
+			break
+		}
+		if err != nil {
+			t.Fatalf("Recv: %v", err)
+		}
+		types = append(types, event.Type)
+		if event.Type == StreamEventToken {
+			tokens += event.Token
+		}
+		if event.Type == StreamEventDone {
+			done = event.Response
+		}
+	}
+
+	want := []StreamEventType{StreamEventToolStart, StreamEventToolEnd, StreamEventToken, StreamEventToken, StreamEventToken, StreamEventDone}
+	if len(types) != len(want) {
+		t.Fatalf("event types = %v, want %v", types, want)
+	}
+	for i := range want {
+		if types[i] != want[i] {
+			t.Fatalf("event types = %v, want %v", types, want)
+		}
+	}
+	if tokens != "planning final answer" {
+		t.Fatalf("tokens = %q", tokens)
+	}
+	if done == nil || done.Reply != "planning\n\nfinal answer" {
+		t.Fatalf("done response = %#v", done)
+	}
+	if calls != 1 {
+		t.Fatalf("Generate calls = %d, want 1", calls)
+	}
+}
+
+func TestStreamAskCloseCancelsInFlightModelCall(t *testing.T) {
+	started := make(chan struct{})
+	fakeGen = func(ctx context.Context, opts ai.Options, req *ai.Request) (*ai.Response, error) {
+		close(started)
+		<-ctx.Done()
+		return nil, ctx.Err()
+	}
+	defer func() { fakeGen = nil }()
+
+	a := newTestAgent(Name("stream-cancel"))
+	stream, err := a.StreamAsk(context.Background(), "cancel me")
+	if err != nil {
+		t.Fatalf("StreamAsk: %v", err)
+	}
+
+	select {
+	case <-started:
+	case <-time.After(time.Second):
+		t.Fatal("model call did not start")
+	}
+
+	closed := make(chan error, 1)
+	go func() { closed <- stream.Close() }()
+	select {
+	case err := <-closed:
+		if err != nil {
+			t.Fatalf("Close: %v", err)
+		}
+	case <-time.After(time.Second):
+		t.Fatal("Close did not cancel the in-flight stream")
+	}
+}
+
+func TestStreamAskHelperRejectsUnsupportedAgent(t *testing.T) {
+	_, err := StreamAsk(context.Background(), unsupportedAgent{}, "hello")
+	if err == nil {
+		t.Fatal("StreamAsk helper should reject unsupported implementations")
+	}
+}
+
+func TestAgentStreamUsesProviderStreamingAndRecordsAssistantMemory(t *testing.T) {
+	var sawRequest bool
+	var sawRunInfo bool
+	fakeStream = func(ctx context.Context, opts ai.Options, req *ai.Request) (ai.Stream, error) {
+		sawRequest = true
+		info, ok := ai.RunInfoFrom(ctx)
+		if !ok || info.RunID == "" || info.Agent != "provider-stream" {
+			t.Fatalf("RunInfo = %#v, %v; want provider stream run metadata", info, ok)
+		}
+		sawRunInfo = true
+		if req.Prompt != "stream the answer" {
+			t.Fatalf("Prompt = %q, want stream the answer", req.Prompt)
+		}
+		if len(req.Messages) != 1 || req.Messages[0].Role != "user" || req.Messages[0].Content != "stream the answer" {
+			t.Fatalf("Messages = %#v, want current user turn in memory", req.Messages)
+		}
+		return &sliceStream{chunks: []string{"hel", "lo"}}, nil
+	}
+	defer func() { fakeStream = nil }()
+
+	a := newTestAgent(Name("provider-stream"))
+	stream, err := a.Stream(context.Background(), "stream the answer")
+	if err != nil {
+		t.Fatalf("Stream: %v", err)
+	}
+
+	var reply string
+	for {
+		chunk, err := stream.Recv()
+		if errors.Is(err, io.EOF) {
+			break
+		}
+		if err != nil {
+			t.Fatalf("Recv: %v", err)
+		}
+		reply += chunk.Reply
+	}
+	if err := stream.Close(); err != nil {
+		t.Fatalf("Close: %v", err)
+	}
+	if !sawRequest {
+		t.Fatal("provider Stream was not called")
+	}
+	if !sawRunInfo {
+		t.Fatal("provider Stream did not receive RunInfo")
+	}
+	if reply != "hello" {
+		t.Fatalf("reply = %q, want hello", reply)
+	}
+	got := a.mem.Messages()
+	if len(got) != 2 || got[0].Role != "user" || got[0].Content != "stream the answer" || got[1].Role != "assistant" || got[1].Content != "hello" {
+		t.Fatalf("memory = %#v, want user turn and streamed assistant reply", got)
+	}
+}
+
+func TestAgentStreamCanceledContextSkipsProviderCallAndMemory(t *testing.T) {
+	calls := 0
+	fakeStream = func(ctx context.Context, opts ai.Options, req *ai.Request) (ai.Stream, error) {
+		calls++
+		return &sliceStream{chunks: []string{"late"}}, nil
+	}
+	defer func() { fakeStream = nil }()
+
+	ctx, cancel := context.WithCancel(context.Background())
+	cancel()
+	a := newTestAgent(Name("provider-stream-cancel"))
+	_, err := a.Stream(ctx, "do not start")
+	if !errors.Is(err, context.Canceled) {
+		t.Fatalf("Stream error = %v, want context canceled", err)
+	}
+	if calls != 0 {
+		t.Fatalf("provider Stream calls = %d, want 0 after caller cancellation", calls)
+	}
+	if got := a.mem.Messages(); len(got) != 0 {
+		t.Fatalf("memory = %#v, want no recorded canceled stream turn", got)
+	}
+}
+
+func TestResumeStreamAskDoesNotReplayCompletedTool(t *testing.T) {
+	ctx := context.Background()
+	cp := flow.StoreCheckpoint(store.NewStore(), "stream-resume-agent")
+	toolRuns := 0
+	first := true
+	fakeGen = func(ctx context.Context, opts ai.Options, req *ai.Request) (*ai.Response, error) {
+		if opts.ToolHandler != nil {
+			res := opts.ToolHandler(ctx, ai.ToolCall{ID: "call-1", Name: "charge", Input: map[string]any{"order": "42"}})
+			if res.Content != "charged" {
+				t.Fatalf("tool result = %q, want charged", res.Content)
+			}
+		}
+		if first {
+			first = false
+			return nil, errors.New("stream disconnected after tool")
+		}
+		return &ai.Response{Reply: "finished from streamed checkpoint"}, nil
+	}
+	defer func() { fakeGen = nil }()
+
+	a := newTestAgent(Name("stream-resume-agent"), WithCheckpoint(cp),
+		WithTool("charge", "charge once", nil, func(context.Context, map[string]any) (string, error) {
+			toolRuns++
+			return "charged", nil
+		}))
+	stream, err := a.StreamAsk(ctx, "charge order 42")
+	if err != nil {
+		t.Fatalf("StreamAsk: %v", err)
+	}
+	for {
+		_, err := stream.Recv()
+		if errors.Is(err, io.EOF) {
+			break
+		}
+		if err != nil {
+			break
+		}
+	}
+	if toolRuns != 1 {
+		t.Fatalf("tool executions after failed StreamAsk = %d, want 1", toolRuns)
+	}
+	runs, err := Pending(ctx, a)
+	if err != nil {
+		t.Fatalf("Pending: %v", err)
+	}
+	if len(runs) != 1 {
+		t.Fatalf("Pending returned %d runs, want 1", len(runs))
+	}
+
+	resumed, err := ResumeStreamAsk(ctx, a, runs[0].ID)
+	if err != nil {
+		t.Fatalf("ResumeStreamAsk: %v", err)
+	}
+	var toolEvents int
+	var done *Response
+	for {
+		event, err := resumed.Recv()
+		if errors.Is(err, io.EOF) {
+			break
+		}
+		if err != nil {
+			t.Fatalf("resumed Recv: %v", err)
+		}
+		if event.Type == StreamEventToolStart || event.Type == StreamEventToolEnd {
+			toolEvents++
+		}
+		if event.Type == StreamEventDone {
+			done = event.Response
+		}
+	}
+	if toolRuns != 1 {
+		t.Fatalf("tool executions after ResumeStreamAsk = %d, want completed tool was not replayed", toolRuns)
+	}
+	if toolEvents != 2 {
+		t.Fatalf("resumed tool events = %d, want start/end for replayed checkpoint result", toolEvents)
+	}
+	if done == nil || done.Reply != "finished from streamed checkpoint" || done.RunID != runs[0].ID {
+		t.Fatalf("done response = %#v", done)
+	}
+}
+
+func TestAgentStreamDoesNotRecordUserWhenProviderStreamingUnsupported(t *testing.T) {
+	fakeStream = func(ctx context.Context, opts ai.Options, req *ai.Request) (ai.Stream, error) {
+		if len(req.Messages) == 0 || req.Messages[len(req.Messages)-1].Role != "user" || req.Messages[len(req.Messages)-1].Content != "stream fallback" {
+			t.Fatalf("stream request messages = %+v, want pending user message", req.Messages)
+		}
+		return nil, ai.ErrStreamingUnsupported
+	}
+	defer func() { fakeStream = nil }()
+
+	mem := NewInMemory(8)
+	a := newTestAgent(Name("stream-fallback"), WithMemory(mem), WithTool("echo", "echo text", nil, func(context.Context, map[string]any) (string, error) {
+		return "ok", nil
+	}))
+
+	_, err := a.Stream(context.Background(), "stream fallback")
+	if !errors.Is(err, ai.ErrStreamingUnsupported) {
+		t.Fatalf("Stream error = %v, want ErrStreamingUnsupported", err)
+	}
+	if got := mem.Messages(); len(got) != 0 {
+		t.Fatalf("memory after unsupported stream = %+v, want no recorded messages", got)
+	}
+}
+
+type unsupportedAgent struct{}
+
+func (unsupportedAgent) Name() string                                      { return "unsupported" }
+func (unsupportedAgent) Init(...Option)                                    {}
+func (unsupportedAgent) Options() Options                                  { return Options{} }
+func (unsupportedAgent) Ask(context.Context, string) (*Response, error)    { return nil, nil }
+func (unsupportedAgent) Stream(context.Context, string) (ai.Stream, error) { return nil, nil }
+func (unsupportedAgent) Run() error                                        { return nil }
+func (unsupportedAgent) Stop() error                                       { return nil }
+func (unsupportedAgent) String() string                                    { return "unsupported" }