Agent-Native/runtime/llm/sse_test.go

package llm

import (
	"context"
	"encoding/json"
	"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)
	}
	if end.FinishReason != "stop" {
		t.Fatalf("finish reason %q, want stop", end.FinishReason)
	}
}

func TestParseSSEFinishReasonLengthKeepsStopEnd(t *testing.T) {
	body := strings.Join([]string{
		`data: {"choices":[{"delta":{"content":"半"},"finish_reason":"length"}]}`,
		`data: [DONE]`,
		"",
	}, "\n")
	var end Event
	for ev := range ParseSSE(context.Background(), strings.NewReader(body)) {
		if ev.Kind == KindMessageEnd {
			end = ev
		}
	}
	if end.FinishReason != "length" {
		t.Fatalf("finish reason %q, want length", end.FinishReason)
	}
	if end.Message.StopReason != StopEnd {
		t.Fatalf("stop reason %s, want end so existing callers stay unchanged", end.Message.StopReason)
	}
}

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)
	}
	if end.TerminalObservation == nil || end.TerminalObservation.ReadOutcome != "done_marker" ||
		!end.TerminalObservation.ToolArgumentsValid || end.TerminalObservation.ToolFragments != 3 {
		t.Fatalf("valid multi-fragment terminal was not classified: %#v", end.TerminalObservation)
	}
	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 TestParseSSEToolFinishWithoutCallFails(t *testing.T) {
	for _, reason := range []string{"tool_calls", "function_call"} {
		t.Run(reason, func(t *testing.T) {
			body := strings.Join([]string{
				`data: {"choices":[{"delta":{},"finish_reason":"` + reason + `"}],"usage":{"prompt_tokens":12,"completion_tokens":5,"total_tokens":17}}`,
				`data: [DONE]`,
				"",
			}, "\n")
			var end Event
			for ev := range ParseSSE(context.Background(), strings.NewReader(body)) {
				if ev.Kind == KindMessageEnd {
					end = ev
				}
			}
			if end.Err == nil || end.Message.StopReason != StopError || end.Usage.TotalTokens != 17 ||
				!end.StreamComplete || len(end.Message.ToolCalls) != 0 {
				t.Fatalf("empty tool finish must fail and retain Usage: %#v", end)
			}
		})
	}
}

func TestParseSSELengthFailsUnfinishedToolCall(t *testing.T) {
	body := strings.Join([]string{
		`data: {"choices":[{"delta":{"tool_calls":[{"index":0,"id":"call_1","function":{"name":"read_file","arguments":"{\"path\":"}}]}}]}`,
		`data: {"choices":[{"delta":{},"finish_reason":"length"}]}`,
		`data: [DONE]`,
		"",
	}, "\n")
	var end Event
	for ev := range ParseSSE(context.Background(), strings.NewReader(body)) {
		if ev.Kind == KindMessageEnd {
			end = ev
		}
	}
	if end.Message.StopReason != StopError || len(end.Message.ToolCalls) != 0 || end.Err == nil {
		t.Fatalf("truncated tool call must fail without execution: %#v", end)
	}
	if end.FinishReason != "length" {
		t.Fatalf("finish %q", end.FinishReason)
	}
}

func TestParseSSEUncleanEndFailsToolCall(t *testing.T) {
	body := `data: {"choices":[{"delta":{"tool_calls":[{"index":0,"id":"call_1","function":{"name":"read_file","arguments":"{\"path\":"}}]}}]}`
	var end Event
	for ev := range ParseSSE(context.Background(), strings.NewReader(body)) {
		if ev.Kind == KindMessageEnd {
			end = ev
		}
	}
	if end.Kind != KindMessageEnd || end.Message.StopReason != StopError || len(end.Message.ToolCalls) != 0 || end.Err == nil {
		t.Fatalf("unclean stream must fail without executing tools: %#v", end)
	}
}

