iop/apps/edge/internal/openai/credential_migration_test.go

package openai

import (
	"encoding/json"
	"net/http"
	"net/http/httptest"
	"reflect"
	"strings"
	"testing"
	"time"

	"iop/apps/edge/internal/authprojection"
	edgeservice "iop/apps/edge/internal/service"
	"iop/packages/go/config"
)

func managedCredentialMigrationServer(t *testing.T, fake *providerFakeRunService) *Server {
	t.Helper()
	now := time.Date(2026, 8, 2, 0, 0, 0, 0, time.UTC)
	cache := authprojection.NewCache(authprojection.DefaultLimits(), func() time.Time { return now })
	projection := makeTestProjection(1, now, time.Hour, map[string]string{
		"managed-iop-token": "principal-one",
	}, map[string]authprojection.Route{
		"route": {
			RouteID: "public-route", PrincipalRef: "principal-one", CredentialSlotRef: "slot-ref-one",
			ProfileID: "openai", UpstreamModel: "served-model", ResourceSelector: "openai-provider",
			RouteRevision: 3, CredentialRevision: 7,
		},
	})
	if err := cache.Apply(projection); err != nil {
		t.Fatal(err)
	}
	srv := NewServer(config.EdgeOpenAIConf{}, fake, nil)
	setManagedPrincipalProjection(srv, cache)
	srv.SetModelCatalog([]config.ModelCatalogEntry{{
		ID: "internal-model", Providers: map[string]string{"openai-provider": "served-model"},
	}})
	return srv
}

func TestManagedCredentialModeRejectsCallerProviderHeadersBeforeEverySurface(t *testing.T) {
	fake := &providerFakeRunService{}
	srv := managedCredentialMigrationServer(t, fake)
	tests := []struct {
		name      string
		path      string
		body      string
		anthropic bool
	}{
		{name: "chat", path: "/v1/chat/completions", body: `{"model":"public-route","messages":[{"role":"user","content":"hello"}]}`},
		{name: "responses", path: "/v1/responses", body: `{"model":"public-route","input":"hello"}`},
		{name: "anthropic", path: "/v1/messages", body: `{"model":"public-route","max_tokens":8,"messages":[{"role":"user","content":"hello"}]}`, anthropic: true},
	}
	for _, tc := range tests {
		t.Run(tc.name, func(t *testing.T) {
			req := httptest.NewRequest(http.MethodPost, tc.path, strings.NewReader(tc.body))
			req.Header.Set("Authorization", "Bearer managed-iop-token")
			req.Header.Set(legacyProviderCredentialHeader, "provider-secret-sentinel")
			if tc.anthropic {
				req.Header.Set(anthropicVersionHeader, anthropicSupportedVersion)
			}
			w := httptest.NewRecorder()
			srv.routes().ServeHTTP(w, req)
			if w.Code != http.StatusBadRequest {
				t.Fatalf("status = %d, want 400; body=%s", w.Code, w.Body.String())
			}
			if strings.Contains(w.Body.String(), "provider-secret-sentinel") {
				t.Fatalf("rejection echoed caller credential: %s", w.Body.String())
			}
			var body map[string]any
			if err := json.Unmarshal(w.Body.Bytes(), &body); err != nil {
				t.Fatalf("decode rejection: %v", err)
			}
			if tc.anthropic {
				if !reflect.DeepEqual(sortedMapKeys(body), []string{"error", "type"}) || body["type"] != "error" {
					t.Fatalf("Anthropic rejection schema changed: %v", body)
				}
			} else if !reflect.DeepEqual(sortedMapKeys(body), []string{"error"}) {
				t.Fatalf("OpenAI rejection schema changed: %v", body)
			}
		})
	}
	if got := fake.poolSubmitCountSnapshot(); got != 0 {
		t.Fatalf("managed caller credential reached provider-pool dispatch %d times", got)
	}
	if got := len(fake.tunnelReqsSnapshot()); got != 0 {
		t.Fatalf("managed caller credential reached provider tunnel %d times", got)
	}
}

func sortedMapKeys(values map[string]any) []string {
	keys := make([]string, 0, len(values))
	for key := range values {
		keys = append(keys, key)
	}
	if len(keys) == 2 && keys[0] > keys[1] {
		keys[0], keys[1] = keys[1], keys[0]
	}
	return keys
}

