Agent-Native/runtime/llm/sse.go

package llm

import (
	"bufio"
	"context"
	"encoding/json"
	"io"
	"strings"
)

type sseChoice struct {
	Delta struct {
		Content   string `json:"content"`
		Reasoning string `json:"reasoning_content"`
		ToolCalls []struct {
			Index    int    `json:"index"`
			ID       string `json:"id"`
			Function struct {
				Name      string `json:"name"`
				Arguments string `json:"arguments"`
			} `json:"function"`
		} `json:"tool_calls"`
	} `json:"delta"`
	FinishReason string `json:"finish_reason"`
}

type sseChunk struct {
	Choices []sseChoice `json:"choices"`
	Usage   *usageWire  `json:"usage"`
	Error   *struct {
		Message string `json:"message"`
	} `json:"error"`
}

type usageWire struct {
	Usage
	Details struct {
		Cached int `json:"cached_tokens"`
	} `json:"prompt_tokens_details"`
}

// ParseSSE reads an OpenAI-compatible server-sent event body and emits
// deltas plus one terminal message_end or error event.
func ParseSSE(ctx context.Context, r io.Reader) <-chan Event {
	out := make(chan Event)
	go func() {
		defer close(out)
		emitParsed(ctx, r, out)
	}()
	return out
}

func emitParsed(ctx context.Context, r io.Reader, out chan<- Event) {
	var content strings.Builder
	calls := map[int]*ToolCall{}
	order := []int{}
	var usage Usage
	finished := ""
	sc := bufio.NewScanner(r)
	sc.Buffer(make([]byte, 0, 64*1024), 4*1024*1024)
	reader := &sseEventReader{sc: sc}
	for {
		data, ok := reader.next()
		if !ok {
			break
		}
		if ctx.Err() != nil {
			emitAbort(out, content.String(), ctx.Err())
			return
		}
		if data == "[DONE]" {
			break
		}
		var chunk sseChunk
		if err := json.Unmarshal([]byte(data), &chunk); err != nil {
			emitError(out, content.String(), false, "malformed stream chunk: "+err.Error())
			return
		}
		if chunk.Error != nil && chunk.Error.Message != "" {
			emitError(out, content.String(), false, chunk.Error.Message)
			return
		}
		if chunk.Usage != nil {
			usage = chunk.Usage.Usage
			if chunk.Usage.Details.Cached > 0 {
				usage.CacheReadTokens = chunk.Usage.Details.Cached
			}
		}
		if len(chunk.Choices) == 0 {
			continue
		}
		choice := chunk.Choices[0]
		if choice.FinishReason != "" {
			finished = choice.FinishReason
		}
		if choice.Delta.Content != "" {
			content.WriteString(choice.Delta.Content)
			select {
			case <-ctx.Done():
				emitAbort(out, content.String(), ctx.Err())
				return
			case out <- Event{Kind: KindTextDelta, Text: choice.Delta.Content}:
			}
		}
		// reasoning_content（国产 OpenAI 兼容模型的思考链）与正文分流：
		// 只作为 reasoning_delta 事件外发，不计入 Message.Content（与 runner
		// stripThinkBlocks"正文无 think"语义对齐）。
		if choice.Delta.Reasoning != "" {
			select {
			case <-ctx.Done():
				emitAbort(out, content.String(), ctx.Err())
				return
			case out <- Event{Kind: KindReasoningDelta, ReasoningDelta: choice.Delta.Reasoning}:
			}
		}
		for _, tc := range choice.Delta.ToolCalls {
			cur, ok := calls[tc.Index]
			if !ok {
				cur = &ToolCall{}
				calls[tc.Index] = cur
				order = append(order, tc.Index)
			}
			if tc.ID != "" {
				cur.ID = tc.ID
			}
			if tc.Function.Name != "" {
				cur.Name += tc.Function.Name
			}
			if tc.Function.Arguments != "" {
				cur.Arguments += tc.Function.Arguments
			}
			select {
			case <-ctx.Done():
				emitAbort(out, content.String(), ctx.Err())
				return
			case out <- Event{
				Kind:           KindToolCallDelta,
				ToolIndex:      tc.Index,
				ToolID:         tc.ID,
				ToolName:       tc.Function.Name,
				ArgumentsDelta: tc.Function.Arguments,
			}:
			}
		}
	}
	if err := sc.Err(); err != nil {
		if ctx.Err() != nil {
			emitAbort(out, content.String(), ctx.Err())
			return
		}
		emitError(out, content.String(), true, err.Error())
		return
	}
	// 取消必须压过正常结束。select 在两个分支都就绪时随机挑选，
	// 竞态下会把 message_end 发给已经取消的调用方。
	if ctx.Err() != nil {
		emitAbort(out, content.String(), ctx.Err())
		return
	}
	msg := Message{Role: RoleAssistant, Content: content.String(), StopReason: mapFinish(finished, len(order) > 0)}
	for _, idx := range order {
		msg.ToolCalls = append(msg.ToolCalls, *calls[idx])
	}
	if msg.StopReason == StopToolCalls && len(msg.ToolCalls) == 0 {
		msg.StopReason = StopEnd
	}
	select {
	case <-ctx.Done():
		emitAbort(out, content.String(), ctx.Err())
	case out <- Event{Kind: KindMessageEnd, Message: msg, Usage: usage, FinishReason: finished}:
	}
}

