CrossRelay/internal/gateway/messages.go

package gateway

import (
	"errors"
	"fmt"
	"io"
	"net/http"

	"gitea.rutua.cn/kate522/CrossRelay/internal/anthropic"
	"gitea.rutua.cn/kate522/CrossRelay/internal/chatcmpl"
)

// handleMessages serves POST /v1/messages in both streaming and
// non-streaming modes.
func (s *Server) handleMessages(w http.ResponseWriter, r *http.Request) {
	var req anthropic.MessagesRequest
	if err := decodeJSON(r, &req, 64<<20); err != nil {
		writeAnthropicError(w, http.StatusBadRequest, "invalid_request_error", "invalid request body: "+err.Error())
		return
	}
	provider := s.store.ActiveClaude()
	if provider == nil {
		writeAnthropicError(w, http.StatusBadRequest, "invalid_request_error",
			"no active provider: activate one in the CrossRelay app")
		return
	}
	model := provider.ResolveModel(req.Model)

	chatReq, err := anthropic.ConvertRequest(&req, model)
	if err != nil {
		writeAnthropicError(w, http.StatusBadRequest, "invalid_request_error", err.Error())
		return
	}

	rec := recordFrom(r)
	rec.update(func(e *LogEntry) {
		e.Provider = provider.Name
		e.Model = model
		e.InputTokens = anthropic.EstimateTokens(&req)
	})

	client := chatcmpl.NewClient(provider.BaseURL, provider.APIKey)
	if req.Stream {
		s.handleStream(w, r, client, chatReq, &req, model, rec)
		return
	}

	resp, err := client.Create(r.Context(), chatReq)
	if err != nil {
		s.failUpstream(w, rec, err)
		return
	}
	out := anthropic.ConvertResponse(resp, req.Model)
	rec.update(func(e *LogEntry) {
		e.OutputTokens = out.Usage.OutputTokens
	})
	writeJSON(w, http.StatusOK, out)
}

// handleStream relays a Chat Completions SSE stream as Anthropic events.
func (s *Server) handleStream(w http.ResponseWriter, r *http.Request, client *chatcmpl.Client,
	chatReq *chatcmpl.Request, req *anthropic.MessagesRequest, model string, rec *record) {

	resp, err := client.CreateStream(r.Context(), chatReq)
	if err != nil {
		s.failUpstream(w, rec, err)
		return
	}
	defer resp.Body.Close()

	conv := anthropic.NewStreamConverter(w, req.Model, anthropic.EstimateTokens(req))
	if err := conv.WriteStart(); err != nil {
		return
	}
	stream := chatcmpl.NewStream(resp.Body)
	for {
		chunk, err := stream.Next()
		if errors.Is(err, io.EOF) {
			break
		}
		if err != nil {
			rec.update(func(e *LogEntry) { e.Error = err.Error() })
			conv.WriteError("api_error", err.Error())
			return
		}
		if err := conv.OnChunk(chunk); err != nil {
			rec.update(func(e *LogEntry) { e.Error = err.Error() })
			conv.WriteError("api_error", err.Error())
			return
		}
	}
	if err := conv.Finish(); err != nil {
		rec.update(func(e *LogEntry) { e.Error = err.Error() })
	}
}

// handleCountTokens answers POST /v1/messages/count_tokens with a
// heuristic estimate so Claude Code's context meter keeps working.
func (s *Server) handleCountTokens(w http.ResponseWriter, r *http.Request) {
	var req anthropic.MessagesRequest
	if err := decodeJSON(r, &req, 64<<20); err != nil {
		writeAnthropicError(w, http.StatusBadRequest, "invalid_request_error", "invalid request body: "+err.Error())
		return
	}
	writeJSON(w, http.StatusOK, anthropic.CountTokensResponse{
		InputTokens: anthropic.EstimateTokens(&req),
	})
}

// failUpstream maps an upstream failure onto an Anthropic error.
func (s *Server) failUpstream(w http.ResponseWriter, rec *record, err error) {
	rec.update(func(e *LogEntry) { e.Error = err.Error() })
	var ue *chatcmpl.UpstreamError
	if errors.As(err, &ue) {
		switch {
		case ue.StatusCode == http.StatusUnauthorized || ue.StatusCode == http.StatusForbidden:
			writeAnthropicError(w, ue.StatusCode, "authentication_error", ue.Message)
		case ue.StatusCode == http.StatusTooManyRequests:
			writeAnthropicError(w, ue.StatusCode, "rate_limit_error", ue.Message)
		case ue.StatusCode == http.StatusBadRequest:
			writeAnthropicError(w, http.StatusBadRequest, "invalid_request_error", ue.Message)
		default:
			writeAnthropicError(w, http.StatusBadGateway, "api_error", ue.Message)
		}
		return
	}
	writeAnthropicError(w, http.StatusBadGateway, "api_error", fmt.Sprintf("upstream request failed: %v", err))
}