Cyrene_v2/internal/runtime/model/anthropic_stream.go

package model

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

// anthropicStreamEvent is one SSE payload from /v1/messages with stream: true.
// Anthropic carries the event kind both in the SSE `event:` field and inside the
// JSON `type` field, so the parser keys off the JSON payload type.
type anthropicStreamEvent struct {
	Type    string                `json:"type"`
	Message *anthropicStartBlock  `json:"message,omitempty"`
	Index   *int                  `json:"index,omitempty"`
	Content *anthropicBlock       `json:"content_block,omitempty"`
	Delta   *anthropicDelta       `json:"delta,omitempty"`
	Usage   *anthropicUsage       `json:"usage,omitempty"`
	Error   *anthropicStreamError `json:"error,omitempty"`
}

// anthropicStartBlock is the payload of a message_start event.
type anthropicStartBlock struct {
	Model string          `json:"model"`
	Usage *anthropicUsage `json:"usage,omitempty"`
}

// anthropicDelta is the incremental payload inside a content_block_delta event.
type anthropicDelta struct {
	Type        string `json:"type"` // text_delta | input_json_delta
	Text        string `json:"text,omitempty"`
	PartialJSON string `json:"partial_json,omitempty"`
}

// anthropicStreamError is the payload of an error event.
type anthropicStreamError struct {
	Type    string `json:"type"`
	Message string `json:"message"`
}

// streamBlock tracks the partial state of one content block during streaming.
type streamBlock struct {
	index       int
	typ         string
	id          string
	name        string
	text        strings.Builder
	input       strings.Builder
	initialArgs map[string]any
}

// scanAnthropicSSE incrementally decodes Anthropic's content-block stream. Text
// deltas are forwarded as text_delta events (segment bracketed by
// text_start/text_end); tool_use input_json_delta fragments are forwarded as
// tool_call_delta events (segment bracketed by tool_call_start/tool_call_end)
// and accumulated per block so they re-assemble into a structured Args map when
// the block stops. The stream must terminate with a message_stop event,
// otherwise ErrStream is returned.
func scanAnthropicSSE(r io.Reader, onEvent func(StreamEvent)) (Response, error) {
	scanner := bufio.NewScanner(r)
	scanner.Buffer(make([]byte, 0, sseBufInit), maxSSELineBytes)
	agg := Response{}
	blocks := map[int]*streamBlock{}
	var order []int
	var text strings.Builder
	sawMessageStop := false

	for scanner.Scan() {
		line := strings.TrimSpace(scanner.Text())
		if line == "" || !strings.HasPrefix(line, sseDataPrefix) {
			continue
		}
		data := strings.TrimSpace(strings.TrimPrefix(line, sseDataPrefix))
		if data == "" {
			continue
		}
		var evt anthropicStreamEvent
		if err := json.Unmarshal([]byte(data), &evt); err != nil {
			return Response{}, fmt.Errorf("model: %w: decode stream chunk: %v", ErrUpstream, err)
		}
		if err := applyAnthropicEvent(&agg, blocks, &order, &text, evt, onEvent); err != nil {
			return Response{}, err
		}
		if evt.Type == "message_stop" {
			sawMessageStop = true
		}
	}
	if err := scanner.Err(); err != nil {
		return Response{}, fmt.Errorf("model: %w: read stream: %v", ErrUpstream, err)
	}
	if !sawMessageStop {
		return Response{}, ErrStream
	}
	finalizeAnthropicStream(&agg, blocks, order, text.String())
	return agg, nil
}