// sseEventReader 按空行切事件，并把同一事件里的多行 data: 用换行拼回。
// 已经是完整 JSON 的 data 行单独成事件，这样单行流不用空行也能解析。
type sseEventReader struct {
	sc    *bufio.Scanner
	carry string
}

func (r *sseEventReader) next() (string, bool) {
	var parts []string
	nextPiece := func() (string, bool) {
		if r.carry != "" {
			piece := r.carry
			r.carry = ""
			return piece, true
		}
		if !r.sc.Scan() {
			return "", false
		}
		return r.sc.Text(), true
	}
	for {
		line, ok := nextPiece()
		if !ok {
			break
		}
		if strings.TrimSpace(line) == "" {
			if len(parts) == 0 {
				continue
			}
			return strings.Join(parts, "\n"), true
		}
		trimmed := strings.TrimSpace(line)
		if strings.HasPrefix(trimmed, ":") {
			continue
		}
		if !strings.HasPrefix(trimmed, "data:") {
			continue
		}
		piece := strings.TrimSpace(strings.TrimPrefix(trimmed, "data:"))
		if len(parts) > 0 && json.Valid([]byte(strings.Join(parts, "\n"))) {
			r.carry = line
			return strings.Join(parts, "\n"), true
		}
		parts = append(parts, piece)
	}
	if len(parts) > 0 {
		return strings.Join(parts, "\n"), true
	}
	return "", false
}

func mapFinish(reason string, hasTools bool) StopReason {
	switch reason {
	case "tool_calls", "function_call":
		return StopToolCalls
	case "content_filter", "error":
		return StopError
	case "stop", "length", "":
		if hasTools {
			return StopToolCalls
		}
		return StopEnd
	default:
		if hasTools {
			return StopToolCalls
		}
		return StopEnd
	}
}

func emitError(out chan<- Event, partial string, retryable bool, message string) {
	out <- Event{
		Kind:         KindError,
		Retryable:    retryable,
		ErrorMessage: message,
		Err:          &StreamError{Message: message, Retryable: retryable},
		Message: Message{
			Role:       RoleAssistant,
			Content:    partial,
			StopReason: StopError,
		},
	}
}

func emitAbort(out chan<- Event, partial string, err error) {
	msg := "aborted"
	if err != nil {
		msg = err.Error()
	}
	out <- Event{
		Kind:         KindError,
		Retryable:    false,
		ErrorMessage: msg,
		Err:          err,
		Message: Message{
			Role:       RoleAssistant,
			Content:    partial,
			StopReason: StopAborted,
		},
	}
}

// StreamError is a provider or transport failure encoded on the event stream.
type StreamError struct {
	Message   string
	Retryable bool
}

func (e *StreamError) Error() string {
	if e == nil {
		return "stream error"
	}
	return e.Message
}