ModelRouter/internal/gateway/messages.go

package gateway

import (
	"context"
	"encoding/json"
	"errors"
	"fmt"
	"net/http"
	"strings"
	"time"

	"llmsproxy/internal/provider"
	"llmsproxy/internal/scheduler"
	"llmsproxy/internal/types"
)

// handleMessages serves the Anthropic Messages protocol (POST /v1/messages).
//
// The point of this endpoint is NOT to add a second scheduling path. It reuses
// everything /v1/chat/completions does — key auth, model scope, per-model
// quotas, the AUTO chain, cooldowns, audit records, plugin hooks — and differs
// in exactly two places:
//
//  1. the request body is parsed from Anthropic's shape into the unified format
//     (messages with content blocks, a top-level `system`, `max_tokens`
//     mandatory, `thinking` optional);
//  2. the response is rendered by the SERVING adapter's serialize_* hooks, so
//     the client sees Anthropic's shape back.
//
// The kernel therefore contains no Anthropic field names on this path: parsing
// is table-driven, and serialization belongs to the adapter.
//
// A source whose adapter does not declare serialize_* falls back to the
// OpenAI shape, which is what the kernel guarantees for an unconfigured
// deployment. That asymmetry is deliberate (see serializer.go): an Anthropic
// client talking to a source that cannot speak Anthropic gets a clear
// OpenAI-shaped answer rather than a malformed Anthropic one.
func (g *Gateway) handleMessages(w http.ResponseWriter, r *http.Request) {
	if r.Method != http.MethodPost {
		writeError(w, http.StatusMethodNotAllowed, "method_not_allowed", "use POST")
		return
	}
	var raw anthropicRequest
	if err := json.NewDecoder(r.Body).Decode(&raw); err != nil {
		writeError(w, http.StatusBadRequest, "invalid_request", "invalid json: "+err.Error())
		return
	}
	req, err := raw.toUnified(g.core.DefaultModel())
	if err != nil {
		writeError(w, http.StatusBadRequest, "invalid_request", err.Error())
		return
	}

	model := req.Model
	ctx := r.Context()
	g.fireStart(ctx, &chatRequest{Messages: req.Messages, Tools: nil},
		chatKind(req.Stream), model, len(req.Messages), 0)

	if isAuto(model) {
		chain, _ := g.core.AutoChainFor(reqKey(ctx))
		if chain == nil || len(chain.Tiers) == 0 {
			writeError(w, http.StatusServiceUnavailable, "no_provider", "no auto slot configured")
			return
		}
		if q := g.checkQuota(ctx, "AUTO"); q != nil {
			g.writeReject(w, q)
			return
		}
		g.serveUnified(ctx, w, chain, req, nil)
		return
	}
	if allow := g.allowedModels(ctx); allow != nil && !g.hasScopeModel(allow, model) {
		writeError(w, http.StatusForbidden, "model_not_allowed",
			"model \""+model+"\" is not allowed for this key")
		return
	}
	cands, effective := g.resolveCands(ctx, &chatRequest{Model: model})
	if len(cands) == 0 {
		writeError(w, http.StatusNotFound, "model_not_found",
			"model \""+model+"\" is not configured")
		return
	}
	if effective == "" {
		effective = firstModel(cands[0])
	}
	if q := g.checkQuota(ctx, effective); q != nil {
		g.writeReject(w, q)
		return
	}
	g.serveUnified(ctx, w, nil, req, &directRoute{cands: cands, effective: effective})
}

// directRoute carries the pre-resolved candidate list for a model-pinned
// request, so serveUnified can run both routes (direct / AUTO) from one body.
type directRoute struct {
	cands     []*provider.Provider
	effective string
}

// serveUnified runs the request through the scheduler and writes the reply via
// the serving adapter's serializer. It is the shared tail of handleChat and
// handleMessages: by this point every gate (auth, scope, quota) has already run,
// so the only remaining questions are "which source serves it" and "what shape
// does the client get".
func (g *Gateway) serveUnified(ctx context.Context, w http.ResponseWriter, chain *scheduler.Chain, req *types.ChatRequest, direct *directRoute) {
	rec := &Req{Key: keyID(reqKey(ctx)), Type: "chat", OK: false}
	if direct != nil {
		rec.Model = direct.effective
		rec.Source = firstSource(direct.cands)
	} else {
		rec.Type = "stream"
	}
	if req.Stream {
		rec.Type = "stream"
	}

	done := g.stats.Begin()
	defer done()

	var (
		resp               *types.UnifiedResponse
		chunks             <-chan types.UnifiedChunk
		usedSrc, usedModel string
		err                error
	)
	t0 := time.Now()
	if direct != nil {
		cands := toScheduler(direct.cands)
		if req.Stream {
			chunks, usedSrc, usedModel, err = g.core.Scheduler().ChatStream(ctx, cands, req)
		} else {
			resp, usedSrc, usedModel, err = g.core.Scheduler().Chat(ctx, cands, req)
		}
	} else {
		if req.Stream {
			chunks, usedSrc, usedModel, err = g.core.Scheduler().ChainChatStream(ctx, chain, req, nil, nil)
		} else {
			resp, usedSrc, usedModel, err = g.core.Scheduler().ChainChat(ctx, chain, req, nil, nil)
		}
	}
	rec.LatMs = time.Since(t0).Milliseconds()
	if err != nil {
		g.failChat(w, rec, err)
		g.writeRec(rec)
		return
	}
	rec.OK = true
	rec.Status = http.StatusOK
	if usedModel != "" {
		rec.Model = usedModel
	}
	rec.Source = usedSrc
	ser := g.serializerFor(usedSrc, directProviderList(direct))

	if chunks != nil {
		g.fireRouted(ctx, "stream", usedSrc, usedModel, tierForRoute(direct), true)
		rec.Prompt = estimatePromptTokens(req)
		g.pumpStream(w, ser, rec, chunks, modelLabel(req, usedModel), t0)
		return
	}
	recordChatUsage(rec, req, resp)
	g.fireRouted(ctx, "chat", usedSrc, usedModel, tierForRoute(direct), false)
	rec.FirstByteMs = rec.LatMs
	g.writeRec(rec)
	writeChatCompletion(w, ser, resp, modelLabel(req, usedModel))
}