func TestParseSSEToolFinishWithoutDone(t *testing.T) {
	body := strings.Join([]string{
		`data: {"choices":[{"delta":{"tool_calls":[{"index":0,"id":"call_1","function":{"name":"echo","arguments":"{}"}}]}}]}`,
		`data: {"choices":[{"delta":{},"finish_reason":"tool_calls"}],"usage":{"prompt_tokens":12,"completion_tokens":5,"total_tokens":17}}`,
		"",
	}, "\n\n")
	var end Event
	for ev := range ParseSSE(context.Background(), strings.NewReader(body)) {
		if ev.Kind == KindMessageEnd {
			end = ev
		}
	}
	if end.Kind != KindMessageEnd || end.Message.StopReason != StopError ||
		len(end.Message.ToolCalls) != 0 || end.Err == nil || end.Retryable ||
		end.Usage.TotalTokens != 17 {
		t.Fatalf("unsafe tool call from unfinished stream: %#v", end)
	}
}

func TestTerminalObservationCleanEOFWithoutDone(t *testing.T) {
	body := strings.Join([]string{
		`data: {"choices":[{"delta":{"content":"PRIVATE_PROMPT_R05","tool_calls":[{"index":0,"id":"call_1","function":{"name":"PRIVATE_TOOL_R05","arguments":"{}"}}]}}]}`,
		`data: {"choices":[{"delta":{},"finish_reason":"tool_calls"}],"usage":{"prompt_tokens":12,"completion_tokens":5,"total_tokens":17}}`,
		"",
	}, "\n\n")
	var end Event
	for ev := range ParseSSE(context.Background(), strings.NewReader(body)) {
		if ev.Kind == KindMessageEnd {
			end = ev
		}
	}
	if end.Err == nil || end.StreamComplete || len(end.Message.ToolCalls) != 0 || end.Usage.TotalTokens != 17 {
		t.Fatalf("unfinished tool stream was released or unmetered: %#v", end)
	}
	meta := end.TerminalObservation
	if meta == nil || meta.ReadOutcome != "clean_eof" || !meta.LastFrameJSONComplete ||
		!meta.UsagePresent || meta.ToolFragments != 1 || meta.ToolArgumentBytes != 2 || !meta.ToolArgumentsValid {
		t.Fatalf("missing content-free clean EOF evidence: %#v", meta)
	}
	raw, err := json.Marshal(meta)
	if err != nil {
		t.Fatal(err)
	}
	if strings.Contains(string(raw), "PRIVATE_PROMPT_R05") || strings.Contains(string(raw), "PRIVATE_TOOL_R05") {
		t.Fatalf("terminal observation retained model content: %s", raw)
	}
}

func TestTerminalObservationMalformedAndReadError(t *testing.T) {
	usageAndTool := `data: {"choices":[{"delta":{"tool_calls":[{"index":0,"id":"call_1","function":{"name":"echo","arguments":"{}"}}]},"finish_reason":"tool_calls"}],"usage":{"prompt_tokens":12,"completion_tokens":5,"total_tokens":17}}` + "\n\n"
	tests := []struct {
		name       string
		ctx        context.Context
		reader     io.Reader
		outcome    string
		usage      int
		frameValid bool
	}{
		{"invalid-frame", context.Background(), strings.NewReader(usageAndTool + "data: {not-json}\n\n"), "invalid_frame", 17, false},
		{"read-error", context.Background(), &errorAfterDataReader{data: strings.NewReader(usageAndTool)}, "read_error", 17, true},
	}
	ctx, cancel := context.WithCancel(context.Background())
	cancel()
	tests = append(tests, struct {
		name       string
		ctx        context.Context
		reader     io.Reader
		outcome    string
		usage      int
		frameValid bool
	}{"cancelled", ctx, strings.NewReader(usageAndTool), "cancelled", 0, false})
	for _, tc := range tests {
		t.Run(tc.name, func(t *testing.T) {
			var end Event
			for ev := range ParseSSE(tc.ctx, tc.reader) {
				if ev.Kind == KindMessageEnd || ev.Kind == KindError {
					end = ev
				}
			}
			if end.Err == nil || len(end.Message.ToolCalls) != 0 || end.Usage.TotalTokens != tc.usage {
				t.Fatalf("failure lost Usage or released tools: %#v", end)
			}
			meta := end.TerminalObservation
			if meta == nil || meta.ReadOutcome != tc.outcome || meta.LastFrameJSONComplete != tc.frameValid {
				t.Fatalf("wrong terminal classification: %#v", meta)
			}
		})
	}
}

