github.com/anthropics/anthropic-sdk-go@v1.70.1/packages/ssestream/ssestream.go

// File generated from our OpenAPI spec by Stainless. See CONTRIBUTING.md for details.

package ssestream

import (
	"bufio"
	"bytes"
	"encoding/json"
	"fmt"
	"io"
	"net/http"
	"strings"

	"github.com/anthropics/anthropic-sdk-go/internal/apierror"
)

type Decoder interface {
	Event() Event
	Next() bool
	Close() error
	Err() error
}

func NewDecoder(res *http.Response) Decoder {
	if res == nil || res.Body == nil {
		return nil
	}

	var decoder Decoder
	contentType := res.Header.Get("content-type")
	if t, ok := decoderTypes[contentType]; ok {
		decoder = t(res.Body)
	} else {
		scn := bufio.NewScanner(res.Body)
		scn.Buffer(nil, bufio.MaxScanTokenSize<<9)
		decoder = &eventStreamDecoder{rc: res.Body, scn: scn}
	}

	// richDecoder needs the http request to provide helpful errors
	// Aside from the error case there should be no difference
	// from the underlying decoder.
	if res.Request != nil {
		return richErrorDecoder{Decoder: decoder, resp: res}
	}

	return decoder
}

// richErrorDecoder wraps a Decoder and carries the original [*http.Response]
// so that it can construct rich API errors from SSE error events.
//
// A richErrorDecoder is used by [Stream] to construct errors whenever possible.
// [Stream] needs to rely on the underlying decoder, because the [Stream] has
// no access to the original http request.
//
// There should be no other differences from the underlying decoder.
type richErrorDecoder struct {
	Decoder
	resp *http.Response
}

func (d *richErrorDecoder) newAPIError(errorJSON []byte) error {
	aerr := &apierror.Error{}
	if d.resp != nil {
		aerr.Request = d.resp.Request
		aerr.Response = d.resp
		aerr.StatusCode = d.resp.StatusCode
		aerr.RequestID = d.resp.Header.Get("request-id")
		aerr.WorkspaceID = d.resp.Header.Get("anthropic-workspace-id")
	}
	if aerr.UnmarshalJSON(errorJSON) != nil {
		return fmt.Errorf("received error while streaming: %s", string(errorJSON))
	}
	return aerr
}

var decoderTypes = map[string](func(io.ReadCloser) Decoder){}

func RegisterDecoder(contentType string, decoder func(io.ReadCloser) Decoder) {
	decoderTypes[strings.ToLower(contentType)] = decoder
}

type Event struct {
	Type string
	Data []byte
}

// A base implementation of a Decoder for text/event-stream.
type eventStreamDecoder struct {
	evt Event
	rc  io.ReadCloser
	scn *bufio.Scanner
	err error
}

func (s *eventStreamDecoder) Next() bool {
	if s.err != nil {
		return false
	}

	event := ""
	data := bytes.NewBuffer(nil)

	for s.scn.Scan() {
		txt := s.scn.Bytes()

		// Dispatch event on an empty line
		if len(txt) == 0 {
			s.evt = Event{
				Type: event,
				Data: data.Bytes(),
			}
			return true
		}

		// Split a string like "event: bar" into name="event" and value=" bar".
		name, value, _ := bytes.Cut(txt, []byte(":"))

		// Consume an optional space after the colon if it exists.
		if len(value) > 0 && value[0] == ' ' {
			value = value[1:]
		}

		switch string(name) {
		case "":
			// An empty line in the for ": something" is a comment and should be ignored.
			continue
		case "event":
			event = string(value)
		case "data":
			_, s.err = data.Write(value)
			if s.err != nil {
				break
			}
			_, s.err = data.WriteRune('\n')
			if s.err != nil {
				break
			}
		}
	}

	if s.scn.Err() != nil {
		s.err = s.scn.Err()
	}

	return false
}

func (s *eventStreamDecoder) Event() Event {
	return s.evt
}

func (s *eventStreamDecoder) Close() error {
	return s.rc.Close()
}

func (s *eventStreamDecoder) Err() error {
	return s.err
}

type Stream[T any] struct {
	decoder Decoder
	cur     T
	err     error
}

func NewStream[T any](decoder Decoder, err error) *Stream[T] {
	return &Stream[T]{
		decoder: decoder,
		err:     err,
	}
}

// Next returns false if the stream has ended or an error occurred.
// Call Stream.Current() to get the current value.
// Call Stream.Err() to get the error.
//
//		for stream.Next() {
//			data := stream.Current()
//		}
//
//	 	if stream.Err() != nil {
//			...
//	 	}
func (s *Stream[T]) Next() bool {
	if s.err != nil {
		return false
	}

	for s.decoder.Next() {
		switch s.decoder.Event().Type {
		case "completion":
			var nxt T
			s.err = json.Unmarshal(s.decoder.Event().Data, &nxt)
			if s.err != nil {
				return false
			}
			s.cur = nxt
			return true
		case "message_start", "message_delta", "message_stop", "content_block_start", "content_block_delta", "content_block_stop", "message", "user.message", "user.interrupt", "user.tool_confirmation", "user.custom_tool_result", "user.tool_result", "agent.message", "agent.thinking", "agent.tool_use", "agent.tool_result", "agent.mcp_tool_use", "agent.mcp_tool_result", "agent.custom_tool_use", "agent.thread_context_compacted", "session.status_running", "session.status_idle", "session.status_rescheduled", "session.status_terminated", "session.error", "session.deleted", "session.updated", "span.model_request_start", "span.model_request_end", "span.outcome_evaluation_start", "span.outcome_evaluation_ongoing", "span.outcome_evaluation_end", "user.define_outcome", "agent.thread_message_received", "agent.thread_message_sent", "agent.session_thread_message_received", "agent.session_thread_message_sent", "session.thread_created", "session.thread_status_created", "session.thread_status_running", "session.thread_status_idle", "session.thread_status_rescheduled", "session.thread_status_terminated", "event_start", "event_delta", "system.message":
			var nxt T
			s.err = json.Unmarshal(s.decoder.Event().Data, &nxt)
			if s.err != nil {
				return false
			}
			s.cur = nxt
			return true
		case "ping":
			continue
		case "error":
			data := s.decoder.Event().Data
			if ed, ok := s.decoder.(richErrorDecoder); ok {
				s.err = ed.newAPIError(data)
			} else {
				s.err = fmt.Errorf("received error while streaming: %s", string(data))
			}
			return false
		}
	}

	// decoder.Next() may be false because of an error
	s.err = s.decoder.Err()

	return false
}

func (s *Stream[T]) Current() T {
	return s.cur
}

func (s *Stream[T]) Err() error {
	return s.err
}

func (s *Stream[T]) Close() error {
	if s.decoder == nil {
		// already closed
		return nil
	}
	return s.decoder.Close()
}