修复代理与日志链路的可靠性问题

This commit is contained in:
2026-07-11 17:23:32 +08:00
parent d39f5a1048
commit 9e6e7ca3c7
21 changed files with 417 additions and 83 deletions
+32
View File
@@ -108,6 +108,38 @@ func TestRequestBodyReadFailureReturnsFixed400WithoutUpstream(t *testing.T) {
}
}
func TestRequestBodyOverLimitReturns413WithoutUpstream(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{MaxRequestBytes: 4})
for _, tc := range []struct {
name string
contentLength int64
}{
{name: "known length", contentLength: 5},
{name: "chunked", contentLength: -1},
} {
t.Run(tc.name, func(t *testing.T) {
req := httptest.NewRequest(http.MethodPost, "http://proxy.test/v1/chat", strings.NewReader("12345"))
req.ContentLength = tc.contentLength
rec := httptest.NewRecorder()
h.ServeHTTP(rec, req)
if rec.Code != http.StatusRequestEntityTooLarge || rec.Body.String() != "request body too large\n" {
t.Fatalf("response=(%d, %q), want fixed 413", rec.Code, rec.Body.String())
}
})
}
if upstreamCalls.Load() != 0 {
t.Fatalf("upstream called %d times", upstreamCalls.Load())
}
}
func TestRequestBodyReadFailureWithZeroContentLengthDoesNotReachUpstream(t *testing.T) {
var upstreamCalls atomic.Int32
upstream := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+52
View File
@@ -643,6 +643,58 @@ func TestGeminiStreamPreservesFunctionCallParts(t *testing.T) {
}
}
func TestGeminiStreamPreservesMultiplePartsByIndex(t *testing.T) {
upstream := newSSEUpstream([]string{
`data: {"candidates":[{"content":{"parts":[{"text":"hello "},{"functionCall":{"name":"lookup","args":{"q":"weather"}}},{"text":"world"}],"role":"model"},"index":0}]}` + "\n\n",
`data: {"candidates":[{"content":{"parts":[{"text":"again"},{"functionCall":{"name":"lookup","args":{"q":"weather"}}},{"text":"!"}],"role":"model"},"finishReason":"STOP","index":0}]}` + "\n\n",
})
e := requestStreamEntry(t, upstream, "/v1beta/models/gemini-1.5-pro:streamGenerateContent")
var captured struct {
Candidates []struct {
Content struct {
Parts []map[string]any `json:"parts"`
} `json:"content"`
} `json:"candidates"`
}
if err := json.Unmarshal(e.ResponseBody, &captured); err != nil {
t.Fatalf("unmarshal assembled Gemini response: %v; body=%s", err, e.ResponseBody)
}
parts := captured.Candidates[0].Content.Parts
if len(parts) != 3 || parts[0]["text"] != "hello again" || parts[2]["text"] != "world!" {
t.Fatalf("multiple Gemini parts not preserved: %+v", parts)
}
if _, ok := parts[1]["functionCall"]; !ok {
t.Fatalf("middle functionCall part missing: %+v", parts)
}
}
func TestGeminiStreamAppendsDifferentPartKindsAcrossChunks(t *testing.T) {
upstream := newSSEUpstream([]string{
`data: {"candidates":[{"content":{"parts":[{"text":"answer"}],"role":"model"},"index":0}]}` + "\n\n",
`data: {"candidates":[{"content":{"parts":[{"functionCall":{"name":"lookup","args":{"q":"weather"}}}],"role":"model"},"finishReason":"STOP","index":0}]}` + "\n\n",
})
e := requestStreamEntry(t, upstream, "/v1beta/models/gemini-1.5-pro:streamGenerateContent")
var captured struct {
Candidates []struct {
Content struct {
Parts []map[string]any `json:"parts"`
} `json:"content"`
} `json:"candidates"`
}
if err := json.Unmarshal(e.ResponseBody, &captured); err != nil {
t.Fatalf("unmarshal assembled Gemini response: %v; body=%s", err, e.ResponseBody)
}
parts := captured.Candidates[0].Content.Parts
if len(parts) != 2 || parts[0]["text"] != "answer" {
t.Fatalf("different Gemini part kinds were merged: %+v", parts)
}
if _, ok := parts[1]["functionCall"]; !ok {
t.Fatalf("functionCall part missing: %+v", parts)
}
}
func TestUnknownSSEKeepsRawBody(t *testing.T) {
e := requestStreamEntry(t, fakeUnknownStreamUpstream(), "/v1/chat/completions")