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

package openai

import (
	"crypto/sha256"
	"encoding/hex"
	"encoding/json"
	"net/http"
	"net/http/httptest"
	"os"
	"reflect"
	"strings"
	"testing"
	"time"

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

func TestAnthropicCanonicalAndAliasRoutes(t *testing.T) {
	fixture := mustReadAnthropicFixture(t, "native_message.json")
	candidate := anthropicTestCandidate(t, "anthropic")
	fake := &providerFakeRunService{
		poolDispatchPath:      string(edgeservice.ProviderPoolPathTunnel),
		poolSelectedCandidate: candidate,
		tunnelServedTarget:    "upstream-claude",
	}
	srv := NewServer(config.EdgeOpenAIConf{BearerToken: "iop-token"}, fake, nil)
	srv.SetModelCatalog([]config.ModelCatalogEntry{{ID: "claude-route", Providers: map[string]string{"anthropic": "upstream-claude"}}})

	for _, path := range []string{"/v1/messages", "/anthropic/v1/messages"} {
		t.Run(path, func(t *testing.T) {
			fake.tunnelFrames = anthropicTunnelFrames(http.StatusOK, "application/json", fixture)
			req := httptest.NewRequest(http.MethodPost, path, strings.NewReader(`{"model":"claude-route","max_tokens":32,"messages":[{"role":"user","content":"hello"}]}`))
			req.Header.Set("X-Api-Key", "iop-token")
			req.Header.Set(anthropicVersionHeader, anthropicSupportedVersion)
			w := httptest.NewRecorder()
			srv.routes().ServeHTTP(w, req)
			if w.Code != http.StatusOK {
				t.Fatalf("status: got %d body=%s", w.Code, w.Body.String())
			}
			if got := w.Body.Bytes(); !reflect.DeepEqual(got, fixture) {
				t.Fatalf("native body changed:\n got: %s\nwant: %s", got, fixture)
			}
		})
	}
	if got := fake.poolSubmitCountSnapshot(); got != 2 {
		t.Fatalf("provider-pool submissions: got %d, want 2", got)
	}
}

func TestAnthropicPrincipalBearerAndAPIKey(t *testing.T) {
	raw := "principal-token"
	sum := sha256.Sum256([]byte(raw))
	cfg := config.EdgeOpenAIConf{PrincipalTokens: []config.OpenAIPrincipalTokenConf{{
		TokenRef: "token-1", TokenHashSHA256: hex.EncodeToString(sum[:]), PrincipalRef: "principal-1",
	}}}
	srv := NewServer(cfg, &fakeRunService{}, nil)
	for _, headers := range []map[string]string{
		{"Authorization": "Bearer " + raw},
		{"X-Api-Key": raw},
		{"Authorization": "Bearer " + raw, "X-Api-Key": raw},
	} {
		req := httptest.NewRequest(http.MethodGet, "/v1/models", nil)
		req.Header.Set(anthropicVersionHeader, anthropicSupportedVersion)
		for key, value := range headers {
			req.Header.Set(key, value)
		}
		w := httptest.NewRecorder()
		srv.routes().ServeHTTP(w, req)
		if w.Code != http.StatusOK {
			t.Fatalf("headers=%v status=%d body=%s", headers, w.Code, w.Body.String())
		}
	}
}

func TestAnthropicPrincipalDualCredentialConflict(t *testing.T) {
	srv := NewServer(config.EdgeOpenAIConf{BearerToken: "first-secret"}, &fakeRunService{}, nil)
	for _, headers := range []map[string]string{
		{"Authorization": "Bearer first-secret", "X-Api-Key": "second-secret"},
		{"Authorization": "Basic first-secret", "X-Api-Key": "first-secret"},
	} {
		req := httptest.NewRequest(http.MethodGet, "/anthropic/v1/models", nil)
		for key, value := range headers {
			req.Header.Set(key, value)
		}
		req.Header.Set(anthropicVersionHeader, anthropicSupportedVersion)
		w := httptest.NewRecorder()
		srv.routes().ServeHTTP(w, req)
		if w.Code != http.StatusUnauthorized {
			t.Fatalf("headers=%v status: got %d body=%s", headers, w.Code, w.Body.String())
		}
		if strings.Contains(w.Body.String(), "first-secret") || strings.Contains(w.Body.String(), "second-secret") {
			t.Fatalf("credential leaked in response: %s", w.Body.String())
		}
	}
}

func TestManagedAnthropicDualHeadersSameToken(t *testing.T) {
	now := time.Date(2026, 8, 1, 12, 0, 0, 0, time.UTC)
	cache := authprojection.NewCache(authprojection.DefaultLimits(), func() time.Time { return now })
	if err := cache.Apply(managedProjectionFixture(1, now, time.Minute, "managed-anthropic-token")); err != nil {
		t.Fatal(err)
	}
	srv := NewServer(config.EdgeOpenAIConf{}, &fakeRunService{}, nil)
	setManagedPrincipalProjection(srv, cache)
	handlerCalls := 0
	handler := srv.withAuth(func(w http.ResponseWriter, r *http.Request) {
		handlerCalls++
		principal, ok := principalFromContext(r.Context())
		if !ok || principal.PrincipalRef != "managed-principal" || principal.Source != principalSourceProjection {
			t.Fatalf("managed principal: %+v ok=%v", principal, ok)
		}
		w.WriteHeader(http.StatusNoContent)
	})

	req := httptest.NewRequest(http.MethodGet, "/anthropic/v1/models", nil)
	req.Header.Set("Authorization", "Bearer managed-anthropic-token")
	req.Header.Set("X-Api-Key", "managed-anthropic-token")
	w := httptest.NewRecorder()
	handler(w, req)
	if w.Code != http.StatusNoContent || handlerCalls != 1 {
		t.Fatalf("status=%d handler_calls=%d body=%s", w.Code, handlerCalls, w.Body.String())
	}
}

func TestManagedAnthropicDualHeadersMismatch(t *testing.T) {
	now := time.Date(2026, 8, 1, 12, 0, 0, 0, time.UTC)
	cache := authprojection.NewCache(authprojection.DefaultLimits(), func() time.Time { return now })
	if err := cache.Apply(managedProjectionFixture(1, now, time.Minute, "managed-anthropic-token")); err != nil {
		t.Fatal(err)
	}
	srv := NewServer(config.EdgeOpenAIConf{BearerToken: "legacy-token"}, &fakeRunService{}, nil)
	setManagedPrincipalProjection(srv, cache)
	handlerCalls := 0
	handler := srv.withAuth(func(w http.ResponseWriter, _ *http.Request) {
		handlerCalls++
		w.WriteHeader(http.StatusNoContent)
	})

	for _, tc := range []struct {
		name   string
		bearer string
		apiKey string
	}{
		{name: "mismatch", bearer: "managed-anthropic-token", apiKey: "different-token"},
		{name: "unregistered", apiKey: "unregistered-token"},
		{name: "legacy source", bearer: "legacy-token", apiKey: "legacy-token"},
	} {
		t.Run(tc.name, func(t *testing.T) {
			req := httptest.NewRequest(http.MethodGet, "/anthropic/v1/models", nil)
			if tc.bearer != "" {
				req.Header.Set("Authorization", "Bearer "+tc.bearer)
			}
			if tc.apiKey != "" {
				req.Header.Set("X-Api-Key", tc.apiKey)
			}
			w := httptest.NewRecorder()
			handler(w, req)
			if w.Code != http.StatusUnauthorized || !strings.Contains(w.Body.String(), `"type":"authentication_error"`) {
				t.Fatalf("status=%d body=%s", w.Code, w.Body.String())
			}
		})
	}
	if err := cache.Apply(managedProjectionFixture(2, now, time.Minute, "")); err != nil {
		t.Fatal(err)
	}
	revokedReq := httptest.NewRequest(http.MethodGet, "/anthropic/v1/models", nil)
	revokedReq.Header.Set("X-Api-Key", "managed-anthropic-token")
	revokedW := httptest.NewRecorder()
	handler(revokedW, revokedReq)
	if revokedW.Code != http.StatusUnauthorized {
		t.Fatalf("revoked status=%d body=%s", revokedW.Code, revokedW.Body.String())
	}
	if handlerCalls != 0 {
		t.Fatalf("rejected managed auth reached handler %d times", handlerCalls)
	}
}

func TestManagedProjectionExpiryRejectsBeforeDispatch(t *testing.T) {
	clock := time.Date(2026, 8, 1, 12, 0, 0, 0, time.UTC)
	cache := authprojection.NewCache(authprojection.DefaultLimits(), func() time.Time { return clock })
	if err := cache.Apply(managedProjectionFixture(1, clock, time.Minute, "managed-anthropic-token")); err != nil {
		t.Fatal(err)
	}
	srv := NewServer(config.EdgeOpenAIConf{}, &fakeRunService{}, nil)
	setManagedPrincipalProjection(srv, cache)
	handlerCalls := 0
	handler := srv.withAuth(func(w http.ResponseWriter, _ *http.Request) {
		handlerCalls++
		w.WriteHeader(http.StatusNoContent)
	})
	clock = clock.Add(time.Minute)
	req := httptest.NewRequest(http.MethodGet, "/anthropic/v1/models", nil)
	req.Header.Set("X-Api-Key", "managed-anthropic-token")
	w := httptest.NewRecorder()
	handler(w, req)
	if w.Code != http.StatusUnauthorized || handlerCalls != 0 {
		t.Fatalf("expired status=%d handler_calls=%d body=%s", w.Code, handlerCalls, w.Body.String())
	}
}

func TestAnthropicModelsVersionSelectsShape(t *testing.T) {
	srv := NewServer(config.EdgeOpenAIConf{}, &fakeRunService{}, nil)
	srv.SetModelCatalog([]config.ModelCatalogEntry{{ID: "claude-route", DisplayName: "Claude Route"}})

	req := httptest.NewRequest(http.MethodGet, "/v1/models", nil)
	req.Header.Set(anthropicVersionHeader, anthropicSupportedVersion)
	req.Header.Set("User-Agent", "not-a-schema-selector")
	w := httptest.NewRecorder()
	srv.routes().ServeHTTP(w, req)
	if w.Code != http.StatusOK {
		t.Fatalf("status: got %d body=%s", w.Code, w.Body.String())
	}
	var got, want any
	if err := json.Unmarshal(w.Body.Bytes(), &got); err != nil {
		t.Fatal(err)
	}
	if err := json.Unmarshal(mustReadAnthropicFixture(t, "models_anthropic.json"), &want); err != nil {
		t.Fatal(err)
	}
	if !reflect.DeepEqual(got, want) {
		t.Fatalf("Anthropic model schema mismatch:\n got: %s\nwant: %s", w.Body.String(), mustReadAnthropicFixture(t, "models_anthropic.json"))
	}
	aliasReq := httptest.NewRequest(http.MethodGet, "/anthropic/v1/models", nil)
	aliasReq.Header.Set(anthropicVersionHeader, anthropicSupportedVersion)
	aliasW := httptest.NewRecorder()
	srv.routes().ServeHTTP(aliasW, aliasReq)
	if aliasW.Code != http.StatusOK || aliasW.Body.String() != w.Body.String() {
		t.Fatalf("alias schema mismatch: status=%d body=%s", aliasW.Code, aliasW.Body.String())
	}
	for _, req := range []*http.Request{
		httptest.NewRequest(http.MethodGet, "/anthropic/v1/models", nil),
		func() *http.Request {
			r := httptest.NewRequest(http.MethodGet, "/v1/models", nil)
			r.Header.Set(anthropicVersionHeader, "2024-01-01")
			return r
		}(),
	} {
		invalidW := httptest.NewRecorder()
		srv.routes().ServeHTTP(invalidW, req)
		if invalidW.Code != http.StatusBadRequest || !strings.Contains(invalidW.Body.String(), `"type":"invalid_request_error"`) {
			t.Fatalf("invalid version accepted: path=%s status=%d body=%s", req.URL.Path, invalidW.Code, invalidW.Body.String())
		}
	}

	openAIReq := httptest.NewRequest(http.MethodGet, "/v1/models", nil)
	openAIReq.Header.Set("User-Agent", "claude-code")
	openAIW := httptest.NewRecorder()
	srv.routes().ServeHTTP(openAIW, openAIReq)
	if !strings.Contains(openAIW.Body.String(), `"object":"list"`) || strings.Contains(openAIW.Body.String(), `"has_more"`) {
		t.Fatalf("User-Agent selected the wrong model schema: %s", openAIW.Body.String())
	}
}

func TestAnthropicCountTokensNativeLocalUnsupported(t *testing.T) {
	requestBody := `{"model":"claude-route","system":"Count this.","messages":[{"role":"user","content":"hello tokens"}]}`

	t.Run("local deterministic", func(t *testing.T) {
		fake := &providerFakeRunService{}
		srv := NewServer(config.EdgeOpenAIConf{}, fake, nil)
		srv.SetModelCatalog([]config.ModelCatalogEntry{{
			ID: "claude-route", Providers: map[string]string{"chat": "served-chat"},
			TokenCounter: &config.TokenCounterConf{Mode: config.TokenCounterDeterministic},
		}})
		w := serveAnthropicRequest(srv, "/v1/messages/count_tokens", requestBody)
		if w.Code != http.StatusOK || !strings.Contains(w.Body.String(), `"input_tokens":`) {
			t.Fatalf("status=%d body=%s", w.Code, w.Body.String())
		}
		if got := fake.poolSubmitCountSnapshot(); got != 0 {
			t.Fatalf("local counter dispatched provider pool %d times", got)
		}
	})

	t.Run("native", func(t *testing.T) {
		fake := &providerFakeRunService{
			poolDispatchPath:      string(edgeservice.ProviderPoolPathTunnel),
			poolSelectedCandidate: anthropicTestCandidate(t, "anthropic"),
			tunnelServedTarget:    "upstream-claude",
			tunnelFrames:          anthropicTunnelFrames(http.StatusOK, "application/json", []byte(`{"input_tokens":7}`)),
		}
		srv := NewServer(config.EdgeOpenAIConf{}, fake, nil)
		srv.SetModelCatalog([]config.ModelCatalogEntry{{ID: "claude-route", Providers: map[string]string{"anthropic": "upstream-claude"}}})
		w := serveAnthropicRequest(srv, "/v1/messages/count_tokens", requestBody)
		if w.Code != http.StatusOK || w.Body.String() != `{"input_tokens":7}` {
			t.Fatalf("status=%d body=%s", w.Code, w.Body.String())
		}
		reqs := fake.tunnelReqsSnapshot()
		if len(reqs) != 1 || reqs[0].Operation != string(config.OperationCountTokens) {
			t.Fatalf("native count request mismatch: %+v", reqs)
		}
	})

	t.Run("unsupported chat profile", func(t *testing.T) {
		fake := &providerFakeRunService{
			poolDispatchPath:      string(edgeservice.ProviderPoolPathTunnel),
			poolSelectedCandidate: anthropicTestCandidate(t, "openai"),
		}
		srv := NewServer(config.EdgeOpenAIConf{}, fake, nil)
		srv.SetModelCatalog([]config.ModelCatalogEntry{{ID: "claude-route", Providers: map[string]string{"chat": "served-chat"}}})
		w := serveAnthropicRequest(srv, "/v1/messages/count_tokens", requestBody)
		if w.Code != http.StatusBadRequest || !strings.Contains(w.Body.String(), `"type":"not_supported_error"`) {
			t.Fatalf("status=%d body=%s", w.Code, w.Body.String())
		}
		if got := len(fake.tunnelReqsSnapshot()); got != 0 {
			t.Fatalf("unsupported count reached provider wire: %d requests", got)
		}
	})
}

func serveAnthropicRequest(srv *Server, path, body string) *httptest.ResponseRecorder {
	req := httptest.NewRequest(http.MethodPost, path, strings.NewReader(body))
	req.Header.Set(anthropicVersionHeader, anthropicSupportedVersion)
	w := httptest.NewRecorder()
	srv.routes().ServeHTTP(w, req)
	return w
}

func anthropicTestCandidate(t *testing.T, profileID string) edgeservice.ProviderPoolCandidate {
	t.Helper()
	profile, err := config.ResolveProtocolProfile(profileID, "", config.BuiltInProtocolProfileCatalog())
	if err != nil {
		t.Fatal(err)
	}
	return edgeservice.ProviderPoolCandidate{
		ActualModel: "served-model", ProviderID: profileID + "-provider",
		ExecutionPath: string(edgeservice.ProviderPoolPathTunnel), ProfileID: profile.ID,
		ProfileDriver: string(profile.Driver), ProfileCapabilities: append([]string(nil), profile.Capabilities...),
		ProtocolProfile: &profile,
	}
}

func anthropicTunnelFrames(status int, contentType string, chunks ...[]byte) chan *iop.ProviderTunnelFrame {
	frames := make(chan *iop.ProviderTunnelFrame, len(chunks)+2)
	frames <- &iop.ProviderTunnelFrame{
		Kind:       iop.ProviderTunnelFrameKind_PROVIDER_TUNNEL_FRAME_KIND_RESPONSE_START,
		StatusCode: int32(status), Headers: map[string]string{"Content-Type": contentType, "X-Request-Id": "req_fixture"},
	}
	for index, chunk := range chunks {
		frames <- &iop.ProviderTunnelFrame{
			Kind:     iop.ProviderTunnelFrameKind_PROVIDER_TUNNEL_FRAME_KIND_BODY,
			Sequence: int64(index + 1), Body: append([]byte(nil), chunk...),
		}
	}
	frames <- &iop.ProviderTunnelFrame{Kind: iop.ProviderTunnelFrameKind_PROVIDER_TUNNEL_FRAME_KIND_END, End: true}
	close(frames)
	return frames
}

func mustReadAnthropicFixture(t *testing.T, name string) []byte {
	t.Helper()
	body, err := os.ReadFile("testdata/anthropic/" + name)
	if err != nil {
		t.Fatal(err)
	}
	return body
}