func TestParseSSEInvalidToolArguments(t *testing.T) {
	body := strings.Join([]string{
		`data: {"choices":[{"delta":{"tool_calls":[{"index":0,"id":"call_1","function":{"name":"publish_work_checkpoint","arguments":"{\"findings\":"}}]}}]}`,
		`data: {"choices":[{"delta":{},"finish_reason":"tool_calls"}],"usage":{"prompt_tokens":12,"completion_tokens":5,"total_tokens":17}}`,
		`data: [DONE]`,
		"",
	}, "\n\n")
	var end Event
	for ev := range ParseSSE(context.Background(), strings.NewReader(body)) {
		if ev.Kind == KindMessageEnd {
			end = ev
		}
	}
	if end.Kind != KindMessageEnd || end.Message.StopReason != StopError ||
		len(end.Message.ToolCalls) != 0 || end.Err == nil || end.Retryable ||
		end.Usage.TotalTokens != 17 {
		t.Fatalf("unsafe tool call with invalid JSON arguments: %#v", end)
	}
	if end.TerminalObservation == nil || end.TerminalObservation.ToolArgumentsValid ||
		end.TerminalObservation.ReadOutcome != "done_marker" {
		t.Fatalf("invalid arguments incorrectly classified: %#v", end.TerminalObservation)
	}
}

func TestParseSSERejectsMissingToolIdentity(t *testing.T) {
	for _, tc := range []struct{ name, firstChunk string }{
		{"missing-id", `data: {"choices":[{"delta":{"tool_calls":[{"index":0,"function":{"name":"echo","arguments":"{}"}}]}}]}`},
		{"missing-name", `data: {"choices":[{"delta":{"tool_calls":[{"index":0,"id":"call_1","function":{"arguments":"{}"}}]}}]}`},
	} {
		t.Run(tc.name, func(t *testing.T) {
			body := strings.Join([]string{
				tc.firstChunk,
				`data: {"choices":[{"delta":{},"finish_reason":"tool_calls"}],"usage":{"prompt_tokens":12,"completion_tokens":5,"total_tokens":17}}`,
				`data: [DONE]`,
				"",
			}, "\n\n")
			var end Event
			for ev := range ParseSSE(context.Background(), strings.NewReader(body)) {
				if ev.Kind == KindMessageEnd {
					end = ev
				}
			}
			if end.Message.StopReason != StopError || len(end.Message.ToolCalls) != 0 || end.Err == nil {
				t.Fatalf("missing tool identity reached execution: %#v", end)
			}
		})
	}
}

func TestParseSSERejectsNonObjectToolArguments(t *testing.T) {
	for _, tc := range []struct{ name, arguments string }{
		{"array", "[]"},
		{"null", "null"},
	} {
		t.Run(tc.name, func(t *testing.T) {
			body := strings.Join([]string{
				`data: {"choices":[{"delta":{"tool_calls":[{"index":0,"id":"call_1","function":{"name":"echo","arguments":"` + tc.arguments + `"}}]}}]}`,
				`data: {"choices":[{"delta":{},"finish_reason":"tool_calls"}],"usage":{"prompt_tokens":12,"completion_tokens":5,"total_tokens":17}}`,
				`data: [DONE]`,
				"",
			}, "\n\n")
			var end Event
			for ev := range ParseSSE(context.Background(), strings.NewReader(body)) {
				if ev.Kind == KindMessageEnd {
					end = ev
				}
			}
			if end.Message.StopReason != StopError || len(end.Message.ToolCalls) != 0 || end.Err == nil {
				t.Fatalf("non-object tool arguments reached execution: %#v", end)
			}
		})
	}
}

