internal/gateway/edge_cases_test.go

Rama main · réplica de lectura del repositorio original.

package gateway

import (
	"context"
	"encoding/json"
	"io"
	"net/http"
	"strconv"
	"strings"
	"testing"

	batchstore "github.com/enterpilot/gomodel/internal/batch"
	"github.com/enterpilot/gomodel/internal/core"
	"github.com/enterpilot/gomodel/internal/usage"
	"github.com/stretchr/testify/require"
)

func TestMergeStoredBatchFromUpstreamPreservesGatewayOwnedMetadata(t *testing.T) {
	stored := &batchstore.StoredBatch{
		Batch: &core.BatchResponse{
			Metadata: map[string]string{
				"provider":          "openai",
				"provider_batch_id": "batch-primary",
				"client":            "original",
			},
		},
	}
	upstream := &core.BatchResponse{
		Metadata: map[string]string{
			"provider":          "anthropic",
			"provider_batch_id": "batch-upstream",
			"client":            "upstream",
		},
	}

	MergeStoredBatchFromUpstream(stored, upstream)
	got := stored.Batch.Metadata["provider"]
	require.Equal(t, "openai", got)
	got = stored.Batch.Metadata["provider_batch_id"]
	require.Equal(t, "batch-primary", got)
	got = stored.Batch.Metadata["client"]
	require.Equal(t, "upstream", got)
}

func TestDetermineBatchExecutionSelectionRejectsNilRequest(t *testing.T) {
	_, err := DetermineBatchExecutionSelectionWithAuthorizerAndInputFileResolver(context.Background(), nil, nil, nil, nil, nil)
	require.Error(t, err)

	var gatewayErr *core.GatewayError
	require.ErrorAs(t, err, &gatewayErr)
	require.Equal(t, core.ErrorTypeInvalidRequest, gatewayErr.Type)
	require.Equal(t, "batch request is required", gatewayErr.Message)
}

func TestExtractTokenTotalsOnlySynthesizesAuthoritativeTotals(t *testing.T) {
	input, output, total, hasUsage, hasTotal := extractTokenTotals(map[string]any{
		"input_tokens": 10,
	})
	require.Equal(t, 10, input)
	require.Equal(t, 0, output)
	require.Equal(t, 0, total)
	require.True(t, hasUsage)
	require.False(t, hasTotal)

	input, output, total, hasUsage, hasTotal = extractTokenTotals(map[string]any{
		"input_tokens":  10,
		"output_tokens": 5,
	})
	require.Equal(t, 10, input)
	require.Equal(t, 5, output)
	require.Equal(t, 15, total)
	require.True(t, hasUsage)
	require.True(t, hasTotal)
}

func TestIntFromFloat64RejectsBoundaryOverflow(t *testing.T) {
	outOfRange := float64(uint64(1) << (strconv.IntSize - 1))
	_, ok := intFromFloat64(outOfRange)
	require.False(t, ok)
	_, ok = intFromFloat64(1.9)
	require.False(t, ok)
}

func TestCloneRequestsForSelectorCopiesMutableFields(t *testing.T) {
	includeUsage := false
	chatReq := &core.ChatRequest{
		Model:         "alias",
		Provider:      "router",
		Messages:      []core.Message{{Role: "user", ToolCalls: []core.ToolCall{{ID: "call-1"}}}},
		Tools:         []map[string]any{{"type": "function"}},
		StreamOptions: &core.StreamOptions{IncludeUsage: includeUsage},
		Reasoning:     &core.Reasoning{Effort: "low"},
	}

	chatClone := CloneChatRequestForSelector(chatReq, core.ModelSelector{Provider: "openai", Model: "gpt-4o-mini"})
	chatClone.Messages[0].Role = "assistant"
	chatClone.Messages[0].ToolCalls[0].ID = "call-2"
	chatClone.Tools[0]["type"] = "changed"
	chatClone.StreamOptions.IncludeUsage = true
	chatClone.Reasoning.Effort = "high"

	require.Equal(t, "user", chatReq.Messages[0].Role)
	require.Equal(t, "call-1", chatReq.Messages[0].ToolCalls[0].ID, "chat messages were shared with clone: %#v", chatReq.Messages)
	got := chatReq.Tools[0]["type"]
	require.Equal(t, "function", got)
	require.False(t, chatReq.StreamOptions.IncludeUsage)
	require.Equal(t, "low", chatReq.Reasoning.Effort)

	responsesReq := &core.ResponsesRequest{
		Model:         "alias",
		Provider:      "router",
		Tools:         []map[string]any{{"type": "function"}},
		Metadata:      map[string]string{"client": "original"},
		StreamOptions: &core.StreamOptions{},
		Reasoning:     &core.Reasoning{Effort: "low"},
	}

	responsesClone := CloneResponsesRequestForSelector(responsesReq, core.ModelSelector{Provider: "openai", Model: "gpt-4o-mini"})
	responsesClone.Tools[0]["type"] = "changed"
	responsesClone.Metadata["client"] = "clone"
	responsesClone.StreamOptions.IncludeUsage = true
	responsesClone.Reasoning.Effort = "high"
	got = responsesReq.Tools[0]["type"]
	require.Equal(t, "function", got)

	require.Equal(t, "original", responsesReq.Metadata["client"])
	require.False(t, responsesReq.StreamOptions.IncludeUsage)
	require.Equal(t, "low", responsesReq.Reasoning.Effort)
}

