268 lines
10 KiB
Go
268 lines
10 KiB
Go
package server
|
|
|
|
import (
|
|
"ai-router/internal/protocol"
|
|
"context"
|
|
"encoding/json"
|
|
"fmt"
|
|
"io"
|
|
"net/http"
|
|
"net/http/httptest"
|
|
"strings"
|
|
"testing"
|
|
"time"
|
|
)
|
|
|
|
const probePath = "/admin/api/models/test"
|
|
const probeChat = `{"id":"c","object":"chat.completion","choices":[{"message":{"role":"assistant","content":"OK"},"finish_reason":"stop"}]}`
|
|
|
|
func saveProbeChannel(t *testing.T, s *Server, cookie *http.Cookie, url, kind string, enabled bool) {
|
|
t.Helper()
|
|
b, _ := json.Marshal(channel{Name: "saved", BaseURL: url, Protocol: kind, Enabled: enabled, Models: map[string]string{"alias": "actual/model"}, APIKey: "secret-key"})
|
|
if w := call(s, "POST", "/admin/api/channels", string(b), cookie); w.Code != 200 {
|
|
t.Fatal(w.Body.String())
|
|
}
|
|
}
|
|
|
|
func probeResult(t *testing.T, w *httptest.ResponseRecorder) modelTestResult {
|
|
t.Helper()
|
|
var result modelTestResult
|
|
if w.Code != 200 || json.Unmarshal(w.Body.Bytes(), &result) != nil {
|
|
t.Fatal(w.Code, w.Body.String())
|
|
}
|
|
if result.DurationMS < 0 || strings.Contains(w.Body.String(), "secret-key") {
|
|
t.Fatal(w.Body.String())
|
|
}
|
|
return result
|
|
}
|
|
|
|
func TestModelProbeProtocols(t *testing.T) {
|
|
for _, tc := range []struct{ kind, path, limit, response string }{
|
|
{"openai-chat", "/chat/completions", "max_completion_tokens", probeChat},
|
|
{"openai-responses", "/responses", "max_output_tokens", `{"id":"r","object":"response","status":"completed","output":[{"type":"message","role":"assistant","content":[{"type":"output_text","text":"OK"}]}]}`},
|
|
{"anthropic", "/v1/messages", "max_tokens", `{"id":"m","type":"message","role":"assistant","stop_reason":"end_turn","content":[{"type":"text","text":"OK"}]}`},
|
|
} {
|
|
t.Run(tc.kind, func(t *testing.T) {
|
|
s := testServer(t)
|
|
cookie := loginCookie(t, s)
|
|
expectedPrompt := "Reply with OK."
|
|
up := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
|
b, _ := io.ReadAll(r.Body)
|
|
var obj map[string]any
|
|
if json.Unmarshal(b, &obj) != nil || obj["model"] != "actual/model" || obj[tc.limit] != float64(128) || obj["stream"] != false || !strings.Contains(string(b), expectedPrompt) {
|
|
t.Error(string(b))
|
|
}
|
|
if tc.kind == "openai-chat" && obj["max_tokens"] != nil {
|
|
t.Error("legacy token limit")
|
|
}
|
|
if r.Method != "POST" || r.URL.Path != tc.path {
|
|
t.Error(r.Method, r.URL.Path)
|
|
}
|
|
if tc.kind == "anthropic" {
|
|
if r.Header.Get("x-api-key") != "secret-key" || r.Header.Get("anthropic-version") != "2023-06-01" {
|
|
t.Error("headers")
|
|
}
|
|
} else if r.Header.Get("Authorization") != "Bearer secret-key" {
|
|
t.Error("authorization")
|
|
}
|
|
fmt.Fprint(w, tc.response)
|
|
}))
|
|
defer up.Close()
|
|
saveProbeChannel(t, s, cookie, up.URL, tc.kind, true)
|
|
for _, prompt := range []string{"", strings.Repeat("中", 2000)} {
|
|
input := map[string]string{"model": "saved/alias"}
|
|
if prompt != "" {
|
|
expectedPrompt = prompt
|
|
input["prompt"] = prompt
|
|
}
|
|
b, _ := json.Marshal(input)
|
|
result := probeResult(t, call(s, "POST", probePath, string(b), cookie))
|
|
if !result.OK || result.Status != 200 || result.Error != "" {
|
|
t.Fatal(result)
|
|
}
|
|
}
|
|
for _, table := range []string{"keys", "logs"} {
|
|
var n int
|
|
if e := s.db.QueryRow("SELECT count(*) FROM " + table).Scan(&n); e != nil || n != 0 {
|
|
t.Fatal(table, n, e)
|
|
}
|
|
}
|
|
if len(s.logQueue) != 0 {
|
|
t.Fatal("queued formal log")
|
|
}
|
|
})
|
|
}
|
|
}
|
|
|
|
func TestModelProbeGuardAndInput(t *testing.T) {
|
|
s := testServer(t)
|
|
cookie := loginCookie(t, s)
|
|
for _, prompt := range []string{"", " ", "\t\r\n", "\u3000\u00a0"} {
|
|
body, _ := json.Marshal(map[string]string{"model": "saved/alias", "prompt": prompt})
|
|
if w := call(s, "POST", probePath, string(body), cookie); w.Code != 400 {
|
|
t.Fatalf("blank prompt %q: status %d", prompt, w.Code)
|
|
}
|
|
}
|
|
up := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { t.Error("unexpected upstream request") }))
|
|
defer up.Close()
|
|
saveProbeChannel(t, s, cookie, up.URL, "openai-chat", false)
|
|
if w := call(s, "POST", probePath, `{"model":"saved/alias"}`, nil); w.Code != 401 {
|
|
t.Fatal(w.Code)
|
|
}
|
|
for _, header := range []string{"Origin", "X-Requested-With"} {
|
|
r := httptest.NewRequest("POST", probePath, strings.NewReader(`{"model":"saved/alias"}`))
|
|
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(w.Code)
|
|
}
|
|
}
|
|
for _, body := range []string{`{"model":"saved/alias"}`, `{"model":"missing/alias"}`} {
|
|
if w := call(s, "POST", probePath, body, cookie); w.Code != 403 {
|
|
t.Fatal(w.Code)
|
|
}
|
|
}
|
|
if _, e := s.db.Exec("UPDATE channels SET enabled=1"); e != nil {
|
|
t.Fatal(e)
|
|
}
|
|
if w := call(s, "POST", probePath, `{"model":"saved/unmapped"}`, cookie); w.Code != 403 {
|
|
t.Fatal(w.Code)
|
|
}
|
|
for _, body := range []string{`{}`, `null`, `{"model":"alias"}`, `{"model":"/alias"}`, `{"model":"saved/"}`, `{"model":"saved/alias","prompt":2}`, `{"model":"saved/alias"} {}`, `{"model":"saved/alias","prompt":"` + strings.Repeat("中", 2001) + `"}`} {
|
|
if w := call(s, "POST", probePath, body, cookie); w.Code != 400 {
|
|
t.Fatal(body, w.Code)
|
|
}
|
|
}
|
|
if w := call(s, "GET", probePath, "", cookie); w.Code != 405 {
|
|
t.Fatal(w.Code)
|
|
}
|
|
}
|
|
|
|
func TestModelProbeFailures(t *testing.T) {
|
|
for _, tc := range []struct {
|
|
name, body string
|
|
status int
|
|
error string
|
|
}{
|
|
{"rejected", "secret-key upstream body", 401, "upstream rejected request"},
|
|
{"redirect", "secret-key", 302, "upstream rejected request"},
|
|
{"html", "<html>secret-key</html>", 200, "invalid upstream response"},
|
|
{"empty", "{}", 200, "invalid upstream response"},
|
|
{"oversize", strings.Repeat("x", modelTestMaxBody+1), 200, "upstream response too large"},
|
|
{"limit", probeChat + strings.Repeat(" ", modelTestMaxBody-len(probeChat)), 200, ""},
|
|
} {
|
|
t.Run(tc.name, func(t *testing.T) {
|
|
s := testServer(t)
|
|
cookie := loginCookie(t, s)
|
|
up := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
|
if r.URL.Path != "/chat/completions" {
|
|
t.Error("followed redirect")
|
|
}
|
|
w.Header().Set("Location", "/redirected")
|
|
w.WriteHeader(tc.status)
|
|
fmt.Fprint(w, tc.body)
|
|
}))
|
|
defer up.Close()
|
|
saveProbeChannel(t, s, cookie, up.URL, "openai-chat", true)
|
|
result := probeResult(t, call(s, "POST", probePath, `{"model":"saved/alias"}`, cookie))
|
|
if result.Status != tc.status || result.Error != tc.error || result.OK != (tc.error == "") {
|
|
t.Fatal(result)
|
|
}
|
|
})
|
|
}
|
|
}
|
|
|
|
func TestModelProbeResponseValidation(t *testing.T) {
|
|
if !validModelTestResponse(protocol.Anthropic, []byte(`{"id":"m","type":"message","role":"assistant","stop_reason":"max_tokens","content":[{"type":"thinking","thinking":"reasoning","signature":"sig"}]}`)) {
|
|
t.Fatal("valid thinking-only response rejected")
|
|
}
|
|
for _, tc := range []struct {
|
|
kind protocol.Kind
|
|
body string
|
|
ok bool
|
|
}{
|
|
{protocol.Chat, `{"id":"c","object":"chat.completion","choices":[{"message":{"role":"assistant","content":null},"finish_reason":"length"}]}`, true},
|
|
{protocol.Responses, `{"id":"r","object":"response","status":"incomplete","incomplete_details":{"reason":"max_output_tokens"},"output":[{"type":"reasoning","summary":[]}]}`, true},
|
|
{protocol.Anthropic, `{"id":"m","type":"message","role":"assistant","stop_reason":"max_tokens","content":[]}`, true},
|
|
{protocol.Chat, `{"id":"c","object":"chat.completion","choices":[{}]}`, false},
|
|
{protocol.Responses, `{"id":"r","object":"response","status":"failed","output":[]}`, false},
|
|
{protocol.Responses, `{"id":"r","object":"response","status":"completed","output":[{}]}`, false},
|
|
{protocol.Anthropic, `{"id":"m","type":"message","role":"assistant","stop_reason":"end_turn","content":[{}]}`, false},
|
|
} {
|
|
if got := validModelTestResponse(tc.kind, []byte(tc.body)); got != tc.ok {
|
|
t.Error(tc.kind, tc.body, got)
|
|
}
|
|
}
|
|
for _, kind := range []protocol.Kind{protocol.Chat, protocol.Responses, protocol.Anthropic} {
|
|
for _, body := range []string{"{}", "null", "[]", probeChat + "{}", `{"error":{"message":"secret-key"}}`} {
|
|
if validModelTestResponse(kind, []byte(body)) {
|
|
t.Error(kind, body)
|
|
}
|
|
}
|
|
}
|
|
}
|
|
|
|
type probeRoundTripper func(*http.Request) (*http.Response, error)
|
|
|
|
func (f probeRoundTripper) RoundTrip(r *http.Request) (*http.Response, error) { return f(r) }
|
|
|
|
func TestModelProbeTimeoutAndCancellation(t *testing.T) {
|
|
for _, cancelRequest := range []bool{false, true} {
|
|
t.Run(fmt.Sprint(cancelRequest), func(t *testing.T) {
|
|
s := testServer(t)
|
|
cookie := loginCookie(t, s)
|
|
up := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {}))
|
|
defer up.Close()
|
|
saveProbeChannel(t, s, cookie, up.URL, "openai-chat", true)
|
|
ctx, cancel := context.WithCancel(context.Background())
|
|
defer cancel()
|
|
s.client.Transport = probeRoundTripper(func(r *http.Request) (*http.Response, error) {
|
|
deadline, ok := r.Context().Deadline()
|
|
if !ok || time.Until(deadline) > 30*time.Second || time.Until(deadline) < 29*time.Second {
|
|
t.Error("missing 30s deadline", deadline)
|
|
}
|
|
if cancelRequest {
|
|
cancel()
|
|
<-r.Context().Done()
|
|
return nil, r.Context().Err()
|
|
}
|
|
return nil, context.DeadlineExceeded
|
|
})
|
|
r := httptest.NewRequest("POST", probePath, strings.NewReader(`{"model":"saved/alias"}`)).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)
|
|
result := probeResult(t, w)
|
|
if result.OK || result.Status != 0 || result.Error == "" {
|
|
t.Fatal(result)
|
|
}
|
|
})
|
|
}
|
|
s := testServer(t)
|
|
canceled := make(chan struct{})
|
|
up := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
|
w.WriteHeader(200)
|
|
w.(http.Flusher).Flush()
|
|
<-r.Context().Done()
|
|
close(canceled)
|
|
}))
|
|
defer up.Close()
|
|
ctx, cancel := context.WithTimeout(context.Background(), 50*time.Millisecond)
|
|
defer cancel()
|
|
result := s.probeModel(ctx, channel{BaseURL: up.URL, Protocol: "openai-chat"}, "actual", "test")
|
|
if result.OK || result.Status != 200 || result.Error != "model test timed out" {
|
|
t.Fatal(result)
|
|
}
|
|
select {
|
|
case <-canceled:
|
|
case <-time.After(time.Second):
|
|
t.Fatal("upstream was not canceled")
|
|
}
|
|
}
|