stream.go

package cursor

import (
	"bufio"
	"bytes"
	"context"
	"encoding/json"
	"errors"
	"io"
	"net/http"
	"net/url"
	"strconv"
	"strings"
)

// StreamOption customizes a run stream request.
type StreamOption func(*streamOptions)

type streamOptions struct {
	lastEventID string
}

// WithLastEventID resumes a stream from the last received SSE id.
func WithLastEventID(id string) StreamOption {
	return func(options *streamOptions) {
		options.lastEventID = id
	}
}

// SDKMessage is the normalized streaming event envelope.
//
// The public Cloud Agents API is beta and some payloads are intentionally
// unstable. Raw always preserves the original JSON data for defensive parsing.
type SDKMessage struct {
	Type       string          `json:"type"`
	Event      string          `json:"-"`
	ID         string          `json:"-"`
	AgentID    string          `json:"agent_id,omitempty"`
	RunID      string          `json:"run_id,omitempty"`
	Text       string          `json:"text,omitempty"`
	Status     RunStatus       `json:"status,omitempty"`
	Message    *SDKChatMessage `json:"message,omitempty"`
	StatusText string          `json:"message_text,omitempty"`
	CallID     string          `json:"call_id,omitempty"`
	Name       string          `json:"name,omitempty"`
	Args       json.RawMessage `json:"args,omitempty"`
	Result     json.RawMessage `json:"result,omitempty"`
	Truncated  json.RawMessage `json:"truncated,omitempty"`
	RequestID  string          `json:"request_id,omitempty"`
	Code       string          `json:"code,omitempty"`
	Raw        json.RawMessage `json:"-"`
}

// SDKChatMessage is a simplified TypeScript-style chat message.
type SDKChatMessage struct {
	Role    string            `json:"role"`
	Content []SDKMessageBlock `json:"content,omitempty"`
}

// SDKMessageBlock is a text or tool-use content block.
type SDKMessageBlock struct {
	Type  string          `json:"type"`
	Text  string          `json:"text,omitempty"`
	ID    string          `json:"id,omitempty"`
	Name  string          `json:"name,omitempty"`
	Input json.RawMessage `json:"input,omitempty"`
}

// RunStream reads Server-Sent Events for one cloud run.
type RunStream struct {
	RetentionSeconds int

	resp   *http.Response
	reader *bufio.Reader
	done   bool
}

// StreamRun opens the Server-Sent Events stream for a run.
func (c *Client) StreamRun(ctx context.Context, agentID, runID string, options ...StreamOption) (*RunStream, error) {
	streamOpts := streamOptions{}
	for _, opt := range options {
		if opt != nil {
			opt(&streamOpts)
		}
	}

	req, err := c.newRequest(ctx, http.MethodGet, "/v1/agents/"+url.PathEscape(agentID)+"/runs/"+url.PathEscape(runID)+"/stream", nil, nil, authBasic)
	if err != nil {
		return nil, err
	}
	req.Header.Set("Accept", "text/event-stream")
	if streamOpts.lastEventID != "" {
		req.Header.Set("Last-Event-ID", streamOpts.lastEventID)
	}

	resp, err := c.httpClient.Do(req)
	if err != nil {
		return nil, err
	}
	if resp.StatusCode < 200 || resp.StatusCode >= 300 {
		defer resp.Body.Close()
		raw, readErr := io.ReadAll(resp.Body)
		if readErr != nil {
			return nil, readErr
		}
		return nil, decodeAPIError(resp.StatusCode, resp.Header, raw)
	}

	retention, _ := strconv.Atoi(resp.Header.Get("X-Cursor-Stream-Retention-Seconds"))
	return &RunStream{
		RetentionSeconds: retention,
		resp:             resp,
		reader:           bufio.NewReader(resp.Body),
	}, nil
}

