iop/apps/edge/internal/openai/single_request_anthropic_stream.go

package openai

import (
	"context"
	"errors"
	"fmt"
	"net/http"
	"strings"
	"sync"
	"time"

	edgeservice "iop/apps/edge/internal/service"
)

const singleRequestAnthropicPingInterval = 15 * time.Second

var (
	errSingleRequestAnthropicStreamUnavailable = errors.New("single-request Anthropic stream is unavailable")
	errSingleRequestAnthropicUnknownProgress   = errors.New("single-request Anthropic stream received an unknown progress stage")
)

type singleRequestAnthropicTerminalKind uint8

const (
	singleRequestAnthropicTerminalFailure singleRequestAnthropicTerminalKind = iota + 1
	singleRequestAnthropicTerminalCancelled
)

// singleRequestAnthropicStream is a privacy-closed projection of the public
// single-request coordinator vocabulary. The mutex owns every byte written to
// the caller, including pings, block indices, and the exclusive terminal.
type singleRequestAnthropicStream struct {
	mu sync.Mutex

	w         http.ResponseWriter
	messageID string
	model     string

	started     bool
	terminal    bool
	terminalErr error
	nextBlock   int
	emitted     map[edgeservice.SingleRequestState]struct{}
}

func newSingleRequestAnthropicStream(
	w http.ResponseWriter,
	requestID string,
	publicModel string,
) (*singleRequestAnthropicStream, error) {
	if w == nil || strings.TrimSpace(requestID) == "" || strings.TrimSpace(publicModel) == "" {
		return nil, errSingleRequestAnthropicStreamUnavailable
	}
	_, ok := w.(http.Flusher)
	if !ok {
		return nil, fmt.Errorf("%w: response writer does not support flushing", errSingleRequestAnthropicStreamUnavailable)
	}
	return &singleRequestAnthropicStream{
		w:         w,
		messageID: "msg_iop_" + strings.TrimPrefix(requestID, "req_"),
		model:     publicModel,
		emitted:   make(map[edgeservice.SingleRequestState]struct{}, 4),
	}, nil
}

func (s *singleRequestAnthropicStream) Start() error {
	s.mu.Lock()
	defer s.mu.Unlock()
	if s.terminal {
		return s.terminalErr
	}
	return s.startLocked()
}

func (s *singleRequestAnthropicStream) startLocked() error {
	if s.started {
		return nil
	}
	s.w.Header().Set("Content-Type", "text/event-stream")
	s.w.Header().Set("Cache-Control", "no-cache")
	s.w.WriteHeader(http.StatusOK)
	message := map[string]any{
		"id":            s.messageID,
		"type":          "message",
		"role":          "assistant",
		"model":         s.model,
		"content":       []any{},
		"stop_reason":   nil,
		"stop_sequence": nil,
		"usage":         anthropicUsage{},
	}
	if err := s.writeEventLocked("message_start", map[string]any{
		"type": "message_start", "message": message,
	}); err != nil {
		s.failWireLocked(err)
		return err
	}
	s.started = true
	return nil
}

// Progress accepts only the coordinator's closed public stage enum. Arbitrary
// progress.Message, Result, and Err values are deliberately ignored.
func (s *singleRequestAnthropicStream) Progress(progress edgeservice.SingleRequestProgress) error {
	s.mu.Lock()
	defer s.mu.Unlock()
	if s.terminal {
		return s.terminalErr
	}
	summary, visible, known := singleRequestAnthropicProgressSummary(progress.Stage)
	if !known {
		return fmt.Errorf("%w: %q", errSingleRequestAnthropicUnknownProgress, progress.Stage)
	}
	if !visible {
		return nil
	}
	if _, ok := s.emitted[progress.Stage]; ok {
		return nil
	}
	if err := s.startLocked(); err != nil {
		return err
	}
	if err := s.writeTextBlockLocked(summary); err != nil {
		s.failWireLocked(err)
		return err
	}
	s.emitted[progress.Stage] = struct{}{}
	return nil
}

