ai-router/internal/server/discovery_test.go

package server

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

const discoveryPath = "/admin/api/channels/discover-models"

func discover(t *testing.T, s *Server, cookie *http.Cookie, c channel) *httptest.ResponseRecorder {
	t.Helper()
	b, e := json.Marshal(c)
	if e != nil {
		t.Fatal(e)
	}
	return call(s, "POST", discoveryPath, string(b), cookie)
}

func TestDiscoveryProtocolsAndReadOnly(t *testing.T) {
	for _, proto := range []string{"openai-chat", "openai-responses", "anthropic"} {
		t.Run(proto, func(t *testing.T) {
			s := testServer(t)
			cookie := loginCookie(t, s)
			calls := 0
			up := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
				calls++
				if r.Method != "GET" || r.URL.Path != "/v1/models" {
					t.Error(r.Method, r.URL)
				}
				if proto == "anthropic" {
					if r.Header.Get("x-api-key") != "secret" || r.Header.Get("anthropic-version") != "2023-06-01" || r.Header.Get("Authorization") != "" {
						t.Error("wrong anthropic headers")
					}
					if r.URL.Query().Get("after_id") == "" {
						fmt.Fprint(w, `{"data":[{"id":"z"}],"has_more":true,"last_id":"cursor/&?"}`)
						return
					}
					if r.URL.Query().Get("after_id") != "cursor/&?" {
						t.Error(r.URL)
					}
				} else if r.Header.Get("Authorization") != "Bearer secret" || r.Header.Get("x-api-key") != "" {
					t.Error("wrong openai headers")
				}
				fmt.Fprint(w, `{"data":[{"id":"z"},{"id":"a"},{"id":"a"}]}`)
			}))
			defer up.Close()
			draft := channel{BaseURL: up.URL + "/v1/", Protocol: proto, APIKey: "secret"}
			if proto == "anthropic" {
				draft.BaseURL = up.URL + "/"
			}
			w := discover(t, s, cookie, draft)
			if w.Code != 200 || strings.TrimSpace(w.Body.String()) != `{"models":["a","z"]}` {
				t.Fatal(w.Code, w.Body.String())
			}
			wantCalls := 1
			if proto == "anthropic" {
				wantCalls = 2
			}
			if calls != wantCalls {
				t.Fatal(calls)
			}
			cs, e := s.loadChannels()
			if e != nil || len(cs) != 0 {
				t.Fatal(cs, e)
			}
			draft.Name = "saved"
			draft.Models = map[string]string{"alias": "a"}
			b, _ := json.Marshal(draft)
			w = call(s, "POST", "/admin/api/channels", string(b), cookie)
			if w.Code != 200 {
				t.Fatal(w.Body.String())
			}
			var saved channel
			if e := json.Unmarshal(w.Body.Bytes(), &saved); e != nil {
				t.Fatal(e)
			}
			before, _ := s.loadChannels()
			draft.ID, draft.APIKey, draft.Name, draft.Models = saved.ID, "", "unsaved", nil
			if w = discover(t, s, cookie, draft); w.Code != 200 {
				t.Fatal(w.Code, w.Body.String())
			}
			after, _ := s.loadChannels()
			if !reflect.DeepEqual(before, after) {
				t.Fatal("discovery changed channel")
			}
			replacement := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
				if proto == "anthropic" {
					if r.Header.Get("x-api-key") != "replacement" {
						t.Error("draft key not used")
					}
				} else if r.Header.Get("Authorization") != "Bearer replacement" {
					t.Error("draft key not used")
				}
				fmt.Fprint(w, `{"data":[]}`)
			}))
			defer replacement.Close()
			draft.BaseURL, draft.APIKey = replacement.URL, "replacement"
			w = discover(t, s, cookie, draft)
			if w.Code != 200 || strings.TrimSpace(w.Body.String()) != `{"models":[]}` {
				t.Fatal(w.Code, w.Body.String())
			}
			after, _ = s.loadChannels()
			if !reflect.DeepEqual(before, after) {
				t.Fatal("draft replacement was saved")
			}
			var logs int
			if e := s.db.QueryRow("SELECT count(*) FROM logs").Scan(&logs); e != nil || logs != 0 {
				t.Fatal(logs, e)
			}
		})
	}
}

func TestDiscoveryAuthAndValidation(t *testing.T) {
	s := testServer(t)
	cookie := loginCookie(t, s)
	if w := call(s, "POST", discoveryPath, "{}", nil); w.Code != 401 {
		t.Fatal(w.Code)
	}
	for _, header := range []string{"Origin", "X-Requested-With"} {
		r := httptest.NewRequest("POST", discoveryPath, strings.NewReader("{}"))
		r.AddCookie(cookie)
		r.Header.Set("Origin", "http://example.com")
		r.Header.Set("X-Requested-With", "ai-router")
		r.Header.Del(header)
		w := httptest.NewRecorder()
		s.ServeHTTP(w, r)
		if w.Code != 403 {
			t.Fatal(header, w.Code)
		}
	}
	for _, method := range []string{"GET", "DELETE", "PUT"} {
		if w := call(s, method, discoveryPath, "{}", cookie); w.Code != 405 {
			t.Fatal(method, w.Code)
		}
	}
	for _, base := range []string{"", "ftp://example.com", "http://", "http://user:[email protected]", "http://example.com?x=1", "http://example.com?", "http://example.com#"} {
		if w := discover(t, s, cookie, channel{BaseURL: base, Protocol: "openai-chat"}); w.Code != 400 {
			t.Fatal(base, w.Code)
		}
	}
	if w := discover(t, s, cookie, channel{BaseURL: "http://example.com", Protocol: "unknown"}); w.Code != 400 {
		t.Fatal(w.Code)
	}
	if w := discover(t, s, cookie, channel{ID: "missing", BaseURL: "http://example.com", Protocol: "anthropic"}); w.Code != 404 {
		t.Fatal(w.Code)
	}
	for _, body := range []string{"{", "{} {}", strings.Repeat(" ", maxBody+1)} {
		if w := call(s, "POST", discoveryPath, body, cookie); w.Code != 400 {
			t.Fatal(w.Code)
		}
	}
}