// ---- Anthropic request shape (in-bound only; the kernel holds no outbound
// Anthropic field names) ----

// anthropicRequest is the client-facing parse target for POST /v1/messages.
//
// Only the fields the gateway must understand to schedule and account for a
// request appear here. Anything else the client sent is preserved verbatim in
// raw so an adapter can read it (the same "kernel does not interpret" rule that
// governs tools).
type anthropicRequest struct {
	Model     string             `json:"model"`
	System    json.RawMessage    `json:"system,omitempty"`
	Messages  []anthropicMessage `json:"messages"`
	MaxTokens int                `json:"max_tokens"`
	Stream    bool               `json:"stream,omitempty"`

	Temperature *float64               `json:"temperature,omitempty"`
	TopP        *float64               `json:"top_p,omitempty"`
	TopK        int                    `json:"top_k,omitempty"`
	StopSeqs    []string               `json:"stop_sequences,omitempty"`
	Tools       json.RawMessage        `json:"tools,omitempty"`
	ToolChoice  json.RawMessage        `json:"tool_choice,omitempty"`
	Thinking    *types.ThinkingConfig  `json:"thinking,omitempty"`
	Metadata    map[string]interface{} `json:"metadata,omitempty"`

	// raw is the untouched body, so an adapter can read fields the gateway
	// has no reason to know (cache_control, service_tier, …).
	raw json.RawMessage
}

// anthropicMessage accepts Anthropic's content as either a plain string or an
// array of typed blocks; both normalize to the unified message shape.
type anthropicMessage struct {
	Role    string          `json:"role"`
	Content json.RawMessage `json:"content"`
}

// toUnified converts the parsed request into the internal format.
func (a *anthropicRequest) toUnified(defaultModel string) (*types.ChatRequest, error) {
	if len(a.Messages) == 0 {
		return nil, errors.New("messages is required")
	}
	msgs := make([]types.ChatMessage, 0, len(a.Messages))
	for i, m := range a.Messages {
		u := types.ChatMessage{Role: m.Role, Content: normalizeAnthropicContent(m.Content)}
		if u.Role == "" {
			return nil, fmt.Errorf("messages[%d].role is required", i)
		}
		msgs = append(msgs, u)
	}
	model := a.Model
	if model == "" {
		model = defaultModel
	}
	th := a.Thinking
	if th == nil {
		// Anthropic's default is thinking OFF, and `thinking: {enabled:false}`
		// is the explicit spelling of that. Emitting it keeps the unified
		// request honest about the client's intent instead of leaving the
		// field absent and letting each adapter guess.
		no := false
		th = &types.ThinkingConfig{Enabled: &no}
	}
	return &types.ChatRequest{
		Model:         model,
		Messages:      msgs,
		Temperature:   a.Temperature,
		MaxTokens:     a.MaxTokens,
		Stream:        a.Stream,
		TopP:          a.TopP,
		TopK:          a.TopK,
		StopSequences: a.StopSeqs,
		Tools:         a.Tools,
		ToolChoice:    a.ToolChoice,
		System:        systemToString(a.System),
		Thinking:      th,
		ExtraBody:     map[string]interface{}{},
	}, nil
}

// normalizeAnthropicContent turns Anthropic content into the unified shape.
//
// A string stays a string. An array of blocks is preserved as-is: the unified
// format deliberately carries multimodal/tool blocks verbatim (ChatMessage
// .Content is json.RawMessage) rather than flattening them, because only an
// adapter knows which block types its upstream accepts.
func normalizeAnthropicContent(raw json.RawMessage) json.RawMessage {
	trimmed := strings.TrimSpace(string(raw))
	if trimmed == "" {
		return types.StringContent("")
	}
	if strings.HasPrefix(trimmed, `"`) {
		return raw
	}
	return raw
}

// systemToString accepts Anthropic's two spellings: a plain string, or an array
// of {type:"text", text} blocks.
func systemToString(raw json.RawMessage) string {
	trimmed := strings.TrimSpace(string(raw))
	if trimmed == "" || trimmed == "null" {
		return ""
	}
	if strings.HasPrefix(trimmed, `"`) {
		var s string
		if err := json.Unmarshal(raw, &s); err == nil {
			return s
		}
		return ""
	}
	var blocks []struct {
		Type string `json:"type"`
		Text string `json:"text"`
	}
	if err := json.Unmarshal(raw, &blocks); err != nil {
		return ""
	}
	var sb strings.Builder
	for _, b := range blocks {
		if b.Type == "text" {
			sb.WriteString(b.Text)
		}
	}
	return sb.String()
}

func chatKind(stream bool) string {
	if stream {
		return "stream"
	}
	return "chat"
}

// directProviderList hands the serving-source resolver its candidate list.
func directProviderList(direct *directRoute) []*provider.Provider {
	if direct == nil {
		return nil
	}
	return direct.cands
}

func tierForRoute(direct *directRoute) int {
	if direct != nil {
		return -1
	}
	return -2
}

func modelLabel(req *types.ChatRequest, usedModel string) string {
	if usedModel != "" {
		return usedModel
	}
	return req.Model
}