444 lines
17 KiB
Go
444 lines
17 KiB
Go
package server
|
|
|
|
import (
|
|
"context"
|
|
"encoding/json"
|
|
"fmt"
|
|
"io"
|
|
"net/http"
|
|
"net/http/httptest"
|
|
"strings"
|
|
"sync"
|
|
"sync/atomic"
|
|
"testing"
|
|
"time"
|
|
)
|
|
|
|
type integrationKey struct{ raw, id string }
|
|
|
|
type integrationStreamRecorder struct {
|
|
*httptest.ResponseRecorder
|
|
flushed chan struct{}
|
|
}
|
|
|
|
func (w *integrationStreamRecorder) Flush() {
|
|
w.ResponseRecorder.Flush()
|
|
select {
|
|
case w.flushed <- struct{}{}:
|
|
default:
|
|
}
|
|
}
|
|
|
|
func integrationAdmin(t *testing.T, s *Server, cookie *http.Cookie, method, path string, body any, out any) {
|
|
t.Helper()
|
|
b, err := json.Marshal(body)
|
|
if err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
w := call(s, method, path, string(b), cookie)
|
|
if w.Code != http.StatusOK {
|
|
t.Fatalf("%s %s: %d %s", method, path, w.Code, w.Body.String())
|
|
}
|
|
if out != nil {
|
|
if err := json.Unmarshal(w.Body.Bytes(), out); err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
}
|
|
}
|
|
|
|
func integrationChannel(t *testing.T, s *Server, cookie *http.Cookie, name, kind, base string) channel {
|
|
t.Helper()
|
|
if kind != "anthropic" {
|
|
base += "/v1"
|
|
}
|
|
var ch channel
|
|
integrationAdmin(t, s, cookie, "POST", "/admin/api/channels", channel{
|
|
Name: name, Protocol: kind, BaseURL: base, Enabled: true,
|
|
APIKey: "upstream-secret", Models: map[string]string{"gpt-5": "gpt-4", "gpt-4": "abc/gpt-4"},
|
|
}, &ch)
|
|
if ch.ID == "" || ch.APIKey != "" {
|
|
t.Fatalf("invalid channel response: %+v", ch)
|
|
}
|
|
return ch
|
|
}
|
|
|
|
func integrationCreateKey(t *testing.T, s *Server, cookie *http.Cookie, name string, ch channel) integrationKey {
|
|
t.Helper()
|
|
var created struct {
|
|
Key string `json:"key"`
|
|
}
|
|
integrationAdmin(t, s, cookie, "POST", "/admin/api/keys", apiKey{Name: name, Channels: []string{ch.ID}}, &created)
|
|
var keys []apiKey
|
|
integrationAdmin(t, s, cookie, "GET", "/admin/api/keys", nil, &keys)
|
|
for _, k := range keys {
|
|
if k.Name == name && created.Key != "" {
|
|
return integrationKey{created.Key, k.ID}
|
|
}
|
|
}
|
|
t.Fatal("created key missing from management API")
|
|
return integrationKey{}
|
|
}
|
|
|
|
func integrationRequest(kind, model string, stream bool) string {
|
|
if kind == "openai-responses" {
|
|
return fmt.Sprintf(`{"model":%q,"input":"hello","max_output_tokens":32,"stream":%t}`, model, stream)
|
|
}
|
|
return fmt.Sprintf(`{"model":%q,"messages":[{"role":"user","content":"hello"}],"max_tokens":32,"stream":%t}`, model, stream)
|
|
}
|
|
|
|
func integrationPath(kind string) string {
|
|
return map[string]string{"openai-chat": "/v1/chat/completions", "openai-responses": "/v1/responses", "anthropic": "/v1/messages"}[kind]
|
|
}
|
|
|
|
func integrationProxy(s *Server, key integrationKey, kind, model string, stream bool) *httptest.ResponseRecorder {
|
|
r := httptest.NewRequest("POST", integrationPath(kind), strings.NewReader(integrationRequest(kind, model, stream)))
|
|
if kind == "anthropic" {
|
|
r.Header.Set("x-api-key", key.raw)
|
|
} else {
|
|
r.Header.Set("Authorization", "Bearer "+key.raw)
|
|
}
|
|
w := httptest.NewRecorder()
|
|
s.ServeHTTP(w, r)
|
|
return w
|
|
}
|
|
|
|
func integrationResponse(kind string) string {
|
|
switch kind {
|
|
case "openai-chat":
|
|
return `{"id":"chat-test","object":"chat.completion","model":"gpt-4","choices":[{"index":0,"message":{"role":"assistant","content":"hello"},"finish_reason":"stop"}],"usage":{"prompt_tokens":12,"completion_tokens":3,"prompt_tokens_details":{"cached_tokens":4}}}`
|
|
case "openai-responses":
|
|
return `{"id":"resp-test","object":"response","model":"gpt-4","status":"completed","output":[{"type":"message","id":"msg-test","role":"assistant","status":"completed","content":[{"type":"output_text","text":"hello","annotations":[]}]}],"usage":{"input_tokens":12,"output_tokens":3,"input_tokens_details":{"cached_tokens":4}}}`
|
|
default:
|
|
return `{"id":"msg-test","type":"message","role":"assistant","model":"gpt-4","content":[{"type":"text","text":"hello"}],"stop_reason":"end_turn","stop_sequence":null,"usage":{"input_tokens":6,"output_tokens":3,"cache_read_input_tokens":4,"cache_creation_input_tokens":2}}`
|
|
}
|
|
}
|
|
|
|
func integrationSSE(kind string) string {
|
|
switch kind {
|
|
case "openai-chat":
|
|
return "data: " + `{"id":"chat-test","choices":[{"index":0,"delta":{"role":"assistant","content":"hello"},"finish_reason":null}]}` + "\n\ndata: " + `{"choices":[{"index":0,"delta":{},"finish_reason":"stop"}],"usage":{"prompt_tokens":12,"completion_tokens":3,"prompt_tokens_details":{"cached_tokens":4}}}` + "\n\ndata: [DONE]\n\n"
|
|
case "openai-responses":
|
|
return "event: response.created\ndata: " + `{"type":"response.created","response":{"id":"resp-test","status":"in_progress"}}` + "\n\n" +
|
|
"event: response.output_item.added\ndata: " + `{"type":"response.output_item.added","output_index":0,"item":{"type":"message","role":"assistant","content":[]}}` + "\n\n" +
|
|
"event: response.output_text.delta\ndata: " + `{"type":"response.output_text.delta","output_index":0,"content_index":0,"delta":"hello"}` + "\n\n" +
|
|
"event: response.completed\ndata: " + `{"type":"response.completed","response":` + integrationResponse(kind) + "}\n\n"
|
|
default:
|
|
return "event: message_start\ndata: " + `{"type":"message_start","message":{"id":"msg-test","role":"assistant","content":[],"usage":{"input_tokens":6,"output_tokens":0,"cache_read_input_tokens":4,"cache_creation_input_tokens":2}}}` + "\n\n" +
|
|
"event: content_block_start\ndata: " + `{"type":"content_block_start","index":0,"content_block":{"type":"text","text":""}}` + "\n\n" +
|
|
"event: content_block_delta\ndata: " + `{"type":"content_block_delta","index":0,"delta":{"type":"text_delta","text":"hello"}}` + "\n\n" +
|
|
"event: content_block_stop\ndata: " + `{"type":"content_block_stop","index":0}` + "\n\n" +
|
|
"event: message_delta\ndata: " + `{"type":"message_delta","delta":{"stop_reason":"end_turn","stop_sequence":null},"usage":{"output_tokens":3}}` + "\n\n" +
|
|
"event: message_stop\ndata: " + `{"type":"message_stop"}` + "\n\n"
|
|
}
|
|
}
|
|
|
|
func integrationLogs(t *testing.T, s *Server, c *http.Cookie, key integrationKey) []requestLog {
|
|
t.Helper()
|
|
if err := s.flushLogs(); err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
var logs []requestLog
|
|
integrationAdmin(t, s, c, "GET", "/admin/api/logs?key_id="+key.id, nil, &logs)
|
|
for _, l := range logs {
|
|
if l.KeyID != key.id {
|
|
t.Fatalf("cross-key log leak: %+v", l)
|
|
}
|
|
}
|
|
return logs
|
|
}
|
|
|
|
func integrationCount(t *testing.T, name string, got *int64, want int64) {
|
|
t.Helper()
|
|
if got == nil || *got != want {
|
|
t.Fatalf("%s: got %v, want %d", name, got, want)
|
|
}
|
|
}
|
|
|
|
func TestIntegrationProtocolMatrix(t *testing.T) {
|
|
kinds := []string{"openai-chat", "openai-responses", "anthropic"}
|
|
for _, upKind := range kinds {
|
|
for _, downKind := range kinds {
|
|
t.Run(downKind+"_to_"+upKind, func(t *testing.T) {
|
|
s := testServer(t)
|
|
cookie := loginCookie(t, s)
|
|
var calls atomic.Int64
|
|
var expectedModel atomic.Value
|
|
up := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
|
calls.Add(1)
|
|
if r.URL.Path != integrationPath(upKind) {
|
|
t.Errorf("upstream path: %s", r.URL.Path)
|
|
}
|
|
if upKind == "anthropic" {
|
|
if r.Header.Get("x-api-key") != "upstream-secret" || r.Header.Get("anthropic-version") == "" {
|
|
t.Error("missing Anthropic credentials/version")
|
|
}
|
|
} else if r.Header.Get("Authorization") != "Bearer upstream-secret" {
|
|
t.Error("wrong upstream credentials")
|
|
}
|
|
var body struct {
|
|
Model string `json:"model"`
|
|
Stream bool `json:"stream"`
|
|
}
|
|
b, err := io.ReadAll(r.Body)
|
|
if err != nil || json.Unmarshal(b, &body) != nil {
|
|
t.Error("invalid upstream request")
|
|
w.WriteHeader(400)
|
|
return
|
|
}
|
|
if body.Model != expectedModel.Load() || !strings.Contains(string(b), "hello") {
|
|
t.Errorf("upstream request: %s; want model %s", b, expectedModel.Load())
|
|
}
|
|
if body.Stream {
|
|
w.Header().Set("Content-Type", "text/event-stream")
|
|
for _, frame := range strings.Split(integrationSSE(upKind), "\n\n") {
|
|
if frame != "" {
|
|
fmt.Fprint(w, frame+"\n\n")
|
|
w.(http.Flusher).Flush()
|
|
}
|
|
}
|
|
} else {
|
|
w.Header().Set("Content-Type", "application/json")
|
|
fmt.Fprint(w, integrationResponse(upKind))
|
|
}
|
|
}))
|
|
defer up.Close()
|
|
abc := integrationChannel(t, s, cookie, "abc", upKind, up.URL)
|
|
other := integrationChannel(t, s, cookie, "other", upKind, up.URL)
|
|
keys := []integrationKey{integrationCreateKey(t, s, cookie, "first", abc), integrationCreateKey(t, s, cookie, "second", other)}
|
|
for i, key := range keys {
|
|
name, forbidden := "abc", "other"
|
|
if i == 1 {
|
|
name, forbidden = forbidden, name
|
|
}
|
|
for _, stream := range []bool{false, true} {
|
|
for _, alias := range []string{"gpt-5", "gpt-4"} {
|
|
expectedModel.Store("gpt-4")
|
|
if alias == "gpt-4" {
|
|
expectedModel.Store("abc/gpt-4")
|
|
}
|
|
w := integrationProxy(s, key, downKind, name+"/"+alias, stream)
|
|
if w.Code != 200 || !strings.Contains(w.Body.String(), "hello") {
|
|
t.Fatalf("stream=%t model=%s: %d %s", stream, alias, w.Code, w.Body.String())
|
|
}
|
|
if stream {
|
|
terminal := map[string]string{"openai-chat": "data: [DONE]", "openai-responses": "event: response.completed", "anthropic": "event: message_stop"}[downKind]
|
|
if !w.Flushed || !strings.Contains(w.Header().Get("Content-Type"), "text/event-stream") || !strings.Contains(w.Body.String(), terminal) {
|
|
t.Fatalf("incomplete SSE: %s", w.Body.String())
|
|
}
|
|
} else {
|
|
var body map[string]any
|
|
if err := json.Unmarshal(w.Body.Bytes(), &body); err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
field := map[string]string{"openai-chat": "choices", "openai-responses": "output", "anthropic": "content"}[downKind]
|
|
if body[field] == nil {
|
|
t.Fatalf("wrong downstream shape: %s", w.Body.String())
|
|
}
|
|
}
|
|
}
|
|
}
|
|
before := calls.Load()
|
|
if w := integrationProxy(s, key, downKind, forbidden+"/gpt-5", false); w.Code != 403 {
|
|
t.Fatalf("permission: %d %s", w.Code, w.Body.String())
|
|
}
|
|
if calls.Load() != before {
|
|
t.Fatal("forbidden request reached upstream")
|
|
}
|
|
}
|
|
if calls.Load() != 8 {
|
|
t.Fatalf("upstream calls: %d", calls.Load())
|
|
}
|
|
for _, key := range keys {
|
|
logs := integrationLogs(t, s, cookie, key)
|
|
if len(logs) != 5 {
|
|
t.Fatalf("logs: %+v", logs)
|
|
}
|
|
for _, l := range logs {
|
|
if l.Status == 403 {
|
|
if l.UsageKnown || l.Input != nil {
|
|
t.Fatal("denied request has usage")
|
|
}
|
|
continue
|
|
}
|
|
actual := "gpt-4"
|
|
if strings.HasSuffix(l.RequestModel, "/gpt-4") {
|
|
actual = "abc/gpt-4"
|
|
}
|
|
if l.Status != 200 || l.Error != "" || !l.UsageKnown || l.ActualModel != actual || l.Downstream != downKind || l.Upstream != upKind {
|
|
t.Fatalf("wrong log: %+v", l)
|
|
}
|
|
integrationCount(t, "input", l.Input, 12)
|
|
integrationCount(t, "output", l.Output, 3)
|
|
integrationCount(t, "cache read", l.CacheRead, 4)
|
|
write := int64(0)
|
|
if upKind == "anthropic" {
|
|
write = 2
|
|
}
|
|
integrationCount(t, "cache write", l.CacheWrite, write)
|
|
}
|
|
var stats struct {
|
|
Summary summary `json:"summary"`
|
|
ByKey []struct {
|
|
KeyID string `json:"key_id"`
|
|
} `json:"by_key"`
|
|
}
|
|
integrationAdmin(t, s, cookie, "GET", "/admin/api/stats?key_id="+key.id, nil, &stats)
|
|
if stats.Summary.Requests != 5 || stats.Summary.Success != 4 || stats.Summary.Failed != 1 || stats.Summary.Unknown != 1 || len(stats.ByKey) != 1 || stats.ByKey[0].KeyID != key.id {
|
|
t.Fatalf("wrong isolated stats: %+v", stats)
|
|
}
|
|
integrationCount(t, "total input", stats.Summary.Input, 48)
|
|
integrationCount(t, "total output", stats.Summary.Output, 12)
|
|
integrationCount(t, "total cache read", stats.Summary.CacheRead, 16)
|
|
write := int64(0)
|
|
if upKind == "anthropic" {
|
|
write = 8
|
|
}
|
|
integrationCount(t, "total cache write", stats.Summary.CacheWrite, write)
|
|
}
|
|
})
|
|
}
|
|
}
|
|
}
|
|
|
|
func TestIntegrationStreamFailures(t *testing.T) {
|
|
for _, tc := range []struct {
|
|
name, body string
|
|
status int
|
|
}{
|
|
{"upstream_http_error", `{"error":"private upstream detail"}`, 429},
|
|
{"error_event", "event: error\ndata: {\"type\":\"error\",\"error\":{\"message\":\"private upstream detail\"}}\n\n", 200},
|
|
{"malformed_event", "data: {invalid}\n\n", 200},
|
|
{"truncated_stream", "data: {\"choices\":[{\"index\":0,\"delta\":{\"content\":\"partial\"},\"finish_reason\":null}]}\n\n", 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) {
|
|
w.Header().Set("Content-Type", "text/event-stream")
|
|
w.WriteHeader(tc.status)
|
|
fmt.Fprint(w, tc.body)
|
|
}))
|
|
defer up.Close()
|
|
ch := integrationChannel(t, s, cookie, "abc", "openai-chat", up.URL)
|
|
key := integrationCreateKey(t, s, cookie, "errors", ch)
|
|
w := integrationProxy(s, key, "openai-responses", "abc/gpt-5", true)
|
|
if w.Code != tc.status {
|
|
t.Fatalf("status: %d body: %s", w.Code, w.Body.String())
|
|
}
|
|
if strings.Contains(w.Body.String(), "response.completed") || strings.Contains(w.Body.String(), "private upstream detail") {
|
|
t.Fatalf("false success or leaked upstream detail: %s", w.Body.String())
|
|
}
|
|
logs := integrationLogs(t, s, cookie, key)
|
|
if len(logs) != 1 {
|
|
t.Fatalf("logs: %+v", logs)
|
|
}
|
|
l := logs[0]
|
|
if l.Status < 400 || l.Error == "" || l.UsageKnown || l.Input != nil || l.Output != nil {
|
|
t.Fatalf("failure not recorded correctly: %+v", l)
|
|
}
|
|
})
|
|
}
|
|
}
|
|
|
|
func TestAnthropicHeaderForwarding(t *testing.T) {
|
|
s := testServer(t)
|
|
cookie := loginCookie(t, s)
|
|
var mu sync.Mutex
|
|
got := map[string]string{}
|
|
up := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
|
mu.Lock()
|
|
got["version"], got["beta"] = r.Header.Get("anthropic-version"), r.Header.Get("anthropic-beta")
|
|
mu.Unlock()
|
|
w.Header().Set("Content-Type", "application/json")
|
|
fmt.Fprint(w, integrationResponse("anthropic"))
|
|
}))
|
|
defer up.Close()
|
|
ch := integrationChannel(t, s, cookie, "abc", "anthropic", up.URL)
|
|
key := integrationCreateKey(t, s, cookie, "beta", ch)
|
|
request := func(version, beta string) {
|
|
r := httptest.NewRequest("POST", "/v1/messages", strings.NewReader(integrationRequest("anthropic", "abc/gpt-5", false)))
|
|
r.Header.Set("x-api-key", key.raw)
|
|
if version != "" {
|
|
r.Header.Set("anthropic-version", version)
|
|
}
|
|
if beta != "" {
|
|
r.Header.Set("anthropic-beta", beta)
|
|
}
|
|
w := httptest.NewRecorder()
|
|
s.ServeHTTP(w, r)
|
|
if w.Code != 200 {
|
|
t.Fatalf("status %d: %s", w.Code, w.Body.String())
|
|
}
|
|
}
|
|
request("2023-06-01", "context-1m-2025-08-07,interleaved-thinking-2025-05-14")
|
|
mu.Lock()
|
|
if got["version"] != "2023-06-01" || got["beta"] != "context-1m-2025-08-07,interleaved-thinking-2025-05-14" {
|
|
t.Fatalf("valid headers dropped: %v", got)
|
|
}
|
|
mu.Unlock()
|
|
request("bad version; injected", `weird"value`)
|
|
mu.Lock()
|
|
if got["version"] != "2023-06-01" || got["beta"] != "" {
|
|
t.Fatalf("unsafe headers forwarded: %v", got)
|
|
}
|
|
mu.Unlock()
|
|
}
|
|
|
|
func TestIntegrationCancellationClosesUpstream(t *testing.T) {
|
|
s := testServer(t)
|
|
cookie := loginCookie(t, s)
|
|
started, stopped := make(chan struct{}), make(chan struct{})
|
|
release := make(chan struct{})
|
|
up := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
|
io.Copy(io.Discard, r.Body)
|
|
w.Header().Set("Content-Type", "text/event-stream")
|
|
fmt.Fprint(w, "data: {\"choices\":[{\"index\":0,\"delta\":{\"content\":\"partial\"},\"finish_reason\":null}]}\n\n")
|
|
w.(http.Flusher).Flush()
|
|
close(started)
|
|
select {
|
|
case <-r.Context().Done():
|
|
close(stopped)
|
|
case <-release:
|
|
}
|
|
}))
|
|
defer up.Close()
|
|
defer close(release)
|
|
ch := integrationChannel(t, s, cookie, "abc", "openai-chat", up.URL)
|
|
key := integrationCreateKey(t, s, cookie, "cancel", ch)
|
|
ctx, cancel := context.WithCancel(context.Background())
|
|
defer cancel()
|
|
r := httptest.NewRequest("POST", "/v1/responses", strings.NewReader(integrationRequest("openai-responses", "abc/gpt-5", true))).WithContext(ctx)
|
|
r.Header.Set("Authorization", "Bearer "+key.raw)
|
|
done := make(chan struct{})
|
|
w := &integrationStreamRecorder{ResponseRecorder: httptest.NewRecorder(), flushed: make(chan struct{}, 1)}
|
|
go func() { s.ServeHTTP(w, r); close(done) }()
|
|
select {
|
|
case <-started:
|
|
case <-time.After(3 * time.Second):
|
|
t.Fatal("upstream did not start")
|
|
}
|
|
select {
|
|
case <-w.flushed:
|
|
case <-time.After(3 * time.Second):
|
|
t.Fatal("downstream did not receive streaming output")
|
|
}
|
|
cancel()
|
|
select {
|
|
case <-done:
|
|
case <-time.After(3 * time.Second):
|
|
t.Fatal("proxy did not stop on cancellation")
|
|
}
|
|
select {
|
|
case <-stopped:
|
|
case <-time.After(3 * time.Second):
|
|
t.Fatal("upstream request was not cancelled")
|
|
}
|
|
logs := integrationLogs(t, s, cookie, key)
|
|
if len(logs) != 1 || logs[0].Status < 400 || logs[0].Error == "" || logs[0].UsageKnown || logs[0].Input != nil {
|
|
t.Fatalf("cancel log: %+v", logs)
|
|
}
|
|
if strings.Contains(w.Body.String(), "response.completed") {
|
|
t.Fatal("cancelled response marked completed")
|
|
}
|
|
}
|