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

package openai

import (
	"bytes"
	"encoding/json"
	"errors"
	"fmt"
	"io"
	"net/http"
	"sync"
	"sync/atomic"

	"iop/packages/go/streamgate"
)

const (
	openAIIngressTypedViewName = "openai.request.semantic"
	openAIRebuiltBodyViewName  = "openai.request.rebuilt"
)

var (
	errOpenAIIngressTooLarge = errors.New("openai request body exceeds ingress limit")
	openAIIngressSequence    atomic.Uint64
)

// openAIIngressSnapshot is the request-local owner of the caller's canonical
// JSON body and its bounded semantic ledger. It deliberately contains no HTTP
// headers, provider credentials, logger fields, or cross-request references.
type openAIIngressSnapshot struct {
	mu          sync.RWMutex
	ref         string
	snapshot    *streamgate.IngressSnapshot
	retainedMax int64
	closed      bool
}

// readOpenAIIngressBody installs the HTTP body limit before reading. The
// standard library reader performs a limit+1 probe internally, so an exact
// limit body succeeds and the first excess byte is reported as overflow.
func readOpenAIIngressBody(w http.ResponseWriter, r *http.Request, maxBytes int64) ([]byte, error) {
	if maxBytes <= 0 {
		return nil, streamgate.ErrIngressSnapshotZeroMaxBytes
	}
	r.Body = http.MaxBytesReader(w, r.Body, maxBytes)
	body, err := io.ReadAll(r.Body)
	if err != nil {
		var maxErr *http.MaxBytesError
		if errors.As(err, &maxErr) {
			return nil, errOpenAIIngressTooLarge
		}
		return nil, err
	}
	return body, nil
}

// buildOpenAIIngressSnapshot transfers the bounded canonical request and a
// serialized semantic view into the shared Core ledger. Equal payloads share
// one backing handle and are counted once.
func buildOpenAIIngressSnapshot(maxBytes int64, canonical []byte, semanticView any) (*openAIIngressSnapshot, error) {
	var semantic []byte
	if raw, ok := semanticView.(json.RawMessage); ok {
		if !json.Valid(raw) {
			return nil, fmt.Errorf("encode OpenAI ingress semantic view: invalid JSON")
		}
		semantic = raw
	} else {
		var err error
		semantic, err = json.Marshal(semanticView)
		if err != nil {
			return nil, fmt.Errorf("encode OpenAI ingress semantic view: %w", err)
		}
	}

	builder, err := streamgate.NewIngressSnapshotBuilder(maxBytes)
	if err != nil {
		return nil, err
	}
	defer builder.Close()

	canonicalHandle, err := builder.IssueOwnedBackingHandle(canonical)
	if err != nil {
		return nil, err
	}
	if err := builder.SetCanonicalWithHandle(canonicalHandle); err != nil {
		return nil, err
	}

	semanticHandle := canonicalHandle
	if !bytes.Equal(semantic, canonical) {
		semanticHandle, err = builder.IssueOwnedBackingHandle(semantic)
		if err != nil {
			return nil, err
		}
	}
	if err := builder.AddTypedView(openAIIngressTypedViewName, semanticHandle); err != nil {
		return nil, err
	}

	snapshot, err := builder.Build()
	if err != nil {
		return nil, err
	}
	ref := fmt.Sprintf("openai.ingress.%d", openAIIngressSequence.Add(1))
	return &openAIIngressSnapshot{
		ref:         ref,
		snapshot:    snapshot,
		retainedMax: snapshot.Accessor().RetainedBytes(),
	}, nil
}

func (s *openAIIngressSnapshot) canonicalBody() ([]byte, error) {
	if s == nil {
		return nil, streamgate.ErrIngressSnapshotClosed
	}
	s.mu.RLock()
	defer s.mu.RUnlock()
	if s.closed || s.snapshot == nil {
		return nil, streamgate.ErrIngressSnapshotClosed
	}
	return s.snapshot.Accessor().CanonicalAlias()
}