func singleRequestAnthropicProgressSummary(stage edgeservice.SingleRequestState) (string, bool, bool) {
	switch stage {
	case edgeservice.SingleRequestStatePlanning:
		return "Planning the requested work.", true, true
	case edgeservice.SingleRequestStateWorking:
		return "Executing the requested work.", true, true
	case edgeservice.SingleRequestStateReviewing:
		return "Reviewing the completed work.", true, true
	case edgeservice.SingleRequestStateRepairing:
		return "Repairing issues found during review.", true, true
	case edgeservice.SingleRequestStateAccepted,
		edgeservice.SingleRequestStateInternalTool,
		edgeservice.SingleRequestStateFinalizing,
		edgeservice.SingleRequestStateCompleted,
		edgeservice.SingleRequestStateFailed,
		edgeservice.SingleRequestStateCancelled:
		return "", false, true
	default:
		return "", false, false
	}
}

func (s *singleRequestAnthropicStream) Ping() error {
	s.mu.Lock()
	defer s.mu.Unlock()
	if s.terminal {
		return s.terminalErr
	}
	if err := s.startLocked(); err != nil {
		return err
	}
	if err := s.writeEventLocked("ping", map[string]any{"type": "ping"}); err != nil {
		s.failWireLocked(err)
		return err
	}
	return nil
}

func (s *singleRequestAnthropicStream) Final(result edgeservice.SingleRequestResult) error {
	s.mu.Lock()
	defer s.mu.Unlock()
	if s.terminal {
		return s.terminalErr
	}
	if err := s.startLocked(); err != nil {
		return err
	}

	// Claim terminal ownership before the first terminal byte. A partial write
	// is never retried as either another success or an error terminal.
	s.terminal = true
	if err := s.writeTextBlockLocked(result.Output); err != nil {
		s.terminalErr = err
		return err
	}
	if err := s.writeEventLocked("message_delta", map[string]any{
		"type":  "message_delta",
		"delta": map[string]any{"stop_reason": "end_turn", "stop_sequence": nil},
		"usage": anthropicUsage{},
	}); err != nil {
		s.terminalErr = err
		return err
	}
	if err := s.writeEventLocked("message_stop", map[string]any{"type": "message_stop"}); err != nil {
		s.terminalErr = err
		return err
	}
	return nil
}

func (s *singleRequestAnthropicStream) Error(kind singleRequestAnthropicTerminalKind) error {
	s.mu.Lock()
	defer s.mu.Unlock()
	if s.terminal {
		return s.terminalErr
	}
	if err := s.startLocked(); err != nil {
		return err
	}
	errorType, message := singleRequestAnthropicError(kind)
	s.terminal = true
	if err := s.writeEventLocked("error", anthropicErrorResponse{
		Type: "error", Error: errorBody{Type: errorType, Message: message},
	}); err != nil {
		s.terminalErr = err
		return err
	}
	return nil
}

func singleRequestAnthropicError(kind singleRequestAnthropicTerminalKind) (string, string) {
	switch kind {
	case singleRequestAnthropicTerminalCancelled:
		return "api_error", "single-request execution was cancelled"
	default:
		return "api_error", "single-request execution failed"
	}
}

func (s *singleRequestAnthropicStream) writeTextBlockLocked(text string) error {
	index := s.nextBlock
	if err := s.writeEventLocked("content_block_start", map[string]any{
		"type": "content_block_start", "index": index,
		"content_block": map[string]any{"type": "text", "text": ""},
	}); err != nil {
		return err
	}
	if err := s.writeEventLocked("content_block_delta", map[string]any{
		"type": "content_block_delta", "index": index,
		"delta": map[string]any{"type": "text_delta", "text": text},
	}); err != nil {
		return err
	}
	if err := s.writeEventLocked("content_block_stop", map[string]any{
		"type": "content_block_stop", "index": index,
	}); err != nil {
		return err
	}
	s.nextBlock++
	return nil
}

func (s *singleRequestAnthropicStream) writeEventLocked(event string, value any) error {
	if err := writeAnthropicSSEEvent(s.w, event, value); err != nil {
		return err
	}
	return http.NewResponseController(s.w).Flush()
}

func (s *singleRequestAnthropicStream) failWireLocked(err error) {
	if s.terminal {
		if s.terminalErr == nil {
			s.terminalErr = err
		}
		return
	}
	s.terminal = true
	s.terminalErr = err
}

