easyai-ai-gateway/apps/api/internal/httpapi/compat_protocol_test.go

package httpapi

import (
	"context"
	"encoding/json"
	"errors"
	"net/http"
	"net/http/httptest"
	"strings"
	"testing"

	"github.com/easyai/easyai-ai-gateway/apps/api/internal/auth"
	"github.com/easyai/easyai-ai-gateway/apps/api/internal/clients"
	"github.com/easyai/easyai-ai-gateway/apps/api/internal/store"
)

func TestProtocolAPIKeyStoreFailureIs503InsteadOf401(t *testing.T) {
	authenticator := auth.New("test-secret", "", "")
	authenticator.LocalAPIKeyVerifier = func(context.Context, string) (*auth.User, error) {
		return nil, errors.New("database unavailable")
	}
	server := &Server{auth: authenticator}
	handler := server.requireProtocolUser(clients.ProtocolGeminiGenerateContent, http.HandlerFunc(
		func(http.ResponseWriter, *http.Request) {
			t.Fatal("unavailable authentication store must not call the protocol handler")
		},
	))
	request := httptest.NewRequest(http.MethodPost, "/v1beta/models/test:generateContent", nil)
	request.Header.Set("Authorization", "Bearer sk-gw-local")
	recorder := httptest.NewRecorder()

	handler.ServeHTTP(recorder, request)

	if recorder.Code != http.StatusServiceUnavailable {
		t.Fatalf("status=%d, want 503; body=%s", recorder.Code, recorder.Body.String())
	}
	if recorder.Header().Get("Retry-After") != "2" {
		t.Fatalf("Retry-After=%q, want 2", recorder.Header().Get("Retry-After"))
	}
	var body map[string]any
	if err := json.Unmarshal(recorder.Body.Bytes(), &body); err != nil {
		t.Fatalf("decode response: %v", err)
	}
	errorBody := requireObject(t, body["error"])
	if errorBody["code"] != float64(http.StatusServiceUnavailable) || errorBody["status"] != "UNAVAILABLE" {
		t.Fatalf("unexpected Gemini unavailable response: %+v", body)
	}
}

func TestProtocolErrorsUseOfficialShapesWithoutGatewayExtensions(t *testing.T) {
	tests := []struct {
		name       string
		protocol   string
		status     int
		assertBody func(*testing.T, map[string]any)
	}{
		{
			name: "openai", protocol: clients.ProtocolOpenAIChatCompletions, status: http.StatusBadRequest,
			assertBody: func(t *testing.T, body map[string]any) {
				errorBody := requireObject(t, body["error"])
				if errorBody["type"] != "invalid_request_error" || errorBody["code"] != "unsupported_kind" {
					t.Fatalf("unexpected OpenAI error: %+v", body)
				}
				assertNoKeys(t, errorBody, "status", "retryable", "taskId", "gateway_status")
			},
		},
		{
			name: "gemini", protocol: clients.ProtocolGeminiGenerateContent, status: http.StatusTooManyRequests,
			assertBody: func(t *testing.T, body map[string]any) {
				errorBody := requireObject(t, body["error"])
				if errorBody["code"] != float64(http.StatusTooManyRequests) || errorBody["status"] != "RESOURCE_EXHAUSTED" {
					t.Fatalf("unexpected Gemini error: %+v", body)
				}
				assertNoKeys(t, errorBody, "retryable", "taskId", "gateway_status")
			},
		},
		{
			name: "volces", protocol: clients.ProtocolVolcesContents, status: http.StatusBadGateway,
			assertBody: func(t *testing.T, body map[string]any) {
				errorBody := requireObject(t, body["error"])
				if errorBody["code"] != "upstream_submission_unknown" {
					t.Fatalf("unexpected Volces error: %+v", body)
				}
				assertNoKeys(t, errorBody, "status", "retryable", "taskId", "gateway_status")
			},
		},
	}
	for _, test := range tests {
		t.Run(test.name, func(t *testing.T) {
			recorder := httptest.NewRecorder()
			code := "unsupported_kind"
			if test.name == "gemini" {
				code = "rate_limit"
			} else if test.name == "volces" {
				code = "upstream_submission_unknown"
			}
			writeProtocolError(recorder, test.protocol, test.status, "failed", nil, code)
			if recorder.Code != test.status {
				t.Fatalf("status=%d, want %d", recorder.Code, test.status)
			}
			var body map[string]any
			if err := json.Unmarshal(recorder.Body.Bytes(), &body); err != nil {
				t.Fatalf("decode response: %v", err)
			}
			test.assertBody(t, body)
		})
	}
}