func (s *openAIIngressSnapshot) semanticBody() ([]byte, error) {
	if s == nil {
		return nil, streamgate.ErrIngressSnapshotClosed
	}
	s.mu.RLock()
	defer s.mu.RUnlock()
	if s.closed || s.snapshot == nil {
		return nil, streamgate.ErrIngressSnapshotClosed
	}
	return s.snapshot.Accessor().TypedViewAlias(openAIIngressTypedViewName)
}

func (s *openAIIngressSnapshot) recoveryRef() (streamgate.RecoveryRequestSnapshotRef, error) {
	if s == nil {
		return streamgate.RecoveryRequestSnapshotRef{}, streamgate.ErrIngressSnapshotClosed
	}
	s.mu.RLock()
	defer s.mu.RUnlock()
	if s.closed || s.snapshot == nil {
		return streamgate.RecoveryRequestSnapshotRef{}, streamgate.ErrIngressSnapshotClosed
	}
	accessor := s.snapshot.Accessor()
	retained := accessor.RetainedBytes()
	peak := retained
	if s.retainedMax > peak {
		peak = s.retainedMax
	}
	return streamgate.NewRecoveryRequestSnapshotRef(
		s.ref,
		uint64(retained),
		uint64(peak),
		uint64(accessor.MaxBytes()),
	)
}

func (s *openAIIngressSnapshot) reserveRebuild(expectedBytes int64) (*streamgate.IngressSnapshotRebuildGuard, error) {
	if s == nil {
		return nil, streamgate.ErrIngressSnapshotClosed
	}
	s.mu.Lock()
	defer s.mu.Unlock()
	if s.closed || s.snapshot == nil {
		return nil, streamgate.ErrIngressSnapshotClosed
	}
	guard, err := s.snapshot.ReserveRebuild(expectedBytes)
	if err != nil {
		return nil, err
	}
	peak := s.snapshot.Accessor().RetainedBytes() + s.snapshot.ReservedTempBytes()
	if peak > s.retainedMax {
		s.retainedMax = peak
	}
	return guard, nil
}

func (s *openAIIngressSnapshot) accessor() (streamgate.IngressSnapshotAccessor, error) {
	if s == nil {
		return streamgate.IngressSnapshotAccessor{}, streamgate.ErrIngressSnapshotClosed
	}
	s.mu.RLock()
	defer s.mu.RUnlock()
	if s.closed || s.snapshot == nil {
		return streamgate.IngressSnapshotAccessor{}, streamgate.ErrIngressSnapshotClosed
	}
	return s.snapshot.Accessor(), nil
}

func (s *openAIIngressSnapshot) Close() {
	if s == nil {
		return
	}
	s.mu.Lock()
	defer s.mu.Unlock()
	if s.closed {
		return
	}
	s.closed = true
	if s.snapshot != nil {
		s.snapshot.Close()
		s.snapshot = nil
	}
}

func (s *openAIIngressSnapshot) isClosed() bool {
	if s == nil {
		return true
	}
	s.mu.RLock()
	defer s.mu.RUnlock()
	return s.closed || s.snapshot == nil
}

func (s *Server) maxIngressSnapshotBytes() int64 {
	s.mu.RLock()
	defer s.mu.RUnlock()
	return int64(s.cfg.StreamEvidenceGate.EffectiveMaxIngressSnapshotBytes())
}

func writeOpenAIIngressError(w http.ResponseWriter, err error) {
	if errors.Is(err, errOpenAIIngressTooLarge) ||
		errors.Is(err, streamgate.ErrIngressSnapshotLimitExceeded) ||
		errors.Is(err, streamgate.ErrIngressSnapshotRebuildOverflow) ||
		errors.Is(err, streamgate.ErrRecoverySnapshotLimitExceeded) {
		writeError(w, http.StatusRequestEntityTooLarge, "invalid_request_error", "request body exceeds configured limit")
		return
	}
	writeError(w, http.StatusBadRequest, "invalid_request_error", "failed to read request body")
}