type singleRequestAnthropicTicker interface {
	Ticks() <-chan time.Time
	Stop()
}

type wallClockSingleRequestAnthropicTicker struct {
	ticker *time.Ticker
}

func (t *wallClockSingleRequestAnthropicTicker) Ticks() <-chan time.Time { return t.ticker.C }
func (t *wallClockSingleRequestAnthropicTicker) Stop()                   { t.ticker.Stop() }

type singleRequestAnthropicTickerFactory func() singleRequestAnthropicTicker

func newWallClockSingleRequestAnthropicTicker() singleRequestAnthropicTicker {
	return &wallClockSingleRequestAnthropicTicker{ticker: time.NewTicker(singleRequestAnthropicPingInterval)}
}

// pumpSingleRequestAnthropicStream owns coordinator progress and the liveness
// worker for one HTTP request. The worker is stopped and joined before every
// terminal attempt or return, so no ping can race after the terminal.
func pumpSingleRequestAnthropicStream(
	ctx context.Context,
	execution edgeservice.SingleRequestExecution,
	stream *singleRequestAnthropicStream,
	tickerFactory singleRequestAnthropicTickerFactory,
) error {
	if execution == nil || stream == nil || tickerFactory == nil {
		return errSingleRequestAnthropicStreamUnavailable
	}
	if err := ctx.Err(); err != nil {
		execution.Cancel()
		return err
	}
	if err := stream.Start(); err != nil {
		execution.Cancel()
		return err
	}
	ticker := tickerFactory()
	if ticker == nil {
		execution.Cancel()
		return errSingleRequestAnthropicStreamUnavailable
	}

	stopPing := make(chan struct{})
	pingDone := make(chan struct{})
	pingErr := make(chan error, 1)
	go func() {
		defer close(pingDone)
		for {
			select {
			case <-stopPing:
				return
			case _, ok := <-ticker.Ticks():
				if !ok {
					return
				}
				if err := stream.Ping(); err != nil {
					select {
					case pingErr <- err:
					default:
					}
					return
				}
			}
		}
	}()

	var stopOnce sync.Once
	stopAndJoinPing := func() {
		stopOnce.Do(func() {
			ticker.Stop()
			close(stopPing)
			<-pingDone
		})
	}
	defer stopAndJoinPing()

	for {
		select {
		case <-ctx.Done():
			stopAndJoinPing()
			execution.Cancel()
			return ctx.Err()
		case err := <-pingErr:
			stopAndJoinPing()
			execution.Cancel()
			return err
		case progress, ok := <-execution.Progress():
			if !ok {
				stopAndJoinPing()
				if ctx.Err() != nil {
					return ctx.Err()
				}
				if execution.State() == edgeservice.SingleRequestStateCompleted {
					return nil
				}
				return stream.Error(singleRequestAnthropicTerminalFailure)
			}

			switch progress.Stage {
			case edgeservice.SingleRequestStateFinalizing:
				stopAndJoinPing()
				if ctx.Err() != nil {
					execution.Cancel()
					return ctx.Err()
				}
				if progress.Result == nil {
					writeErr := stream.Error(singleRequestAnthropicTerminalFailure)
					ackErr := execution.AcknowledgeTerminal(false)
					return errors.Join(writeErr, ackErr)
				}
				writeErr := stream.Final(*progress.Result)
				ackErr := execution.AcknowledgeTerminal(writeErr == nil)
				return errors.Join(writeErr, ackErr)
			case edgeservice.SingleRequestStateFailed:
				stopAndJoinPing()
				if ctx.Err() != nil {
					return ctx.Err()
				}
				return stream.Error(singleRequestAnthropicTerminalFailure)
			case edgeservice.SingleRequestStateCancelled:
				stopAndJoinPing()
				if ctx.Err() != nil {
					return ctx.Err()
				}
				return stream.Error(singleRequestAnthropicTerminalCancelled)
			default:
				if err := stream.Progress(progress); err != nil {
					stopAndJoinPing()
					_ = stream.Error(singleRequestAnthropicTerminalFailure)
					execution.Cancel()
					return err
				}
			}
		}
	}
}