func TestCloneEmbeddingRequestForSelectorPreservesRequestFields(t *testing.T) {
	dimensions := 256
	req := &core.EmbeddingRequest{
		Model:          "alias",
		Provider:       "router",
		Input:          []any{"first", map[string]any{"text": "second"}},
		Dimensions:     &dimensions,
		EncodingFormat: "base64",
		ExtraFields: core.UnknownJSONFieldsFromMap(map[string]json.RawMessage{
			"x_custom": json.RawMessage(`{"mode":"test"}`),
		}),
	}

	cloned := CloneEmbeddingRequestForSelector(req, core.ModelSelector{Provider: "openai-backup", Model: "embedding-model"})
	cloned.ExtraFields = core.UnknownJSONFieldsFromMap(map[string]json.RawMessage{
		"x_custom": json.RawMessage(`{"mode":"changed"}`),
	})

	require.Equal(t, "embedding-model", cloned.Model)
	require.Equal(t, "openai-backup", cloned.Provider)
	require.Equal(t, "first", cloned.Input.([]any)[0])
	require.Equal(t, "second", req.Input.([]any)[1].(map[string]any)["text"])
	require.NotSame(t, req, cloned)
	require.Same(t, req.Dimensions, cloned.Dimensions)
	require.Equal(t, 256, *req.Dimensions)
	require.Equal(t, req.Input, cloned.Input)
	require.Equal(t, "base64", cloned.EncodingFormat)
	require.JSONEq(t, `{"mode":"test"}`, string(req.ExtraFields.Lookup("x_custom")))
}

func TestCloneEmbeddingRequestForSelectorPreservesTypedInputs(t *testing.T) {
	for _, input := range []any{[]string{"first", "second"}, []int{2, 3}, [][]int{{2, 3}, {5}}} {
		req := &core.EmbeddingRequest{Input: input}
		cloned := CloneEmbeddingRequestForSelector(req, core.ModelSelector{Model: "backup"})
		require.IsType(t, input, cloned.Input)
		require.Equal(t, req.Input, cloned.Input, "cloning must preserve JSON value and number semantics")
	}
}

func TestShouldEnforceReturningUsageDataRequiresEnabledLogger(t *testing.T) {
	orchestrator := NewInferenceOrchestrator(InferenceConfig{
		UsageLogger: &usageCaptureLogger{
			config: usage.Config{
				Enabled:                   false,
				EnforceReturningUsageData: true,
			},
		},
	})

	require.False(t, orchestrator.ShouldEnforceReturningUsageData())
}

func TestStreamResponsesRejectsNilRequest(t *testing.T) {
	orchestrator := NewInferenceOrchestrator(InferenceConfig{Provider: &providerTypeResolverStub{}})

	_, err := orchestrator.StreamResponses(context.Background(), nil, nil)
	require.Error(t, err)

	var gatewayErr *core.GatewayError
	require.ErrorAs(t, err, &gatewayErr)
	require.Equal(t, core.ErrorTypeInvalidRequest, gatewayErr.Type)
}