// applyAnthropicEvent folds a single stream event into the running aggregation.
func applyAnthropicEvent(
	agg *Response,
	blocks map[int]*streamBlock,
	order *[]int,
	text *strings.Builder,
	evt anthropicStreamEvent,
	onEvent func(StreamEvent),
) error {
	switch evt.Type {
	case "message_start":
		if evt.Message != nil {
			agg.Model = evt.Message.Model
			if evt.Message.Usage != nil {
				agg.PromptTokens = evt.Message.Usage.InputTokens
			}
		}
	case "content_block_start":
		if evt.Index == nil || evt.Content == nil {
			return nil
		}
		b := &streamBlock{
			index:       *evt.Index,
			typ:         evt.Content.Type,
			id:          evt.Content.ID,
			name:        evt.Content.Name,
			initialArgs: evt.Content.Input,
		}
		switch evt.Content.Type {
		case "text":
			if onEvent != nil {
				onEvent(StreamEvent{Kind: StreamKindTextStart, Index: *evt.Index})
			}
			if evt.Content.Text != "" {
				b.text.WriteString(evt.Content.Text)
				text.WriteString(evt.Content.Text)
				if onEvent != nil {
					onEvent(StreamEvent{Kind: StreamKindTextDelta, Delta: evt.Content.Text, Index: *evt.Index})
				}
			}
		case "tool_use":
			if onEvent != nil {
				onEvent(StreamEvent{Kind: StreamKindToolCallStart, Index: *evt.Index, ID: evt.Content.ID})
			}
		}
		if _, ok := blocks[*evt.Index]; !ok {
			*order = append(*order, *evt.Index)
		}
		blocks[*evt.Index] = b
	case "content_block_delta":
		if evt.Index == nil || evt.Delta == nil {
			return nil
		}
		b := blocks[*evt.Index]
		if b == nil {
			return nil
		}
		switch evt.Delta.Type {
		case "text_delta":
			b.text.WriteString(evt.Delta.Text)
			text.WriteString(evt.Delta.Text)
			if onEvent != nil {
				onEvent(StreamEvent{Kind: StreamKindTextDelta, Delta: evt.Delta.Text, Index: *evt.Index})
			}
		case "input_json_delta":
			b.input.WriteString(evt.Delta.PartialJSON)
			if onEvent != nil {
				onEvent(StreamEvent{Kind: StreamKindToolCallDelta, Delta: evt.Delta.PartialJSON, Index: *evt.Index, ID: b.id})
			}
		}
	case "content_block_stop":
		if evt.Index == nil {
			return nil
		}
		b := blocks[*evt.Index]
		if b == nil {
			return nil
		}
		if b.typ == "text" {
			if onEvent != nil {
				onEvent(StreamEvent{Kind: StreamKindTextEnd, Index: *evt.Index})
			}
		} else if b.typ == "tool_use" {
			if onEvent != nil {
				onEvent(StreamEvent{Kind: StreamKindToolCallEnd, Index: *evt.Index, ID: b.id})
			}
		}
	case "message_delta":
		if evt.Usage != nil {
			agg.CompletionTokens = evt.Usage.OutputTokens
		}
	case "error":
		if evt.Error != nil {
			return fmt.Errorf("model: %w: anthropic stream error: %s: %s", ErrUpstream, evt.Error.Type, evt.Error.Message)
		}
	}
	return nil
}

// finalizeAnthropicStream assembles the aggregated Response from the accumulated
// blocks: concatenated text plus any tool_use ToolCalls derived per block.
func finalizeAnthropicStream(agg *Response, blocks map[int]*streamBlock, order []int, content string) {
	agg.Content = content
	agg.Message = Message{Role: RoleAssistant, Content: []Content{{Type: ContentTypeText, Text: content}}}
	agg.TotalTokens = agg.PromptTokens + agg.CompletionTokens
	for _, idx := range order {
		b := blocks[idx]
		if b == nil || b.typ != "tool_use" {
			continue
		}
		args := b.initialArgs
		if b.input.Len() > 0 {
			decoded := map[string]any{}
			if err := json.Unmarshal([]byte(b.input.String()), &decoded); err == nil {
				args = decoded
			} else {
				args = map[string]any{}
			}
		}
		agg.ToolCalls = append(agg.ToolCalls, ToolCall{Name: b.name, Args: args})
	}
}