m1saka/token_thief
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

fix: restore reviewable migration evidence

Browse Source
This commit is contained in:
MiMoCode
2026-07-10 18:26:48 +08:00
commit d1bbb5370c
42 changed files with 6419 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files
tests/proxy/error_test.go
+139
View File
@@ -0,0 +1,139 @@
package proxy_test
import (
"io"
"net/http"
"net/http/httptest"
"net/url"
"testing"
"time"
"git.misaka.ren/M1saka/token_thief/config"
"git.misaka.ren/M1saka/token_thief/logger"
"git.misaka.ren/M1saka/token_thief/proxy"
)
type sliceSubmitter struct{ entries []*logger.LogEntry }
func (s *sliceSubmitter) Submit(e *logger.LogEntry) { s.entries = append(s.entries, e) }
// TestUpstreamErrorRecorded 验证上游不可达时 502 响应被记录、错误信息进入 LogEntry.Error。
func TestUpstreamErrorRecorded(t *testing.T) {
// 指向一个一定不可用的端口
badURL, _ := url.Parse("http://127.0.0.1:1") // port 1 几乎肯定 connection refused
sub := &sliceSubmitter{}
filter, err := config.NewFilter(config.FilterDisabled, nil)
if err != nil {
t.Fatal(err)
}
h := proxy.New(badURL, filter, sub, 1024)
srv := httptest.NewServer(h)
defer srv.Close()
resp, err := http.Get(srv.URL + "/v1/chat/completions")
if err != nil {
t.Fatal(err)
}
defer resp.Body.Close()
body, _ := io.ReadAll(resp.Body)
if resp.StatusCode != http.StatusBadGateway {
t.Errorf("status=%d want 502, body=%q", resp.StatusCode, body)
}
if resp.Header.Get("X-Request-Id") == "" {
t.Errorf("missing X-Request-Id header")
}
deadline := time.Now().Add(2 * time.Second)
for len(sub.entries) == 0 && time.Now().Before(deadline) {
time.Sleep(10 * time.Millisecond)
}
if len(sub.entries) == 0 {
t.Fatal("no log entry captured")
}
e := sub.entries[0]
if e.StatusCode != http.StatusBadGateway {
t.Errorf("LogEntry.StatusCode=%d want 502", e.StatusCode)
}
if e.Error == "" {
t.Errorf("LogEntry.Error should be set, got empty")
}
if string(e.ResponseBody) != "bad gateway\n" {
t.Errorf("response_body should be fixed bad gateway, got %q", e.ResponseBody)
}
if e.RequestID == "" {
t.Errorf("LogEntry.RequestID empty")
}
}
// TestModifyResponseSetsRequestID 验证正常上游响应也会带上 X-Request-Id。
func TestModifyResponseSetsRequestID(t *testing.T) {
upstream := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
w.Header().Set("Content-Type", "application/json")
w.WriteHeader(200)
_, _ = w.Write([]byte(`{"ok":true}`))
}))
defer upstream.Close()
u, _ := url.Parse(upstream.URL)
sub := &sliceSubmitter{}
filter, err := config.NewFilter(config.FilterDisabled, nil)
if err != nil {
t.Fatal(err)
}
h := proxy.New(u, filter, sub, 1024)
srv := httptest.NewServer(h)
defer srv.Close()
resp, err := http.Get(srv.URL + "/v1/anything")
if err != nil {
t.Fatal(err)
}
defer resp.Body.Close()
rid := resp.Header.Get("X-Request-Id")
if len(rid) != 32 {
t.Errorf("X-Request-Id length=%d want 32, value=%q", len(rid), rid)
}
deadline := time.Now().Add(2 * time.Second)
for len(sub.entries) == 0 && time.Now().Before(deadline) {
time.Sleep(10 * time.Millisecond)
}
if len(sub.entries) == 0 {
t.Fatal("no log entry")
}
if sub.entries[0].RequestID != rid {
t.Errorf("LogEntry.RequestID=%q response header=%q (should match)", sub.entries[0].RequestID, rid)
}
}
func TestUpstreamTLSInsecureSkipVerifyAllowsSelfSigned(t *testing.T) {
upstream := httptest.NewTLSServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
w.WriteHeader(http.StatusOK)
_, _ = w.Write([]byte(`{"ok":true}`))
}))
defer upstream.Close()
u, _ := url.Parse(upstream.URL)
sub := &sliceSubmitter{}
filter, err := config.NewFilter(config.FilterDisabled, nil)
if err != nil {
t.Fatal(err)
}
h := proxy.NewWithOptions(u, filter, sub, 1024, proxy.Options{UpstreamTLSInsecureSkipVerify: true})
srv := httptest.NewServer(h)
defer srv.Close()
resp, err := http.Get(srv.URL + "/v1/anything")
if err != nil {
t.Fatal(err)
}
defer resp.Body.Close()
if resp.StatusCode != http.StatusOK {
body, _ := io.ReadAll(resp.Body)
t.Fatalf("status=%d want 200, body=%q", resp.StatusCode, body)
}
}
tests/proxy/robustness_test.go
+336
View File
@@ -0,0 +1,336 @@
package proxy_test
import (
"context"
"errors"
"io"
"net/http"
"net/http/httptest"
"net/netip"
"net/url"
"strings"
"sync/atomic"
"testing"
"time"
"git.misaka.ren/M1saka/token_thief/config"
"git.misaka.ren/M1saka/token_thief/proxy"
)
type failingBody struct{ err error }
func (b failingBody) Read([]byte) (int, error) { return 0, b.err }
func (failingBody) Close() error { return nil }
type lateFailingBody struct {
remaining int
err error
}
func (b *lateFailingBody) Read(p []byte) (int, error) {
if b.remaining == 0 {
return 0, b.err
}
n := min(len(p), b.remaining)
for i := range p[:n] {
p[i] = 'x'
}
b.remaining -= n
return n, nil
}
func (*lateFailingBody) Close() error { return nil }
type failingResponseWriter struct {
header http.Header
short bool
}
func (w *failingResponseWriter) Header() http.Header { return w.header }
func (*failingResponseWriter) WriteHeader(int) {}
func (w *failingResponseWriter) Write(p []byte) (int, error) {
if w.short {
return len(p) - 1, nil
}
return 0, errors.New("write failed")
}
func newTestHandler(t *testing.T, upstream *url.URL, sub *captureSubmitter, opts proxy.Options) *proxy.Handler {
t.Helper()
filter, err := config.NewFilter(config.FilterDisabled, nil)
if err != nil {
t.Fatal(err)
}
return proxy.NewWithOptions(upstream, filter, sub, 1024*1024, opts)
}
func TestRequestBodyReadFailureReturnsFixed400WithoutUpstream(t *testing.T) {
var upstreamCalls atomic.Int32
upstream := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
upstreamCalls.Add(1)
w.WriteHeader(http.StatusOK)
}))
defer upstream.Close()
u, _ := url.Parse(upstream.URL)
h := newTestHandler(t, u, &captureSubmitter{}, proxy.Options{})
req := httptest.NewRequest(http.MethodPost, "http://proxy.test/v1/chat", nil)
req.Body = failingBody{err: errors.New("secret read failure")}
req.ContentLength = 1
rec := httptest.NewRecorder()
h.ServeHTTP(rec, req)
if rec.Code != http.StatusBadRequest || rec.Body.String() != "bad request\n" {
t.Fatalf("response=(%d, %q), want fixed 400 bad request", rec.Code, rec.Body.String())
}
if upstreamCalls.Load() != 0 {
t.Fatalf("upstream called %d times", upstreamCalls.Load())
}
}
func TestRequestBodyReadFailureAfterCaptureLimitDoesNotReachUpstream(t *testing.T) {
var upstreamCalls atomic.Int32
upstream := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
upstreamCalls.Add(1)
_, _ = io.Copy(io.Discard, r.Body)
}))
defer upstream.Close()
u, _ := url.Parse(upstream.URL)
h := newTestHandler(t, u, &captureSubmitter{}, proxy.Options{})
req := httptest.NewRequest(http.MethodPost, "http://proxy.test/v1/chat", nil)
req.Body = &lateFailingBody{remaining: 1024*1024 + 1, err: errors.New("late read failure")}
req.ContentLength = 1024*1024 + 2
rec := httptest.NewRecorder()
h.ServeHTTP(rec, req)
if rec.Code != http.StatusBadRequest || rec.Body.String() != "bad request\n" {
t.Fatalf("response=(%d, %q), want fixed 400 bad request", rec.Code, rec.Body.String())
}
if upstreamCalls.Load() != 0 {
t.Fatalf("upstream called %d times", upstreamCalls.Load())
}
}
func TestFilteredRequestBodyReadFailureDoesNotReachUpstream(t *testing.T) {
var upstreamCalls atomic.Int32
upstream := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
upstreamCalls.Add(1)
}))
defer upstream.Close()
u, _ := url.Parse(upstream.URL)
filter, err := config.NewFilter(config.FilterBlacklist, []string{"/ignored"})
if err != nil {
t.Fatal(err)
}
h := proxy.NewWithOptions(u, filter, &captureSubmitter{}, 1024, proxy.Options{})
req := httptest.NewRequest(http.MethodPost, "http://proxy.test/ignored", nil)
req.Body = failingBody{err: errors.New("read failure")}
req.ContentLength = 1
rec := httptest.NewRecorder()
h.ServeHTTP(rec, req)
if rec.Code != http.StatusBadRequest || rec.Body.String() != "bad request\n" {
t.Fatalf("response=(%d, %q), want fixed 400 bad request", rec.Code, rec.Body.String())
}
if upstreamCalls.Load() != 0 {
t.Fatalf("upstream called %d times", upstreamCalls.Load())
}
}
func TestBadGatewayResponseDoesNotLeakUpstreamError(t *testing.T) {
badURL, _ := url.Parse("http://127.0.0.1:1")
h := newTestHandler(t, badURL, &captureSubmitter{}, proxy.Options{})
rec := httptest.NewRecorder()
h.ServeHTTP(rec, httptest.NewRequest(http.MethodGet, "http://proxy.test/v1/test", nil))
if rec.Code != http.StatusBadGateway || rec.Body.String() != "bad gateway\n" {
t.Fatalf("response=(%d, %q), want fixed 502 bad gateway", rec.Code, rec.Body.String())
}
}
func TestResponseWriteFailurePreventsCommit(t *testing.T) {
upstream := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
_, _ = io.WriteString(w, "response")
}))
defer upstream.Close()
u, _ := url.Parse(upstream.URL)
for _, short := range []bool{false, true} {
sub := &captureSubmitter{}
h := newTestHandler(t, u, sub, proxy.Options{})
w := &failingResponseWriter{header: make(http.Header), short: short}
h.ServeHTTP(w, httptest.NewRequest(http.MethodGet, "http://proxy.test/v1/test", nil))
if sub.Len() != 0 {
t.Fatalf("short=%v: failed response write was committed", short)
}
}
}
func TestTrustedProxyClientIP(t *testing.T) {
upstream := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
_, _ = io.WriteString(w, "ok")
}))
defer upstream.Close()
u, _ := url.Parse(upstream.URL)
trusted := netip.MustParsePrefix("10.0.0.0/8")
tests := []struct {
name string
peer string
xff string
wantIP string
}{
{name: "untrusted peer ignores xff", peer: "203.0.113.9:1234", xff: "198.51.100.1", wantIP: "203.0.113.9"},
{name: "strip trusted from right", peer: "10.0.0.2:1234", xff: "198.51.100.7, 10.0.0.3", wantIP: "198.51.100.7"},
{name: "ipv6", peer: "[2001:db8::2]:1234", xff: "198.51.100.7", wantIP: "2001:db8::2"},
}
for _, tc := range tests {
t.Run(tc.name, func(t *testing.T) {
sub := &captureSubmitter{}
h := newTestHandler(t, u, sub, proxy.Options{TrustedProxies: []netip.Prefix{trusted}})
req := httptest.NewRequest(http.MethodGet, "http://proxy.test/v1/test", nil)
req.RemoteAddr = tc.peer
req.Header.Set("X-Forwarded-For", tc.xff)
h.ServeHTTP(httptest.NewRecorder(), req)
if sub.Len() != 1 || sub.Entry(0).ClientIP != tc.wantIP {
t.Fatalf("ClientIP=%q, want %q", sub.Entry(0).ClientIP, tc.wantIP)
}
})
}
}
func TestTrustedProxyUsesValidXRealIPWithoutXFF(t *testing.T) {
upstream := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
w.WriteHeader(http.StatusOK)
}))
defer upstream.Close()
u, _ := url.Parse(upstream.URL)
sub := &captureSubmitter{}
trusted := netip.MustParsePrefix("192.0.2.0/24")
h := newTestHandler(t, u, sub, proxy.Options{TrustedProxies: []netip.Prefix{trusted}})
req := httptest.NewRequest(http.MethodGet, "http://proxy.test/x", nil)
req.RemoteAddr = "192.0.2.10:1234"
req.Header.Set("X-Real-IP", "198.51.100.20")
h.ServeHTTP(httptest.NewRecorder(), req)
if sub.Len() != 1 || sub.Entry(0).ClientIP != "198.51.100.20" {
t.Fatalf("entries=%d", sub.Len())
}
}
func TestResponseBodyTimeoutPreventsCommit(t *testing.T) {
release := make(chan struct{})
upstream := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
w.Header().Set("Content-Type", "application/json")
w.WriteHeader(http.StatusOK)
w.(http.Flusher).Flush()
<-release
}))
u, _ := url.Parse(upstream.URL)
sub := &captureSubmitter{}
h := newTestHandler(t, u, sub, proxy.Options{ResponseTimeout: 50 * time.Millisecond})
srv := httptest.NewServer(h)
resp, err := http.Get(srv.URL + "/v1/test")
if err == nil {
_, _ = io.ReadAll(resp.Body)
_ = resp.Body.Close()
}
time.Sleep(50 * time.Millisecond)
if sub.Len() != 0 {
t.Fatal("timed out response must not be committed")
}
close(release)
srv.Close()
upstream.Close()
}
func TestSSEIdleTimeoutResetsAfterReads(t *testing.T) {
upstream := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
w.Header().Set("Content-Type", "text/event-stream")
f := w.(http.Flusher)
for _, event := range []string{"data: one\n\n", "data: two\n\n", "data: [DONE]\n\n"} {
_, _ = io.WriteString(w, event)
f.Flush()
time.Sleep(30 * time.Millisecond)
}
}))
defer upstream.Close()
u, _ := url.Parse(upstream.URL)
sub := &captureSubmitter{}
h := newTestHandler(t, u, sub, proxy.Options{SSEIdleTimeout: 60 * time.Millisecond})
srv := httptest.NewServer(h)
defer srv.Close()
resp, err := http.Get(srv.URL + "/v1/test")
if err != nil {
t.Fatal(err)
}
_, _ = io.ReadAll(resp.Body)
_ = resp.Body.Close()
if sub.Len() != 1 {
t.Fatalf("entries=%d, want completed SSE commit", sub.Len())
}
}
func TestIncompleteSSEEventPreventsCommit(t *testing.T) {
upstream := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
w.Header().Set("Content-Type", "text/event-stream")
_, _ = io.WriteString(w, "data: partial")
}))
defer upstream.Close()
u, _ := url.Parse(upstream.URL)
sub := &captureSubmitter{}
h := newTestHandler(t, u, sub, proxy.Options{})
h.ServeHTTP(httptest.NewRecorder(), httptest.NewRequest(http.MethodGet, "http://proxy.test/v1/test", nil))
if sub.Len() != 0 {
t.Fatal("SSE ending mid-event must not be committed as complete")
}
}
func TestIncompleteSSEEventAfterTerminalEventPreventsCommit(t *testing.T) {
upstream := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
w.Header().Set("Content-Type", "text/event-stream")
_, _ = io.WriteString(w, "data: [DONE]\n\n")
w.(http.Flusher).Flush()
_, _ = io.WriteString(w, "data: partial")
}))
defer upstream.Close()
u, _ := url.Parse(upstream.URL)
sub := &captureSubmitter{}
h := newTestHandler(t, u, sub, proxy.Options{})
h.ServeHTTP(httptest.NewRecorder(), httptest.NewRequest(http.MethodGet, "http://proxy.test/v1/test", nil))
if sub.Len() != 0 {
t.Fatal("SSE ending mid-event after a terminal event must not be committed")
}
}
func TestShutdownReturnsWithoutHijackedConnections(t *testing.T) {
upstream := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { w.WriteHeader(http.StatusOK) }))
defer upstream.Close()
u, _ := url.Parse(upstream.URL)
h := newTestHandler(t, u, &captureSubmitter{}, proxy.Options{})
ctx, cancel := context.WithTimeout(context.Background(), time.Second)
defer cancel()
if err := h.Shutdown(ctx); err != nil {
t.Fatal(err)
}
}
func TestTerminalTextInNonSSEBodyDoesNotAffectCommit(t *testing.T) {
upstream := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
w.Header().Set("Content-Type", "application/json")
_, _ = io.WriteString(w, `{"text":"data: [DONE]"}`)
}))
defer upstream.Close()
u, _ := url.Parse(upstream.URL)
sub := &captureSubmitter{}
h := newTestHandler(t, u, sub, proxy.Options{})
h.ServeHTTP(httptest.NewRecorder(), httptest.NewRequest(http.MethodGet, "http://proxy.test/v1/test", strings.NewReader("")))
if sub.Len() != 1 {
t.Fatalf("entries=%d, want 1", sub.Len())
}
}
tests/proxy/sse_capture_test.go
+824
View File
@@ -0,0 +1,824 @@
package proxy_test
import (
"bufio"
"encoding/json"
"fmt"
"io"
"net"
"net/http"
"net/http/httptest"
"net/url"
"strings"
"sync"
"testing"
"time"
"git.misaka.ren/M1saka/token_thief/config"
"git.misaka.ren/M1saka/token_thief/logger"
"git.misaka.ren/M1saka/token_thief/proxy"
)
type captureSubmitter struct {
mu sync.Mutex
entries []*logger.LogEntry
}
func (c *captureSubmitter) Submit(e *logger.LogEntry) {
c.mu.Lock()
defer c.mu.Unlock()
c.entries = append(c.entries, e)
}
func (c *captureSubmitter) Len() int {
c.mu.Lock()
defer c.mu.Unlock()
return len(c.entries)
}
func (c *captureSubmitter) Entry(i int) *logger.LogEntry {
c.mu.Lock()
defer c.mu.Unlock()
return c.entries[i]
}
// fakeOpenAIStreamUpstream 模拟一个 OpenAI 兼容的 SSE 上游：
// 分 5 次往响应里 write 一行 SSE 数据，每次都 Flush。
func fakeOpenAIStreamUpstream() *httptest.Server {
chunks := []string{
`data: {"id":"chatcmpl-1","object":"chat.completion.chunk","choices":[{"index":0,"delta":{"role":"assistant"},"finish_reason":null}]}` + "\n\n",
`data: {"id":"chatcmpl-1","object":"chat.completion.chunk","choices":[{"index":0,"delta":{"content":"你"},"finish_reason":null}]}` + "\n\n",
`data: {"id":"chatcmpl-1","object":"chat.completion.chunk","choices":[{"index":0,"delta":{"content":"好"},"finish_reason":null}]}` + "\n\n",
`data: {"id":"chatcmpl-1","object":"chat.completion.chunk","choices":[{"index":0,"delta":{},"finish_reason":"stop"}]}` + "\n\n",
"data: [DONE]\n\n",
}
return httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
w.Header().Set("Content-Type", "text/event-stream")
w.Header().Set("Cache-Control", "no-cache")
w.WriteHeader(http.StatusOK)
flusher := w.(http.Flusher)
for _, ch := range chunks {
_, _ = io.WriteString(w, ch)
flusher.Flush()
time.Sleep(5 * time.Millisecond)
}
}))
}
func fakeAnthropicStreamUpstream() *httptest.Server {
chunks := []string{
`event: message_start` + "\n" + `data: {"type":"message_start","message":{"id":"msg-1","type":"message","role":"assistant","content":[],"model":"claude-3-5-sonnet","stop_reason":null,"stop_sequence":null,"usage":{"input_tokens":10,"output_tokens":1}}}` + "\n\n",
`event: content_block_start` + "\n" + `data: {"type":"content_block_start","index":0,"content_block":{"type":"text","text":""}}` + "\n\n",
`event: content_block_delta` + "\n" + `data: {"type":"content_block_delta","index":0,"delta":{"type":"text_delta","text":"你"}}` + "\n\n",
`event: content_block_delta` + "\n" + `data: {"type":"content_block_delta","index":0,"delta":{"type":"text_delta","text":"好"}}` + "\n\n",
`event: message_delta` + "\n" + `data: {"type":"message_delta","delta":{"stop_reason":"end_turn","stop_sequence":null},"usage":{"output_tokens":3}}` + "\n\n",
`event: message_stop` + "\n" + `data: {"type":"message_stop"}` + "\n\n",
}
return newSSEUpstream(chunks)
}
func fakeGeminiStreamUpstream() *httptest.Server {
chunks := []string{
`data: {"candidates":[{"content":{"parts":[{"text":"你"}],"role":"model"},"index":0}]}` + "\n\n",
`data: {"candidates":[{"content":{"parts":[{"text":"好"}],"role":"model"},"finishReason":"STOP","index":0}],"usageMetadata":{"promptTokenCount":2,"candidatesTokenCount":2,"totalTokenCount":4}}` + "\n\n",
}
return newSSEUpstream(chunks)
}
func fakeUnknownStreamUpstream() *httptest.Server {
return newSSEUpstream([]string{
`event: custom` + "\n" + `data: not-json` + "\n\n",
})
}
func fakeOpenAIStreamUpstreamThatStaysOpen(release <-chan struct{}) *httptest.Server {
chunks := []string{
`data: {"id":"chatcmpl-hang","object":"chat.completion.chunk","choices":[{"index":0,"delta":{"role":"assistant"},"finish_reason":null}]}` + "\n\n",
`data: {"id":"chatcmpl-hang","object":"chat.completion.chunk","choices":[{"index":0,"delta":{"content":"ok"},"finish_reason":null}]}` + "\n\n",
`data: {"id":"chatcmpl-hang","object":"chat.completion.chunk","choices":[{"index":0,"delta":{},"finish_reason":"stop"}]}` + "\n\n",
"data: [DONE]\n\n",
}
return httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
w.Header().Set("Content-Type", "text/event-stream")
w.WriteHeader(http.StatusOK)
flusher := w.(http.Flusher)
for _, ch := range chunks {
_, _ = io.WriteString(w, ch)
flusher.Flush()
}
<-release
}))
}
func newSSEUpstream(chunks []string) *httptest.Server {
return httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
w.Header().Set("Content-Type", "text/event-stream")
w.Header().Set("Cache-Control", "no-cache")
w.WriteHeader(http.StatusOK)
flusher := w.(http.Flusher)
for _, ch := range chunks {
_, _ = io.WriteString(w, ch)
flusher.Flush()
time.Sleep(5 * time.Millisecond)
}
}))
}
func TestSSEChunkAssembly(t *testing.T) {
upstream := fakeOpenAIStreamUpstream()
defer upstream.Close()
u, _ := url.Parse(upstream.URL)
sub := &captureSubmitter{}
filter, err := config.NewFilter(config.FilterDisabled, nil)
if err != nil {
t.Fatal(err)
}
h := proxy.New(u, filter, sub, 1024*1024)
proxySrv := httptest.NewServer(h)
defer proxySrv.Close()
// 客户端走原始 TCP，逐字节读 + 打印，验证流式实时到达
pu, _ := url.Parse(proxySrv.URL)
conn, err := net.DialTimeout("tcp", pu.Host, 3*time.Second)
if err != nil {
t.Fatal(err)
}
defer conn.Close()
fmt.Fprintf(conn, "POST /v1/chat/completions HTTP/1.1\r\nHost: %s\r\nContent-Length: 0\r\n\r\n", pu.Host)
_ = conn.SetReadDeadline(time.Now().Add(3 * time.Second))
br := bufio.NewReader(conn)
resp, err := http.ReadResponse(br, nil)
if err != nil {
t.Fatal(err)
}
clientBody, _ := io.ReadAll(resp.Body)
// 等待 proxy.ServeHTTP 返回并把 entry 提交
deadline := time.Now().Add(2 * time.Second)
for sub.Len() == 0 && time.Now().Before(deadline) {
time.Sleep(10 * time.Millisecond)
}
if sub.Len() == 0 {
t.Fatal("no log entry captured")
}
e := sub.Entry(0)
t.Logf("\n========== 客户端收到的字节 ==========\n%s", clientBody)
t.Logf("\n========== 数据库 response_body 字段（按字节原样存储）==========\n%s", e.ResponseBody)
t.Logf("\n========== 元数据 ==========")
t.Logf("is_stream = %v", e.IsStream)
t.Logf("status_code = %d", e.StatusCode)
t.Logf("len(body) = %d bytes", len(e.ResponseBody))
t.Logf("response_truncated = %v", e.ResponseTruncated)
if !strings.Contains(string(clientBody), `"content":"你"`) ||
!strings.Contains(string(clientBody), `"content":"好"`) ||
!strings.Contains(string(clientBody), "[DONE]") {
t.Errorf("client response body 缺少预期 chunk 内容")
}
var captured struct {
Choices []struct {
Message struct {
Role string `json:"role"`
Content string `json:"content"`
} `json:"message"`
FinishReason string `json:"finish_reason"`
} `json:"choices"`
}
if err := json.Unmarshal(e.ResponseBody, &captured); err != nil {
t.Fatalf("stream response_body should be assembled JSON: %v; body=%q", err, e.ResponseBody)
}
if len(captured.Choices) != 1 {
t.Fatalf("assembled JSON choices length=%d, want 1", len(captured.Choices))
}
if captured.Choices[0].Message.Role != "assistant" {
t.Errorf("assembled role=%q, want assistant", captured.Choices[0].Message.Role)
}
if captured.Choices[0].Message.Content != "你好" {
t.Errorf("assembled content=%q, want 你好", captured.Choices[0].Message.Content)
}
if captured.Choices[0].FinishReason != "stop" {
t.Errorf("assembled finish_reason=%q, want stop", captured.Choices[0].FinishReason)
}
if !e.IsStream {
t.Errorf("is_stream 应为 true")
}
}
func TestOpenAIStreamPreservesReasoningAndMetadata(t *testing.T) {
upstream := newSSEUpstream([]string{
`data: {"id":"resp-1","object":"chat.completion.chunk","created":1779335544,"model":"gpt-5.4-mini-2026-03-17","choices":[{"index":0,"delta":{"role":"assistant","reasoning_content":"think "},"finish_reason":null,"native_finish_reason":null}]}` + "\n\n",
`data: {"id":"resp-1","object":"chat.completion.chunk","created":1779335544,"model":"gpt-5.4-mini-2026-03-17","choices":[{"index":0,"delta":{"reasoning_content":"hard"},"finish_reason":null,"native_finish_reason":null}]}` + "\n\n",
`data: {"id":"resp-1","object":"chat.completion.chunk","created":1779335544,"model":"gpt-5.4-mini-2026-03-17","choices":[{"index":0,"delta":{"content":"final"},"finish_reason":null,"native_finish_reason":null}]}` + "\n\n",
`data: {"id":"resp-1","object":"chat.completion.chunk","created":1779335544,"model":"gpt-5.4-mini-2026-03-17","choices":[{"index":0,"delta":{},"finish_reason":"stop","native_finish_reason":"stop"}]}` + "\n\n",
"data: [DONE]\n\n",
})
e := requestStreamEntry(t, upstream, "/v1/chat/completions")
var captured struct {
ID string `json:"id"`
Object string `json:"object"`
Created int64 `json:"created"`
Model string `json:"model"`
Choices []struct {
Index int `json:"index"`
Message struct {
Role string `json:"role"`
Content string `json:"content"`
ReasoningContent string `json:"reasoning_content"`
} `json:"message"`
FinishReason string `json:"finish_reason"`
NativeFinishReason string `json:"native_finish_reason"`
} `json:"choices"`
}
if err := json.Unmarshal(e.ResponseBody, &captured); err != nil {
t.Fatalf("openai stream response_body should be JSON: %v; body=%q", err, e.ResponseBody)
}
if captured.ID != "resp-1" || captured.Object != "chat.completion" || captured.Created != 1779335544 || captured.Model != "gpt-5.4-mini-2026-03-17" {
t.Fatalf("unexpected metadata: %+v", captured)
}
if len(captured.Choices) != 1 {
t.Fatalf("choices length=%d, want 1", len(captured.Choices))
}
choice := captured.Choices[0]
if choice.Message.Role != "assistant" || choice.Message.Content != "final" || choice.Message.ReasoningContent != "think hard" {
t.Fatalf("unexpected message: %+v", choice.Message)
}
if choice.FinishReason != "stop" || choice.NativeFinishReason != "stop" {
t.Fatalf("unexpected finish reasons: %+v", choice)
}
}
func TestOpenAIStreamDoneDoesNotSubmitBeforeUpstreamCloses(t *testing.T) {
release := make(chan struct{})
upstream := fakeOpenAIStreamUpstreamThatStaysOpen(release)
u, _ := url.Parse(upstream.URL)
sub := &captureSubmitter{}
filter, err := config.NewFilter(config.FilterDisabled, nil)
if err != nil {
t.Fatal(err)
}
h := proxy.New(u, filter, sub, 1024*1024)
proxySrv := httptest.NewServer(h)
released := false
defer func() {
if !released {
close(release)
}
proxySrv.Close()
upstream.Close()
}()
clientDone := make(chan error, 1)
go func() {
resp, err := http.Post(proxySrv.URL+"/v1/chat/completions", "application/json", strings.NewReader(`{"stream":true}`))
if err != nil {
clientDone <- err
return
}
_, _ = io.ReadAll(resp.Body)
_ = resp.Body.Close()
clientDone <- nil
}()
time.Sleep(100 * time.Millisecond)
if sub.Len() != 0 {
t.Fatal("terminal SSE event must not submit before ReverseProxy returns")
}
close(release)
released = true
select {
case err := <-clientDone:
if err != nil {
t.Fatal(err)
}
if sub.Len() != 1 {
t.Fatalf("stream should be submitted once, got %d entries", sub.Len())
}
case <-time.After(2 * time.Second):
t.Fatal("client did not finish after upstream closed")
}
}
func TestOpenAIStreamPreservesUsageChunk(t *testing.T) {
upstream := newSSEUpstream([]string{
`data: {"id":"resp-usage","object":"chat.completion.chunk","created":1779335544,"model":"gpt-5.4-mini","system_fingerprint":"fp_123","choices":[{"index":0,"delta":{"role":"assistant","content":"ok"},"finish_reason":null}]}` + "\n\n",
`data: {"id":"resp-usage","object":"chat.completion.chunk","created":1779335544,"model":"gpt-5.4-mini","system_fingerprint":"fp_123","choices":[{"index":0,"delta":{},"finish_reason":"stop"}],"usage":null}` + "\n\n",
`data: {"id":"resp-usage","object":"chat.completion.chunk","created":1779335544,"model":"gpt-5.4-mini","system_fingerprint":"fp_123","choices":[],"usage":{"prompt_tokens":5,"completion_tokens":2,"total_tokens":7}}` + "\n\n",
"data: [DONE]\n\n",
})
e := requestStreamEntry(t, upstream, "/v1/chat/completions")
var captured struct {
ID string `json:"id"`
SystemFingerprint string `json:"system_fingerprint"`
Choices []struct {
Message struct {
Content string `json:"content"`
} `json:"message"`
} `json:"choices"`
Usage struct {
PromptTokens int `json:"prompt_tokens"`
CompletionTokens int `json:"completion_tokens"`
TotalTokens int `json:"total_tokens"`
} `json:"usage"`
}
if err := json.Unmarshal(e.ResponseBody, &captured); err != nil {
t.Fatalf("openai stream with usage should be assembled JSON: %v; body=%q", err, e.ResponseBody)
}
if captured.ID != "resp-usage" || captured.SystemFingerprint != "fp_123" {
t.Fatalf("metadata not preserved: %+v", captured)
}
if len(captured.Choices) != 1 || captured.Choices[0].Message.Content != "ok" {
t.Fatalf("choices not assembled: %+v", captured.Choices)
}
if captured.Usage.PromptTokens != 5 || captured.Usage.CompletionTokens != 2 || captured.Usage.TotalTokens != 7 {
t.Fatalf("usage not preserved: %+v", captured.Usage)
}
}
func TestOpenAIStreamAssemblesToolCalls(t *testing.T) {
upstream := newSSEUpstream([]string{
`data: {"id":"resp-tools","object":"chat.completion.chunk","created":1779335544,"model":"gpt-5.4-mini","choices":[{"index":0,"delta":{"role":"assistant","tool_calls":[{"index":0,"id":"call_1","type":"function","function":{"name":"lookup","arguments":"{\"q\":"}}]},"finish_reason":null}]}` + "\n\n",
`data: {"id":"resp-tools","object":"chat.completion.chunk","created":1779335544,"model":"gpt-5.4-mini","choices":[{"index":0,"delta":{"tool_calls":[{"index":0,"function":{"arguments":"\"weather\"}"}}]},"finish_reason":null}]}` + "\n\n",
`data: {"id":"resp-tools","object":"chat.completion.chunk","created":1779335544,"model":"gpt-5.4-mini","choices":[{"index":0,"delta":{},"finish_reason":"tool_calls"}]}` + "\n\n",
"data: [DONE]\n\n",
})
e := requestStreamEntry(t, upstream, "/v1/chat/completions")
var captured struct {
Choices []struct {
Message struct {
Role string `json:"role"`
Content string `json:"content"`
ToolCalls []struct {
ID string `json:"id"`
Type string `json:"type"`
Function struct {
Name string `json:"name"`
Arguments string `json:"arguments"`
} `json:"function"`
} `json:"tool_calls"`
} `json:"message"`
FinishReason string `json:"finish_reason"`
} `json:"choices"`
}
if err := json.Unmarshal(e.ResponseBody, &captured); err != nil {
t.Fatalf("openai tool stream should be assembled JSON: %v; body=%q", err, e.ResponseBody)
}
if len(captured.Choices) != 1 || captured.Choices[0].FinishReason != "tool_calls" {
t.Fatalf("unexpected choices: %+v", captured.Choices)
}
message := captured.Choices[0].Message
if message.Role != "assistant" || message.Content != "" {
t.Fatalf("unexpected message basics: %+v", message)
}
if len(message.ToolCalls) != 1 {
t.Fatalf("tool_calls length=%d, want 1; body=%s", len(message.ToolCalls), e.ResponseBody)
}
tool := message.ToolCalls[0]
if tool.ID != "call_1" || tool.Type != "function" || tool.Function.Name != "lookup" || tool.Function.Arguments != `{"q":"weather"}` {
t.Fatalf("unexpected tool call: %+v", tool)
}
}
func TestOpenAIResponsesStreamAssemblesCompletedResponse(t *testing.T) {
upstream := newSSEUpstream([]string{
`event: response.created` + "\n" + `data: {"type":"response.created","response":{"id":"resp-1","object":"response","status":"in_progress","model":"gpt-5.4-mini","output":[]}}` + "\n\n",
`event: response.output_text.delta` + "\n" + `data: {"type":"response.output_text.delta","item_id":"msg-1","output_index":0,"content_index":0,"delta":"hello"}` + "\n\n",
`event: response.completed` + "\n" + `data: {"type":"response.completed","response":{"id":"resp-1","object":"response","status":"completed","model":"gpt-5.4-mini","output":[{"id":"msg-1","type":"message","role":"assistant","content":[{"type":"output_text","text":"hello"}]}],"usage":{"input_tokens":3,"output_tokens":1,"total_tokens":4}}}` + "\n\n",
})
e := requestStreamEntry(t, upstream, "/v1/responses")
var captured struct {
ID string `json:"id"`
Object string `json:"object"`
Status string `json:"status"`
Model string `json:"model"`
Output []struct {
Type string `json:"type"`
Role string `json:"role"`
Content []struct {
Type string `json:"type"`
Text string `json:"text"`
} `json:"content"`
} `json:"output"`
Usage struct {
InputTokens int `json:"input_tokens"`
OutputTokens int `json:"output_tokens"`
TotalTokens int `json:"total_tokens"`
} `json:"usage"`
}
if err := json.Unmarshal(e.ResponseBody, &captured); err != nil {
t.Fatalf("responses stream should be completed response JSON: %v; body=%q", err, e.ResponseBody)
}
if captured.ID != "resp-1" || captured.Object != "response" || captured.Status != "completed" || captured.Model != "gpt-5.4-mini" {
t.Fatalf("unexpected response metadata: %+v", captured)
}
if len(captured.Output) != 1 || len(captured.Output[0].Content) != 1 || captured.Output[0].Content[0].Text != "hello" {
t.Fatalf("unexpected response output: %+v", captured.Output)
}
if captured.Usage.TotalTokens != 4 {
t.Fatalf("usage not preserved: %+v", captured.Usage)
}
}
func TestOpenAICompletionsStreamAssemblesText(t *testing.T) {
upstream := newSSEUpstream([]string{
`data: {"id":"cmpl-1","object":"text_completion","created":1779335544,"model":"gpt-5.4-mini","choices":[{"index":0,"text":"hello","finish_reason":null}]}` + "\n\n",
`data: {"id":"cmpl-1","object":"text_completion","created":1779335544,"model":"gpt-5.4-mini","choices":[{"index":0,"text":" world","finish_reason":null}]}` + "\n\n",
`data: {"id":"cmpl-1","object":"text_completion","created":1779335544,"model":"gpt-5.4-mini","choices":[{"index":0,"text":"","finish_reason":"stop"}],"usage":{"prompt_tokens":2,"completion_tokens":2,"total_tokens":4}}` + "\n\n",
"data: [DONE]\n\n",
})
e := requestStreamEntry(t, upstream, "/v1/completions")
var captured struct {
ID string `json:"id"`
Object string `json:"object"`
Created int64 `json:"created"`
Model string `json:"model"`
Choices []struct {
Index int `json:"index"`
Text string `json:"text"`
FinishReason string `json:"finish_reason"`
} `json:"choices"`
Usage struct {
TotalTokens int `json:"total_tokens"`
} `json:"usage"`
}
if err := json.Unmarshal(e.ResponseBody, &captured); err != nil {
t.Fatalf("completions stream should be assembled JSON: %v; body=%q", err, e.ResponseBody)
}
if captured.ID != "cmpl-1" || captured.Object != "text_completion" || captured.Model != "gpt-5.4-mini" || captured.Created != 1779335544 {
t.Fatalf("unexpected metadata: %+v", captured)
}
if len(captured.Choices) != 1 || captured.Choices[0].Text != "hello world" || captured.Choices[0].FinishReason != "stop" {
t.Fatalf("unexpected choices: %+v", captured.Choices)
}
if captured.Usage.TotalTokens != 4 {
t.Fatalf("usage not preserved: %+v", captured.Usage)
}
}
func TestAnthropicSSEAssemblesNativeMessageJSON(t *testing.T) {
e := requestStreamEntry(t, fakeAnthropicStreamUpstream(), "/v1/messages")
var captured struct {
ID string `json:"id"`
Type string `json:"type"`
Role string `json:"role"`
Content []struct {
Type string `json:"type"`
Text string `json:"text"`
} `json:"content"`
StopReason string `json:"stop_reason"`
Usage struct {
InputTokens int `json:"input_tokens"`
OutputTokens int `json:"output_tokens"`
} `json:"usage"`
}
if err := json.Unmarshal(e.ResponseBody, &captured); err != nil {
t.Fatalf("anthropic stream response_body should be native JSON: %v; body=%q", err, e.ResponseBody)
}
if captured.ID != "msg-1" || captured.Type != "message" || captured.Role != "assistant" {
t.Fatalf("unexpected anthropic message metadata: %+v", captured)
}
if len(captured.Content) != 1 || captured.Content[0].Type != "text" || captured.Content[0].Text != "你好" {
t.Fatalf("unexpected anthropic content: %+v", captured.Content)
}
if captured.StopReason != "end_turn" {
t.Errorf("stop_reason=%q, want end_turn", captured.StopReason)
}
if captured.Usage.InputTokens != 10 || captured.Usage.OutputTokens != 3 {
t.Errorf("usage=%+v, want input=10 output=3", captured.Usage)
}
}
func TestAnthropicSSEAssemblesToolUseContent(t *testing.T) {
upstream := newSSEUpstream([]string{
`event: message_start` + "\n" + `data: {"type":"message_start","message":{"id":"msg-tool","type":"message","role":"assistant","content":[],"model":"claude-3-5-sonnet","stop_reason":null,"stop_sequence":null,"usage":{"input_tokens":10,"output_tokens":1}}}` + "\n\n",
`event: content_block_start` + "\n" + `data: {"type":"content_block_start","index":0,"content_block":{"type":"tool_use","id":"toolu_1","name":"lookup","input":{}}}` + "\n\n",
`event: content_block_delta` + "\n" + `data: {"type":"content_block_delta","index":0,"delta":{"type":"input_json_delta","partial_json":"{\"q\":"}}` + "\n\n",
`event: content_block_delta` + "\n" + `data: {"type":"content_block_delta","index":0,"delta":{"type":"input_json_delta","partial_json":"\"weather\"}"}}` + "\n\n",
`event: message_delta` + "\n" + `data: {"type":"message_delta","delta":{"stop_reason":"tool_use","stop_sequence":null},"usage":{"output_tokens":8}}` + "\n\n",
`event: message_stop` + "\n" + `data: {"type":"message_stop"}` + "\n\n",
})
e := requestStreamEntry(t, upstream, "/v1/messages")
var captured struct {
Content []struct {
Type string `json:"type"`
ID string `json:"id"`
Name string `json:"name"`
Input map[string]any `json:"input"`
} `json:"content"`
StopReason string `json:"stop_reason"`
}
if err := json.Unmarshal(e.ResponseBody, &captured); err != nil {
t.Fatalf("anthropic tool stream should be JSON: %v; body=%q", err, e.ResponseBody)
}
if len(captured.Content) != 1 {
t.Fatalf("content length=%d, want 1; body=%s", len(captured.Content), e.ResponseBody)
}
tool := captured.Content[0]
if tool.Type != "tool_use" || tool.ID != "toolu_1" || tool.Name != "lookup" || tool.Input["q"] != "weather" {
t.Fatalf("unexpected tool content: %+v", tool)
}
if captured.StopReason != "tool_use" {
t.Fatalf("stop_reason=%q, want tool_use", captured.StopReason)
}
}
func TestGeminiSSEAssemblesNativeGenerateContentJSON(t *testing.T) {
e := requestStreamEntry(t, fakeGeminiStreamUpstream(), "/v1beta/models/gemini-1.5-pro:generateContent")
var captured struct {
Candidates []struct {
Content struct {
Role string `json:"role"`
Parts []struct {
Text string `json:"text"`
} `json:"parts"`
} `json:"content"`
FinishReason string `json:"finishReason"`
Index int `json:"index"`
} `json:"candidates"`
UsageMetadata struct {
PromptTokenCount int `json:"promptTokenCount"`
CandidatesTokenCount int `json:"candidatesTokenCount"`
TotalTokenCount int `json:"totalTokenCount"`
} `json:"usageMetadata"`
}
if err := json.Unmarshal(e.ResponseBody, &captured); err != nil {
t.Fatalf("gemini stream response_body should be native JSON: %v; body=%q", err, e.ResponseBody)
}
if len(captured.Candidates) != 1 {
t.Fatalf("candidates length=%d, want 1", len(captured.Candidates))
}
candidate := captured.Candidates[0]
if candidate.Content.Role != "model" || len(candidate.Content.Parts) != 1 || candidate.Content.Parts[0].Text != "你好" {
t.Fatalf("unexpected gemini content: %+v", candidate.Content)
}
if candidate.FinishReason != "STOP" {
t.Errorf("finishReason=%q, want STOP", candidate.FinishReason)
}
if captured.UsageMetadata.TotalTokenCount != 4 {
t.Errorf("usageMetadata=%+v, want totalTokenCount=4", captured.UsageMetadata)
}
}
func TestGeminiStreamPreservesSafetyRatingsAndUsageMetadata(t *testing.T) {
upstream := newSSEUpstream([]string{
`data: {"candidates":[{"content":{"parts":[{"text":"你"}],"role":"model"},"finishReason":null,"index":0,"safetyRatings":[{"category":"HARM_CATEGORY_HARASSMENT","probability":"NEGLIGIBLE"}]}],"usageMetadata":{"promptTokenCount":8,"toolUsePromptTokenCount":0,"candidatesTokenCount":0,"totalTokenCount":8,"thoughtsTokenCount":10}}` + "\n\n",
`data: {"candidates":[{"content":{"parts":[{"text":"好"}],"role":"model"},"finishReason":"STOP","index":0,"safetyRatings":[{"category":"HARM_CATEGORY_HARASSMENT","probability":"NEGLIGIBLE"}]}],"usageMetadata":{"promptTokenCount":8,"toolUsePromptTokenCount":0,"candidatesTokenCount":2,"totalTokenCount":20,"thoughtsTokenCount":10}}` + "\n\n",
})
e := requestStreamEntry(t, upstream, "/v1beta/models/gemini-1.5-pro:streamGenerateContent")
var captured struct {
Candidates []struct {
Content struct {
Parts []struct {
Text string `json:"text"`
} `json:"parts"`
} `json:"content"`
SafetyRatings []struct {
Category string `json:"category"`
Probability string `json:"probability"`
} `json:"safetyRatings"`
} `json:"candidates"`
UsageMetadata struct {
PromptTokenCount int `json:"promptTokenCount"`
ToolUsePromptTokenCount int `json:"toolUsePromptTokenCount"`
CandidatesTokenCount int `json:"candidatesTokenCount"`
TotalTokenCount int `json:"totalTokenCount"`
ThoughtsTokenCount int `json:"thoughtsTokenCount"`
} `json:"usageMetadata"`
}
if err := json.Unmarshal(e.ResponseBody, &captured); err != nil {
t.Fatalf("gemini stream response_body should be JSON: %v; body=%q", err, e.ResponseBody)
}
if len(captured.Candidates) != 1 || len(captured.Candidates[0].SafetyRatings) != 1 {
t.Fatalf("expected safetyRatings to be preserved, got %+v", captured.Candidates)
}
if captured.Candidates[0].SafetyRatings[0].Category != "HARM_CATEGORY_HARASSMENT" {
t.Fatalf("unexpected safetyRatings: %+v", captured.Candidates[0].SafetyRatings)
}
if captured.UsageMetadata.ToolUsePromptTokenCount != 0 || captured.UsageMetadata.ThoughtsTokenCount != 10 || captured.UsageMetadata.TotalTokenCount != 20 {
t.Fatalf("usageMetadata fields not preserved: %+v", captured.UsageMetadata)
}
}
func TestGeminiStreamPreservesFunctionCallParts(t *testing.T) {
upstream := newSSEUpstream([]string{
`data: {"candidates":[{"content":{"parts":[{"functionCall":{"name":"lookup","args":{"q":"weather"}}}],"role":"model"},"finishReason":"STOP","index":0}],"usageMetadata":{"promptTokenCount":2,"candidatesTokenCount":2,"totalTokenCount":4}}` + "\n\n",
})
e := requestStreamEntry(t, upstream, "/v1beta/models/gemini-1.5-pro:streamGenerateContent")
var captured struct {
Candidates []struct {
Content struct {
Parts []struct {
FunctionCall struct {
Name string `json:"name"`
Args map[string]any `json:"args"`
} `json:"functionCall"`
} `json:"parts"`
} `json:"content"`
} `json:"candidates"`
}
if err := json.Unmarshal(e.ResponseBody, &captured); err != nil {
t.Fatalf("gemini functionCall stream should be JSON: %v; body=%q", err, e.ResponseBody)
}
call := captured.Candidates[0].Content.Parts[0].FunctionCall
if call.Name != "lookup" || call.Args["q"] != "weather" {
t.Fatalf("functionCall not preserved: %+v; body=%s", call, e.ResponseBody)
}
}
func TestUnknownSSEKeepsRawBody(t *testing.T) {
e := requestStreamEntry(t, fakeUnknownStreamUpstream(), "/v1/chat/completions")
if string(e.ResponseBody) != "event: custom\ndata: not-json\n\n" {
t.Fatalf("unknown stream should keep raw body, got %q", e.ResponseBody)
}
}
func TestTruncatedSSEKeepsCapturedRawBody(t *testing.T) {
upstream := fakeOpenAIStreamUpstream()
defer upstream.Close()
u, _ := url.Parse(upstream.URL)
sub := &captureSubmitter{}
filter, err := config.NewFilter(config.FilterDisabled, nil)
if err != nil {
t.Fatal(err)
}
h := proxy.New(u, filter, sub, 520)
proxySrv := httptest.NewServer(h)
defer proxySrv.Close()
resp, err := http.Post(proxySrv.URL+"/v1/chat/completions", "application/json", nil)
if err != nil {
t.Fatal(err)
}
_, _ = io.ReadAll(resp.Body)
_ = resp.Body.Close()
deadline := time.Now().Add(2 * time.Second)
for sub.Len() == 0 && time.Now().Before(deadline) {
time.Sleep(10 * time.Millisecond)
}
if sub.Len() == 0 {
t.Fatal("no log entry captured")
}
e := sub.Entry(0)
if !e.ResponseTruncated {
t.Fatal("response should be marked truncated")
}
if !strings.HasPrefix(string(e.ResponseBody), "data: ") {
t.Fatalf("truncated stream should keep captured raw body, got %q", e.ResponseBody)
}
if json.Valid(e.ResponseBody) {
t.Fatalf("truncated stream should not be assembled as JSON, got %q", e.ResponseBody)
}
}
func TestMultimodalStreamRequestBodyIsCaptured(t *testing.T) {
upstream := newSSEUpstream([]string{
`data: {"id":"chatcmpl-image","object":"chat.completion.chunk","choices":[{"index":0,"delta":{"role":"assistant"},"finish_reason":null}]}` + "\n\n",
`data: {"id":"chatcmpl-image","object":"chat.completion.chunk","choices":[{"index":0,"delta":{"content":"ok"},"finish_reason":null}]}` + "\n\n",
`data: {"id":"chatcmpl-image","object":"chat.completion.chunk","choices":[{"index":0,"delta":{},"finish_reason":"stop"}]}` + "\n\n",
"data: [DONE]\n\n",
})
defer upstream.Close()
u, _ := url.Parse(upstream.URL)
sub := &captureSubmitter{}
filter, err := config.NewFilter(config.FilterDisabled, nil)
if err != nil {
t.Fatal(err)
}
h := proxy.New(u, filter, sub, 1024*1024)
proxySrv := httptest.NewServer(h)
defer proxySrv.Close()
reqBody := `{"model":"gpt-5.4-mini","messages":[{"role":"user","content":[{"type":"text","text":"describe"},{"type":"image_url","image_url":{"url":"data:image/png;base64,AAAA"}}]}],"stream":true}`
resp, err := http.Post(proxySrv.URL+"/v1/chat/completions", "application/json", strings.NewReader(reqBody))
if err != nil {
t.Fatal(err)
}
_, _ = io.ReadAll(resp.Body)
_ = resp.Body.Close()
deadline := time.Now().Add(2 * time.Second)
for sub.Len() == 0 && time.Now().Before(deadline) {
time.Sleep(10 * time.Millisecond)
}
if sub.Len() == 0 {
t.Fatal("no log entry captured")
}
e := sub.Entry(0)
if e.RequestTruncated {
t.Fatal("multimodal request should not be truncated")
}
requestText := string(e.RequestBody)
if !strings.Contains(requestText, `"image_url"`) || !strings.Contains(requestText, `data:image/png;base64,AAAA`) {
t.Fatalf("request_body should contain image input, got %q", requestText)
}
if !e.IsStream {
t.Fatal("response should still be marked stream")
}
}
func TestChunkedMultimodalRequestBodyIsCaptured(t *testing.T) {
upstream := newSSEUpstream([]string{
`data: {"id":"chatcmpl-image","object":"chat.completion.chunk","choices":[{"index":0,"delta":{"content":"ok"},"finish_reason":"stop"}]}` + "\n\n",
"data: [DONE]\n\n",
})
defer upstream.Close()
u, _ := url.Parse(upstream.URL)
sub := &captureSubmitter{}
filter, err := config.NewFilter(config.FilterDisabled, nil)
if err != nil {
t.Fatal(err)
}
h := proxy.New(u, filter, sub, 1024*1024)
proxySrv := httptest.NewServer(h)
defer proxySrv.Close()
reqBody := `{"model":"gpt-5.4-mini","messages":[{"role":"user","content":[{"type":"text","text":"describe"},{"type":"image_url","image_url":{"url":"data:image/png;base64,AAAA"}}]}],"stream":true}`
req, err := http.NewRequest(http.MethodPost, proxySrv.URL+"/v1/chat/completions", strings.NewReader(reqBody))
if err != nil {
t.Fatal(err)
}
req.ContentLength = -1
req.Header.Set("Content-Type", "application/json")
req.TransferEncoding = []string{"chunked"}
resp, err := http.DefaultClient.Do(req)
if err != nil {
t.Fatal(err)
}
_, _ = io.ReadAll(resp.Body)
_ = resp.Body.Close()
deadline := time.Now().Add(2 * time.Second)
for sub.Len() == 0 && time.Now().Before(deadline) {
time.Sleep(10 * time.Millisecond)
}
if sub.Len() == 0 {
t.Fatal("no log entry captured")
}
e := sub.Entry(0)
requestText := string(e.RequestBody)
if !strings.Contains(requestText, `"image_url"`) || !strings.Contains(requestText, `data:image/png;base64,AAAA`) {
t.Fatalf("chunked request_body should contain image input, got %q", requestText)
}
}
func requestStreamEntry(t *testing.T, upstream *httptest.Server, path string) *logger.LogEntry {
t.Helper()
defer upstream.Close()
u, _ := url.Parse(upstream.URL)
sub := &captureSubmitter{}
filter, err := config.NewFilter(config.FilterDisabled, nil)
if err != nil {
t.Fatal(err)
}
h := proxy.New(u, filter, sub, 1024*1024)
proxySrv := httptest.NewServer(h)
defer proxySrv.Close()
resp, err := http.Post(proxySrv.URL+path, "application/json", nil)
if err != nil {
t.Fatal(err)
}
_, _ = io.ReadAll(resp.Body)
_ = resp.Body.Close()
deadline := time.Now().Add(2 * time.Second)
for sub.Len() == 0 && time.Now().Before(deadline) {
time.Sleep(10 * time.Millisecond)
}
if sub.Len() == 0 {
t.Fatal("no log entry captured")
}
return sub.Entry(0)
}
tests/proxy/websocket_test.go
+299
View File
@@ -0,0 +1,299 @@
package proxy_test
import (
"bufio"
"context"
"io"
"net"
"net/http"
"net/http/httptest"
"net/url"
"strings"
"testing"
"time"
"git.misaka.ren/M1saka/token_thief/config"
"git.misaka.ren/M1saka/token_thief/logger"
"git.misaka.ren/M1saka/token_thief/proxy"
)
type noopSubmitter struct{}
func (noopSubmitter) Submit(*logger.LogEntry) {}
type chanSubmitter chan *logger.LogEntry
func (c chanSubmitter) Submit(e *logger.LogEntry) { c <- e }
// fakeUpstream 模拟一个最简 WebSocket 升级：
// 收到 GET + Upgrade: websocket 后回 101，然后做字节回声直到对端关闭。
func fakeUpstream(t *testing.T) *httptest.Server {
t.Helper()
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
if strings.ToLower(r.Header.Get("Upgrade")) != "websocket" {
http.Error(w, "expected websocket upgrade", http.StatusBadRequest)
return
}
hj, ok := w.(http.Hijacker)
if !ok {
http.Error(w, "no hijack", http.StatusInternalServerError)
return
}
conn, brw, err := hj.Hijack()
if err != nil {
t.Errorf("upstream hijack: %v", err)
return
}
defer conn.Close()
// 直接回 101 握手响应（简化版，不做真正 Sec-WebSocket-Accept 计算）
_, _ = brw.WriteString("HTTP/1.1 101 Switching Protocols\r\n" +
"Upgrade: websocket\r\n" +
"Connection: Upgrade\r\n" +
"\r\n")
_ = brw.Flush()
// echo
buf := make([]byte, 1024)
for {
n, err := conn.Read(buf)
if err != nil {
return
}
if _, err := conn.Write(buf[:n]); err != nil {
return
}
}
}))
return srv
}
func TestWebSocketHandshakeCapturedAndShutdownClosesConnection(t *testing.T) {
upstream := fakeUpstream(t)
defer upstream.Close()
u, _ := url.Parse(upstream.URL)
filter, err := config.NewFilter(config.FilterDisabled, nil)
if err != nil {
t.Fatal(err)
}
entries := make(chanSubmitter, 1)
h := proxy.New(u, filter, entries, 1024)
proxySrv := httptest.NewServer(h)
defer proxySrv.Close()
pu, _ := url.Parse(proxySrv.URL)
conn, err := net.DialTimeout("tcp", pu.Host, time.Second)
if err != nil {
t.Fatal(err)
}
defer conn.Close()
_, _ = io.WriteString(conn, "GET /v1/realtime HTTP/1.1\r\nHost: "+pu.Host+"\r\nUpgrade: websocket\r\nConnection: Upgrade\r\n\r\n")
br := bufio.NewReader(conn)
resp, err := http.ReadResponse(br, nil)
if err != nil {
t.Fatal(err)
}
if resp.StatusCode != http.StatusSwitchingProtocols {
t.Fatalf("status=%d", resp.StatusCode)
}
select {
case entry := <-entries:
if entry.StatusCode != http.StatusSwitchingProtocols || len(entry.ResponseBody) != 0 {
t.Fatalf("handshake entry status=%d body=%q", entry.StatusCode, entry.ResponseBody)
}
case <-time.After(time.Second):
t.Fatal("websocket handshake was not captured")
}
payload := []byte("frame-data-must-not-be-logged")
if _, err := conn.Write(payload); err != nil {
t.Fatal(err)
}
echo := make([]byte, len(payload))
if _, err := io.ReadFull(br, echo); err != nil {
t.Fatalf("read websocket payload: %v", err)
}
if string(echo) != string(payload) {
t.Fatalf("echo mismatch: got %q want %q", echo, payload)
}
select {
case entry := <-entries:
t.Fatalf("websocket frame produced an extra log entry: %+v", entry)
case <-time.After(50 * time.Millisecond):
}
ctx, cancel := context.WithTimeout(context.Background(), time.Second)
defer cancel()
if err := h.Shutdown(ctx); err != nil {
t.Fatal(err)
}
_ = conn.SetReadDeadline(time.Now().Add(time.Second))
if _, err := conn.Read(make([]byte, 1)); err == nil {
t.Fatal("connection remains open after Shutdown")
}
if err := h.Shutdown(ctx); err != nil {
t.Fatalf("second Shutdown: %v", err)
}
}
func TestWebSocketNaturalCloseUnregistersConnection(t *testing.T) {
upstream := fakeUpstream(t)
defer upstream.Close()
u, _ := url.Parse(upstream.URL)
filter, err := config.NewFilter(config.FilterDisabled, nil)
if err != nil {
t.Fatal(err)
}
h := proxy.New(u, filter, noopSubmitter{}, 1024)
proxySrv := httptest.NewServer(h)
defer proxySrv.Close()
pu, _ := url.Parse(proxySrv.URL)
conn, err := net.DialTimeout("tcp", pu.Host, time.Second)
if err != nil {
t.Fatal(err)
}
_, _ = io.WriteString(conn, "GET /v1/realtime HTTP/1.1\r\nHost: "+pu.Host+"\r\nUpgrade: websocket\r\nConnection: Upgrade\r\n\r\n")
resp, err := http.ReadResponse(bufio.NewReader(conn), nil)
if err != nil {
t.Fatal(err)
}
if resp.StatusCode != http.StatusSwitchingProtocols {
t.Fatalf("status=%d", resp.StatusCode)
}
if err := conn.Close(); err != nil {
t.Fatal(err)
}
deadline := time.Now().Add(time.Second)
for {
ctx, cancel := context.WithCancel(context.Background())
cancel()
err := h.Shutdown(ctx)
if err == nil {
break
}
if time.Now().After(deadline) {
t.Fatalf("connection was not unregistered after natural close: %v", err)
}
time.Sleep(10 * time.Millisecond)
}
}
func TestNonWebSocketUpgradeIsManagedButNotLogged(t *testing.T) {
upstream := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
hj := w.(http.Hijacker)
conn, brw, err := hj.Hijack()
if err != nil {
t.Errorf("upstream hijack: %v", err)
return
}
defer conn.Close()
_, _ = brw.WriteString("HTTP/1.1 101 Switching Protocols\r\nUpgrade: test-protocol\r\nConnection: Upgrade\r\n\r\n")
_ = brw.Flush()
_, _ = io.Copy(io.Discard, conn)
}))
defer upstream.Close()
u, _ := url.Parse(upstream.URL)
filter, err := config.NewFilter(config.FilterDisabled, nil)
if err != nil {
t.Fatal(err)
}
entries := make(chanSubmitter, 1)
h := proxy.New(u, filter, entries, 1024)
proxySrv := httptest.NewServer(h)
defer proxySrv.Close()
pu, _ := url.Parse(proxySrv.URL)
conn, err := net.DialTimeout("tcp", pu.Host, time.Second)
if err != nil {
t.Fatal(err)
}
defer conn.Close()
_, _ = io.WriteString(conn, "GET /upgrade HTTP/1.1\r\nHost: "+pu.Host+"\r\nUpgrade: test-protocol\r\nConnection: Upgrade\r\n\r\n")
resp, err := http.ReadResponse(bufio.NewReader(conn), nil)
if err != nil {
t.Fatal(err)
}
if resp.StatusCode != http.StatusSwitchingProtocols {
t.Fatalf("status=%d", resp.StatusCode)
}
select {
case entry := <-entries:
t.Fatalf("non-WebSocket upgrade was logged: %+v", entry)
case <-time.After(50 * time.Millisecond):
}
ctx, cancel := context.WithTimeout(context.Background(), time.Second)
defer cancel()
if err := h.Shutdown(ctx); err != nil {
t.Fatal(err)
}
}
// TestWebSocketProxyPassthrough 验证 Upgrade 请求能正确透传，
// 确认 captureWriter 的 Hijacker 实现没破坏 ReverseProxy 的 WS 行为。
func TestWebSocketProxyPassthrough(t *testing.T) {
upstream := fakeUpstream(t)
defer upstream.Close()
u, err := url.Parse(upstream.URL)
if err != nil {
t.Fatal(err)
}
// disabled 模式下 ShouldLog 返回 true（全量记录），会进入捕获分支。
filter, err := config.NewFilter(config.FilterDisabled, nil)
if err != nil {
t.Fatal(err)
}
h := proxy.New(u, filter, noopSubmitter{}, 1024)
proxySrv := httptest.NewServer(h)
defer proxySrv.Close()
// 建立到代理的 TCP 连接，手写 Upgrade 请求
pu, _ := url.Parse(proxySrv.URL)
d := net.Dialer{Timeout: 3 * time.Second}
conn, err := d.DialContext(context.Background(), "tcp", pu.Host)
if err != nil {
t.Fatal(err)
}
defer conn.Close()
req := "GET /v1/realtime HTTP/1.1\r\n" +
"Host: " + pu.Host + "\r\n" +
"Upgrade: websocket\r\n" +
"Connection: Upgrade\r\n" +
"Sec-WebSocket-Key: dGhlIHNhbXBsZSBub25jZQ==\r\n" +
"Sec-WebSocket-Version: 13\r\n" +
"\r\n"
if _, err := io.WriteString(conn, req); err != nil {
t.Fatal(err)
}
_ = conn.SetReadDeadline(time.Now().Add(3 * time.Second))
br := bufio.NewReader(conn)
resp, err := http.ReadResponse(br, nil)
if err != nil {
t.Fatalf("read upgrade response: %v", err)
}
if resp.StatusCode != http.StatusSwitchingProtocols {
t.Fatalf("expected 101, got %d", resp.StatusCode)
}
if !strings.EqualFold(resp.Header.Get("Upgrade"), "websocket") {
t.Fatalf("expected Upgrade: websocket, got %q", resp.Header.Get("Upgrade"))
}
// echo 测试
payload := "hello-websocket"
if _, err := io.WriteString(conn, payload); err != nil {
t.Fatal(err)
}
got := make([]byte, len(payload))
if _, err := io.ReadFull(br, got); err != nil {
t.Fatalf("read echo: %v", err)
}
if string(got) != payload {
t.Fatalf("echo mismatch: got %q want %q", got, payload)
}
}