func TestDispatchChatCompletionRejectsEmptyProviderResponse(t *testing.T) {
	orchestrator := NewInferenceOrchestrator(InferenceConfig{Provider: &providerTypeResolverStub{}})

	_, _, err := orchestrator.DispatchChatCompletion(context.Background(), nil, &core.ChatRequest{Model: "gpt-4o-mini"})
	require.Error(t, err)

	var gatewayErr *core.GatewayError
	require.ErrorAs(t, err, &gatewayErr)
	require.Equal(t, core.ErrorTypeProvider, gatewayErr.Type)
	require.Equal(t, http.StatusBadGateway, gatewayErr.HTTPStatusCode())
}

func TestStreamResponsesRejectsEmptyProviderStream(t *testing.T) {
	orchestrator := NewInferenceOrchestrator(InferenceConfig{Provider: &providerTypeResolverStub{}})

	_, err := orchestrator.StreamResponses(context.Background(), nil, &core.ResponsesRequest{Model: "gpt-4o-mini"})
	require.Error(t, err)

	var gatewayErr *core.GatewayError
	require.ErrorAs(t, err, &gatewayErr)
	require.Equal(t, core.ErrorTypeProvider, gatewayErr.Type)
	require.Equal(t, http.StatusBadGateway, gatewayErr.HTTPStatusCode())
}

func TestStreamResponsesFallsBackAfterEmptyPrimaryStream(t *testing.T) {
	provider := &streamFailoverProvider{
		streamsByModel: map[string]io.ReadCloser{
			"fallback": io.NopCloser(strings.NewReader("data: {}\n\n")),
		},
	}
	orchestrator := NewInferenceOrchestrator(InferenceConfig{
		Provider: provider,
		FailoverResolver: failoverResolverFunc(func(*core.RequestModelResolution, core.Operation) []core.ModelSelector {
			return []core.ModelSelector{{Provider: "openai", Model: "fallback"}}
		}),
	})
	workflow := &core.Workflow{
		Endpoint: core.DescribeEndpoint(http.MethodPost, "/v1/responses"),
		Resolution: &core.RequestModelResolution{
			ResolvedSelector: core.ModelSelector{Provider: "openai", Model: "primary"},
			ProviderType:     "openai",
		},
		Policy: &core.ResolvedWorkflowPolicy{
			VersionID: "workflow-fallback",
			Features: core.WorkflowFeatures{
				Cache:      true,
				Audit:      true,
				Usage:      true,
				Guardrails: true,
				Failover:   true,
			},
		},
	}

	result, err := orchestrator.StreamResponses(context.Background(), workflow, &core.ResponsesRequest{Model: "primary"})
	require.NoError(t, err)

	defer result.Stream.Close()

	require.True(t, result.Meta.UsedFailover)
	require.Equal(t, "openai/fallback", result.Meta.FailoverModel)
	got := strings.Join(provider.responseStreamCalls, ",")
	require.Equal(t, "primary,fallback", got)
}

type failoverResolverFunc func(*core.RequestModelResolution, core.Operation) []core.ModelSelector

func (f failoverResolverFunc) ResolveFailovers(resolution *core.RequestModelResolution, op core.Operation) []core.ModelSelector {
	return f(resolution, op)
}

type streamFailoverProvider struct {
	streamsByModel      map[string]io.ReadCloser
	responseStreamCalls []string
}

func (p *streamFailoverProvider) ChatCompletion(context.Context, *core.ChatRequest) (*core.ChatResponse, error) {
	return nil, nil
}

func (p *streamFailoverProvider) StreamChatCompletion(context.Context, *core.ChatRequest) (io.ReadCloser, error) {
	return nil, nil
}

func (p *streamFailoverProvider) ListModels(context.Context) (*core.ModelsResponse, error) {
	return nil, nil
}

func (p *streamFailoverProvider) Responses(context.Context, *core.ResponsesRequest) (*core.ResponsesResponse, error) {
	return nil, nil
}

func (p *streamFailoverProvider) StreamResponses(_ context.Context, req *core.ResponsesRequest) (io.ReadCloser, error) {
	p.responseStreamCalls = append(p.responseStreamCalls, req.Model)
	return p.streamsByModel[req.Model], nil
}

func (p *streamFailoverProvider) Embeddings(context.Context, *core.EmbeddingRequest) (*core.EmbeddingResponse, error) {
	return nil, nil
}

func (p *streamFailoverProvider) Supports(string) bool { return true }

func (p *streamFailoverProvider) GetProviderType(model string) string {
	selector, err := core.ParseModelSelector(model, "")
	if err == nil && selector.Provider != "" {
		return selector.Provider
	}
	return ""
}