func TestParseSSECapacityErrorIsRetryable(t *testing.T) {
	body := "data: {\"error\":{\"message\":\"server is at capacity\"}}\n\n"
	var got Event
	for ev := range ParseSSE(context.Background(), strings.NewReader(body)) {
		got = ev
	}
	if got.Kind != KindError || !got.Retryable || !CapacityLimited(got.ErrorMessage) {
		t.Fatalf("event %#v", got)
	}
}

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 TestParseSSEMalformedChunkAfterUsageKeepsUsage(t *testing.T) {
	body := strings.Join([]string{
		`data: {"choices":[{"delta":{"tool_calls":[{"index":0,"id":"call_1","function":{"name":"echo","arguments":"{}"}}]},"finish_reason":"tool_calls"}],"usage":{"prompt_tokens":12,"completion_tokens":5,"total_tokens":17}}`,
		`data: {not-json}`,
		"",
	}, "\n\n")
	var end Event
	for ev := range ParseSSE(context.Background(), strings.NewReader(body)) {
		end = ev
	}
	if end.Kind != KindMessageEnd || end.Message.StopReason != StopError ||
		len(end.Message.ToolCalls) != 0 || end.Err == nil || end.Retryable ||
		end.Usage.TotalTokens != 17 {
		t.Fatalf("malformed chunk dropped already returned usage: %#v", end)
	}
}

func TestParseSSEProviderErrorAfterUsageKeepsUsage(t *testing.T) {
	body := strings.Join([]string{
		`data: {"choices":[{"delta":{"tool_calls":[{"index":0,"id":"call_1","function":{"name":"echo","arguments":"{}"}}]}}],"usage":{"prompt_tokens":12,"completion_tokens":5,"total_tokens":17}}`,
		`data: {"error":{"message":"server is at capacity"}}`,
		"",
	}, "\n\n")
	var end Event
	for ev := range ParseSSE(context.Background(), strings.NewReader(body)) {
		end = ev
	}
	if end.Kind != KindMessageEnd || end.Message.StopReason != StopError ||
		len(end.Message.ToolCalls) != 0 || end.Err == nil || end.Retryable ||
		end.Usage.TotalTokens != 17 {
		t.Fatalf("provider error after usage retried or lost its charge: %#v", end)
	}
}

type errorAfterDataReader struct{ data *strings.Reader }

func (r *errorAfterDataReader) Read(p []byte) (int, error) {
	n, err := r.data.Read(p)
	if err == io.EOF {
		return 0, io.ErrUnexpectedEOF
	}
	return n, err
}

func TestParseSSEReaderErrorAfterUsageKeepsUsage(t *testing.T) {
	body := strings.Join([]string{
		`data: {"choices":[{"delta":{"tool_calls":[{"index":0,"id":"call_1","function":{"name":"echo","arguments":"{}"}}]}}],"usage":{"prompt_tokens":12,"completion_tokens":5,"total_tokens":17}}`,
		"",
	}, "\n\n")
	var end Event
	for ev := range ParseSSE(context.Background(), &errorAfterDataReader{data: strings.NewReader(body)}) {
		end = ev
	}
	if end.Kind != KindMessageEnd || end.Message.StopReason != StopError ||
		len(end.Message.ToolCalls) != 0 || end.Err == nil || end.Retryable ||
		end.Usage.TotalTokens != 17 {
		t.Fatalf("reader error after usage retried or lost its charge: %#v", end)
	}
}

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")
	}
}

func TestParseSSEReasoningDelta(t *testing.T) {
	body := strings.Join([]string{
		`data: {"choices":[{"delta":{"reasoning_content":"先分析需求"},"index":0}]}`,
		`data: {"choices":[{"delta":{"reasoning_content":"再检索标准"},"index":0}]}`,
		`data: {"choices":[{"delta":{"content":"结论"},"finish_reason":"stop"}],"usage":{"prompt_tokens":1,"completion_tokens":2,"total_tokens":3}}`,
		`data: [DONE]`,
		"",
	}, "\n")
	var reasoning []string
	var text []string
	var end Event
	for ev := range ParseSSE(context.Background(), strings.NewReader(body)) {
		switch ev.Kind {
		case KindReasoningDelta:
			reasoning = append(reasoning, ev.ReasoningDelta)
		case KindTextDelta:
			text = append(text, ev.Text)
		case KindMessageEnd:
			end = ev
		}
	}
	if strings.Join(reasoning, "") != "先分析需求再检索标准" {
		t.Fatalf("reasoning deltas %v", reasoning)
	}
	// 正文与思考链分流：Content 只含 text，不含 reasoning
	if strings.Join(text, "") != "结论" {
		t.Fatalf("text deltas %v", text)
	}
	if end.Message.Content != "结论" {
		t.Fatalf("end content %q must not include reasoning", end.Message.Content)
	}
}