// Next reads the next stream event. It returns io.EOF when the stream closes.
func (s *RunStream) Next() (*SDKMessage, error) {
	if s == nil || s.reader == nil {
		return nil, errors.New("stream is closed")
	}
	if s.done {
		return nil, io.EOF
	}

	for {
		id, event, data, err := s.readEvent()
		if err != nil {
			return nil, err
		}
		if event == "" && len(bytes.TrimSpace(data)) == 0 && id == "" {
			continue
		}
		message := decodeSDKMessage(id, event, data)
		if message.Type == "done" {
			s.done = true
		}
		return message, nil
	}
}

// Close closes the underlying HTTP response body.
func (s *RunStream) Close() error {
	if s == nil || s.resp == nil || s.resp.Body == nil {
		return nil
	}
	err := s.resp.Body.Close()
	s.reader = nil
	return err
}

func (s *RunStream) readEvent() (string, string, []byte, error) {
	var id string
	var event string
	var data []string
	sawField := false

	for {
		line, err := s.reader.ReadString('\n')
		if err != nil && !errors.Is(err, io.EOF) {
			return "", "", nil, err
		}
		line = strings.TrimSuffix(line, "\n")
		line = strings.TrimSuffix(line, "\r")

		if line == "" {
			if !sawField && err == nil {
				continue
			}
			return id, event, []byte(strings.Join(data, "\n")), nil
		}

		if strings.HasPrefix(line, ":") {
			sawField = true
		} else {
			field, value, found := strings.Cut(line, ":")
			if found && strings.HasPrefix(value, " ") {
				value = strings.TrimPrefix(value, " ")
			}
			if !found {
				value = ""
			}
			sawField = true
			switch field {
			case "id":
				id = value
			case "event":
				event = value
			case "data":
				data = append(data, value)
			}
		}

		if errors.Is(err, io.EOF) {
			if !sawField {
				return "", "", nil, io.EOF
			}
			return id, event, []byte(strings.Join(data, "\n")), nil
		}
	}
}

func decodeSDKMessage(id, event string, data []byte) *SDKMessage {
	if event == "" {
		event = "message"
	}

	message := &SDKMessage{
		Type:  event,
		Event: event,
		ID:    id,
		Raw:   append(json.RawMessage(nil), data...),
	}

	if len(bytes.TrimSpace(data)) == 0 {
		return message
	}

	var payload struct {
		Type               string          `json:"type"`
		AgentID            string          `json:"agent_id"`
		RunIDSnake         string          `json:"run_id"`
		RunIDCamel         string          `json:"runId"`
		Text               string          `json:"text"`
		Status             RunStatus       `json:"status"`
		Message            json.RawMessage `json:"message"`
		ThinkingDurationMS int64           `json:"thinking_duration_ms"`
		CallID             string          `json:"call_id"`
		Name               string          `json:"name"`
		Args               json.RawMessage `json:"args"`
		Result             json.RawMessage `json:"result"`
		Truncated          json.RawMessage `json:"truncated"`
		RequestID          string          `json:"request_id"`
		Code               string          `json:"code"`
	}
	if err := json.Unmarshal(data, &payload); err != nil {
		return message
	}

	if payload.Type != "" {
		message.Type = payload.Type
	}
	message.AgentID = payload.AgentID
	if payload.RunIDSnake != "" {
		message.RunID = payload.RunIDSnake
	} else {
		message.RunID = payload.RunIDCamel
	}
	message.Text = payload.Text
	message.Status = payload.Status
	if len(payload.Message) > 0 {
		var chatMessage SDKChatMessage
		if err := json.Unmarshal(payload.Message, &chatMessage); err == nil && chatMessage.Role != "" {
			message.Message = &chatMessage
		} else {
			var statusText string
			if err := json.Unmarshal(payload.Message, &statusText); err == nil {
				message.StatusText = statusText
			}
		}
	}
	message.CallID = payload.CallID
	message.Name = payload.Name
	message.Args = payload.Args
	message.Result = payload.Result
	message.Truncated = payload.Truncated
	message.RequestID = payload.RequestID
	message.Code = payload.Code

	return message
}