Atlas/internal/provider/stream.go

package provider

import (
	"bufio"
	"bytes"
	"context"
	"fmt"
	"io"
	"net/http"
	"strings"
	"time"

	"github.com/rose_cat707/Atlas/internal/service/reqlog"
)

type CallMeta struct {
	SessionID string
	Method    string
	Path      string
	Model     string
}

func (r *Router) ProxyWithMeta(ctx context.Context, cfg *Config, method, path string, body []byte, headers map[string]string, meta CallMeta) (*ProxyResult, error) {
	start := time.Now()
	res, err := r.proxy(ctx, cfg, method, path, body, headers)
	r.record(meta, cfg, start, res, err, body, "")
	return res, err
}

func (r *Router) ProxyStream(ctx context.Context, cfg *Config, method, path string, body []byte, headers map[string]string, meta CallMeta, w io.Writer, flush func()) error {
	start := time.Now()
	url := cfg.BaseURL + path
	req, err := http.NewRequestWithContext(ctx, method, url, bytes.NewReader(body))
	if err != nil {
		r.record(meta, cfg, start, nil, err, body, "")
		return err
	}
	req.Header.Set("Authorization", "Bearer "+cfg.APIKey)
	req.Header.Set("Content-Type", "application/json")
	req.Header.Set("Accept", "text/event-stream")
	for k, v := range headers {
		req.Header.Set(k, v)
	}

	resp, err := r.streamClient.Do(req)
	if err != nil {
		r.record(meta, cfg, start, nil, err, body, "")
		return err
	}
	defer resp.Body.Close()

	if resp.StatusCode >= 400 {
		data, _ := io.ReadAll(resp.Body)
		res := &ProxyResult{StatusCode: resp.StatusCode, Body: data, Headers: resp.Header}
		r.record(meta, cfg, start, res, fmt.Errorf("upstream %d", resp.StatusCode), body, "")
		return fmt.Errorf("upstream %d: %s", resp.StatusCode, string(data))
	}

	var capture strings.Builder
	reader := bufio.NewReader(resp.Body)
	for {
		line, readErr := reader.ReadString('\n')
		if len(line) > 0 {
			capture.WriteString(line)
			if _, err := io.WriteString(w, line); err != nil {
				return err
			}
			if flush != nil && (line == "\n" || strings.HasPrefix(line, "data:")) {
				flush()
			}
		}
		if readErr == io.EOF {
			break
		}
		if readErr != nil {
			r.record(meta, cfg, start, &ProxyResult{StatusCode: resp.StatusCode}, readErr, body, capture.String())
			return readErr
		}
	}
	r.record(meta, cfg, start, &ProxyResult{StatusCode: resp.StatusCode}, nil, body, capture.String())
	return nil
}

func (r *Router) record(meta CallMeta, cfg *Config, start time.Time, res *ProxyResult, callErr error, reqBody []byte, streamCapture string) {
	if r.logger == nil {
		return
	}
	entry := reqlog.Entry{
		SessionID:      meta.SessionID,
		Method:         meta.Method,
		Path:           meta.Path,
		Protocol:       cfg.Protocol,
		ProviderID:     reqlog.ProviderIDPtr(cfg.ID),
		Model:          meta.Model,
		LatencyMS:      reqlog.Since(start),
		RequestSummary: reqlog.Summarize(reqBody, 512),
	}
	if res != nil {
		entry.StatusCode = res.StatusCode
		if streamCapture != "" {
			entry.ResponseSummary = reqlog.SummarizeSSE(streamCapture, 512)
		} else {
			entry.ResponseSummary = reqlog.Summarize(res.Body, 512)
		}
	}
	if callErr != nil {
		entry.Error = callErr.Error()
		if entry.StatusCode == 0 {
			entry.StatusCode = 502
		}
	}
	r.logger.Log(context.Background(), entry)
}

func (r *Router) proxy(ctx context.Context, cfg *Config, method, path string, body []byte, headers map[string]string) (*ProxyResult, error) {
	url := cfg.BaseURL + path
	req, err := http.NewRequestWithContext(ctx, method, url, bytes.NewReader(body))
	if err != nil {
		return nil, err
	}
	req.Header.Set("Authorization", "Bearer "+cfg.APIKey)
	if len(body) > 0 {
		req.Header.Set("Content-Type", "application/json")
	}
	for k, v := range headers {
		req.Header.Set(k, v)
	}
	resp, err := r.client.Do(req)
	if err != nil {
		return nil, err
	}
	defer resp.Body.Close()
	data, err := io.ReadAll(resp.Body)
	if err != nil {
		return nil, err
	}
	return &ProxyResult{StatusCode: resp.StatusCode, Body: data, Headers: resp.Header}, nil
}

// Proxy keeps backward compatibility without logging.
func (r *Router) Proxy(ctx context.Context, cfg *Config, method, path string, body []byte, headers map[string]string) (*ProxyResult, error) {
	return r.proxy(ctx, cfg, method, path, body, headers)
}