func TestDiscoveryUpstreamErrors(t *testing.T) {
	for _, tc := range []struct {
		name, body    string
		status, calls int
	}{
		{"status", "secret upstream details", 401, 1},
		{"json", "secret", 200, 1},
		{"missing data", "{}", 200, 1},
		{"null", `{"data":null}`, 200, 1},
		{"empty id", `{"data":[{"id":""}]}`, 200, 1},
		{"trailing", `{"data":[]}{}`, 200, 1},
		{"oversized", strings.Repeat("x", discoveryMaxBody+1), 200, 1},
		{"missing cursor", `{"data":[],"has_more":true}`, 200, 1},
		{"repeated cursor", `{"data":[],"has_more":true,"last_id":"same"}`, 200, 2},
		{"page limit", "", 200, discoveryMaxPages},
	} {
		t.Run(tc.name, func(t *testing.T) {
			s := testServer(t)
			cookie := loginCookie(t, s)
			calls := 0
			up := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
				calls++
				w.WriteHeader(tc.status)
				if tc.name == "page limit" {
					fmt.Fprintf(w, `{"data":[],"has_more":true,"last_id":"%d"}`, calls)
				} else {
					fmt.Fprint(w, tc.body)
				}
			}))
			defer up.Close()
			w := discover(t, s, cookie, channel{BaseURL: up.URL, Protocol: "anthropic", APIKey: "secret"})
			if w.Code != 502 || strings.TrimSpace(w.Body.String()) != `{"error":"model discovery failed"}` || calls != tc.calls {
				t.Fatal(w.Code, w.Body.String(), calls)
			}
		})
	}
}

func TestDiscoveryRedirectAndDeadline(t *testing.T) {
	s := testServer(t)
	cookie := loginCookie(t, s)
	target := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { t.Error("redirect followed") }))
	defer target.Close()
	up := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { http.Redirect(w, r, target.URL, 307) }))
	defer up.Close()
	for _, proto := range []string{"openai-chat", "anthropic"} {
		if w := discover(t, s, cookie, channel{BaseURL: up.URL, Protocol: proto, APIKey: "secret"}); w.Code != 502 {
			t.Fatal(w.Code)
		}
	}
	slow := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { <-r.Context().Done() }))
	defer slow.Close()
	ctx, cancel := context.WithTimeout(context.Background(), 20*time.Millisecond)
	defer cancel()
	b, _ := json.Marshal(channel{BaseURL: slow.URL, Protocol: "openai-chat"})
	r := httptest.NewRequest("POST", discoveryPath, strings.NewReader(string(b))).WithContext(ctx)
	r.AddCookie(cookie)
	r.Header.Set("Origin", "http://example.com")
	r.Header.Set("X-Requested-With", "ai-router")
	w := httptest.NewRecorder()
	s.ServeHTTP(w, r)
	if w.Code != 504 {
		t.Fatal(w.Code, w.Body.String())
	}
}

func TestModelsListSelectedAliasesAndPermissions(t *testing.T) {
	s := testServer(t)
	cookie := loginCookie(t, s)
	var allowed string
	for _, name := range []string{"selected", "forbidden", "disabled"} {
		b, _ := json.Marshal(channel{Name: name, BaseURL: "http://example.com", Protocol: "openai-chat", Enabled: name != "disabled", Models: map[string]string{"z": "vendor/real", "alias": "vendor/real"}})
		w := call(s, "POST", "/admin/api/channels", string(b), cookie)
		if w.Code != 200 {
			t.Fatal(w.Body.String())
		}
		var c channel
		json.Unmarshal(w.Body.Bytes(), &c)
		if name != "forbidden" {
			if allowed != "" {
				allowed += ","
			}
			allowed += fmt.Sprintf("%q", c.ID)
		}
	}
	w := call(s, "POST", "/admin/api/keys", `{"name":"client","channels":[`+allowed+`]}`, cookie)
	if w.Code != 200 {
		t.Fatal(w.Body.String())
	}
	var key map[string]string
	json.Unmarshal(w.Body.Bytes(), &key)
	r := httptest.NewRequest("GET", "/v1/models", nil)
	r.Header.Set("Authorization", "Bearer "+key["key"])
	w = httptest.NewRecorder()
	s.ServeHTTP(w, r)
	var result struct{ Data []struct{ ID string } }
	if e := json.Unmarshal(w.Body.Bytes(), &result); e != nil || w.Code != 200 {
		t.Fatal(w.Code, e)
	}
	if len(result.Data) != 2 || result.Data[0].ID != "selected/alias" || result.Data[1].ID != "selected/z" {
		t.Fatal(w.Body.String())
	}
	for _, model := range []string{"forbidden/alias", "disabled/alias"} {
		r = httptest.NewRequest("POST", "/v1/chat/completions", strings.NewReader(fmt.Sprintf(`{"model":%q,"messages":[]}`, model)))
		r.Header.Set("Authorization", "Bearer "+key["key"])
		w = httptest.NewRecorder()
		s.ServeHTTP(w, r)
		if w.Code != 403 {
			t.Fatal(model, w.Code)
		}
	}
}