func TestWireResponsePassthroughPreservesStatusUnknownFieldsAndAllowedHeaders(t *testing.T) {
	recorder := httptest.NewRecorder()
	wire := &clients.WireResponse{
		Protocol:   clients.ProtocolOpenAIResponses,
		StatusCode: http.StatusAccepted,
		Headers: map[string][]string{
			"Content-Type": {"application/json"},
			"X-Request-Id": {"req_official"},
			"Retry-After":  {"2"},
		},
		RawJSON: []byte(`{"id":"resp_1","future_official_field":{"v":1}}`),
	}
	writeWireResponse(recorder, wire)
	if recorder.Code != http.StatusAccepted || recorder.Header().Get("X-Request-Id") != "req_official" || recorder.Header().Get("Retry-After") != "2" {
		t.Fatalf("wire metadata was not preserved: status=%d headers=%v", recorder.Code, recorder.Header())
	}
	if got := strings.TrimSpace(recorder.Body.String()); got != string(wire.RawJSON) {
		t.Fatalf("wire body changed: %s", got)
	}
	if !wireResponseMatches(wire, clients.ProtocolOpenAIResponses) {
		t.Fatal("native unconverted wire response must be eligible for passthrough")
	}
	wire.Converted = true
	if wireResponseMatches(wire, clients.ProtocolOpenAIResponses) {
		t.Fatal("converted wire response must not be eligible for passthrough")
	}
}

func TestCompatibilityStatusMappings(t *testing.T) {
	for internal, want := range map[string]string{
		"queued": "queued", "running": "running", "succeeded": "succeeded", "failed": "failed", "cancelled": "cancelled",
	} {
		if got := volcesCompatibleTaskStatus(internal); got != want {
			t.Fatalf("Volces status %q=%q, want %q", internal, got, want)
		}
	}
	for internal, want := range map[string]string{
		"queued": "submitted", "running": "processing", "succeeded": "succeed", "failed": "failed", "cancelled": "failed",
	} {
		if got := klingV1Status(internal); got != want {
			t.Fatalf("Kling V1 status %q=%q, want %q", internal, got, want)
		}
	}
	for internal, want := range map[string]string{
		"queued": "submitted", "running": "processing", "succeeded": "succeeded", "failed": "failed", "cancelled": "failed",
	} {
		if got := klingV2Status(internal); got != want {
			t.Fatalf("Kling V2 status %q=%q, want %q", internal, got, want)
		}
	}
}

func TestNativeOpenAIStreamPreservesUnknownEventFields(t *testing.T) {
	recorder := httptest.NewRecorder()
	request := httptest.NewRequest(http.MethodPost, "/v1/chat/completions", nil)
	executor := &fakeTaskExecutor{
		deltas: []clients.StreamDeltaEvent{{
			Event: map[string]any{
				"id": "chatcmpl_1", "object": "chat.completion.chunk",
				"choices":               []any{map[string]any{"index": 0, "delta": map[string]any{"content": "hi"}}},
				"future_official_field": map[string]any{"enabled": true},
			},
			WireProtocol: clients.ProtocolOpenAIChatCompletions,
			WireHeaders:  map[string][]string{"X-Request-Id": {"req_stream_1"}},
		}},
		output: map[string]any{"id": "chatcmpl_1"},
	}
	writeProtocolCompatibleTaskResponse(
		context.Background(), recorder, request, executor,
		"chat.completions", "gpt-test", clients.ProtocolOpenAIChatCompletions,
		store.GatewayTask{ID: "gateway-task"}, &auth.User{}, true, false,
	)
	if recorder.Header().Get("X-Request-Id") != "req_stream_1" {
		t.Fatalf("official stream header was lost: %+v", recorder.Header())
	}
	body := recorder.Body.String()
	if !strings.Contains(body, `"future_official_field":{"enabled":true}`) || strings.Count(body, "data: [DONE]") != 1 {
		t.Fatalf("native stream event was reconstructed or duplicated: %s", body)
	}
	if strings.Contains(body, "gateway-task") {
		t.Fatalf("native stream leaked gateway task id: %s", body)
	}
}

func requireObject(t *testing.T, value any) map[string]any {
	t.Helper()
	result, ok := value.(map[string]any)
	if !ok {
		t.Fatalf("value is not an object: %#v", value)
	}
	return result
}

func assertNoKeys(t *testing.T, object map[string]any, keys ...string) {
	t.Helper()
	for _, key := range keys {
		if _, ok := object[key]; ok {
			t.Fatalf("unexpected gateway extension %q in %+v", key, object)
		}
	}
}