func TestLegacyCredentialModeForwardsOnlyTheDedicatedProviderHeader(t *testing.T) {
	for _, tc := range []struct {
		name      string
		path      string
		body      string
		anthropic bool
	}{
		{name: "chat", path: "/v1/chat/completions", body: `{"model":"legacy-model","messages":[{"role":"user","content":"hello"}]}`},
		{name: "responses", path: "/v1/responses", body: `{"model":"legacy-model","input":"hello"}`},
		{name: "anthropic", path: "/v1/messages", body: `{"model":"legacy-model","max_tokens":8,"messages":[{"role":"user","content":"hello"}]}`, anthropic: true},
	} {
		t.Run(tc.name, func(t *testing.T) {
			fake := &providerFakeRunService{poolDispatchPath: string(edgeservice.ProviderPoolPathTunnel)}
			if tc.anthropic {
				fake.poolSelectedCandidate = anthropicTestCandidate(t, "anthropic")
				fake.tunnelFrames = anthropicTunnelFrames(http.StatusOK, "application/json", []byte(`{"type":"message","role":"assistant","content":[],"stop_reason":"end_turn","usage":{"input_tokens":1,"output_tokens":1}}`))
			} else {
				fake.tunnelFrames = staticProviderTunnelFrames(`{"id":"fixture","object":"response"}`)
			}
			cfg := config.EdgeOpenAIConf{
				BearerToken: "legacy-iop-token",
				ProviderAuth: config.EdgeOpenAIProviderAuthConf{
					Enabled: true, FromHeader: legacyProviderCredentialHeader, TargetHeader: "Authorization", Scheme: "Bearer", Required: true,
				},
			}
			srv := NewServer(cfg, fake, nil)
			srv.SetModelCatalog([]config.ModelCatalogEntry{{
				ID: "legacy-model", Providers: map[string]string{"legacy-provider": "served-model"},
			}})
			req := httptest.NewRequest(http.MethodPost, tc.path, strings.NewReader(tc.body))
			req.Header.Set("Authorization", "Bearer legacy-iop-token")
			req.Header.Set(legacyProviderCredentialHeader, "legacy-provider-token")
			if tc.anthropic {
				req.Header.Set(anthropicVersionHeader, anthropicSupportedVersion)
			}
			w := httptest.NewRecorder()
			srv.routes().ServeHTTP(w, req)
			if w.Code != http.StatusOK {
				t.Fatalf("status = %d, want 200; body=%s", w.Code, w.Body.String())
			}
			requests := fake.tunnelReqsSnapshot()
			if len(requests) != 1 {
				t.Fatalf("provider tunnel calls = %d, want 1", len(requests))
			}
			found := 0
			for header, value := range requests[0].Headers {
				if strings.Contains(value, "legacy-provider-token") {
					found++
					if tc.anthropic && !strings.EqualFold(header, "x-api-key") {
						t.Fatalf("Anthropic credential header = %q, want x-api-key", header)
					}
				}
				if strings.Contains(value, "legacy-iop-token") {
					t.Fatalf("inbound IOP token was reused as provider credential in %q", header)
				}
			}
			if found != 1 {
				t.Fatalf("provider credential occurrences = %d, want 1; headers=%v", found, requests[0].Headers)
			}
		})
	}
}

func TestManagedResponsesProviderSchemaRemainsUnchanged(t *testing.T) {
	profile, err := config.ResolveProtocolProfile("openai", "", config.BuiltInProtocolProfileCatalog())
	if err != nil {
		t.Fatal(err)
	}
	want := `{"id":"resp-managed","object":"response","status":"completed","output":[{"type":"message","content":[{"type":"output_text","text":"ok"}]}],"usage":{"input_tokens":2,"output_tokens":1}}`
	fake := &providerFakeRunService{
		tunnelFrames: staticProviderTunnelFrames(want),
		poolSelectedCandidate: edgeservice.ProviderPoolCandidate{
			ProviderID: "openai-provider", ActualModel: "served-model", ExecutionPath: string(edgeservice.ProviderPoolPathTunnel),
			ProfileID: profile.ID, ProfileDriver: string(profile.Driver),
			ProfileCapabilities: append([]string(nil), profile.Capabilities...), ProtocolProfile: &profile,
		},
	}
	srv := managedCredentialMigrationServer(t, fake)
	req := httptest.NewRequest(http.MethodPost, "/v1/responses", strings.NewReader(`{"model":"public-route","input":"hello"}`))
	req.Header.Set("Authorization", "Bearer managed-iop-token")
	w := httptest.NewRecorder()
	srv.routes().ServeHTTP(w, req)
	if w.Code != http.StatusOK || w.Body.String() != want {
		t.Fatalf("managed Responses schema/body changed: status=%d\n got=%s\nwant=%s", w.Code, w.Body.String(), want)
	}
}