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

.env.example

@@ -11,6 +11,8 @@ CLICKHOUSE_URL=
 
 # 单条请求/响应 body 最大记录字节数
 MAX_BODY_BYTES=1048576
+# 单个代理请求体硬上限，超出时返回 413
+MAX_REQUEST_BYTES=16777216
 
 # 异步日志队列容量
 LOG_QUEUE_SIZE=256

README.md

@@ -79,6 +79,7 @@ go vet ./...
 | `UPSTREAM_TLS_INSECURE_SKIP_VERIFY` | 跳过上游 HTTPS 证书校验；仅开发/可信内网自签证书场景使用 | `false` |
 | `CLICKHOUSE_URL` | ClickHouse 原生协议地址（必填）；支持严格 `host:port`、`clickhouse://` 或 TLS `clickhouses://` | - |
 | `MAX_BODY_BYTES` | 单个请求体和响应体的记录上限 | `1048576` |
+| `MAX_REQUEST_BYTES` | 单个代理请求体硬上限，超出时返回 413 | `16777216` |
 | `LOG_QUEUE_SIZE` | 异步队列条数上限 | `256` |
 | `LOG_QUEUE_BYTES` | 异步队列总字节预算 | `67108864` |
 | `LOG_BATCH_SIZE` | 批量写入条数 | `50` |
@@ -164,7 +165,7 @@ ORDER BY (started_at, request_id);
 | `client_ip` | `String` | 否 | 默认记录 TCP 对端；仅直接对端命中 `TRUSTED_PROXIES` 时，从右向左解析 `X-Forwarded-For`，无 XFF 时使用有效的 `X-Real-IP` |
 | `request_headers` | `String` | 否 | 完整请求头，序列化为 `{"Header-Name": ["value1", "value2"], ...}` 的 JSON 对象；**注意 Authorization、Cookie、API-Key 等敏感头未脱敏**，按设计原样存储 |
 | `request_body` | `String` | 否 | 请求体内容。最多保留 `MAX_BODY_BYTES`（默认 1 MiB）字节，超出部分丢弃；读取失败时请求不会转发 |
-| `request_truncated` | `Bool` | 否 | 请求体是否被 `MAX_BODY_BYTES` 截断。截断只影响数据库记录；代理会将已完整读入内存的请求体通过 `bytes.Reader` 重放给下游 newapi |
+| `request_truncated` | `Bool` | 否 | 请求体是否被 `MAX_BODY_BYTES` 截断。截断只影响数据库记录；请求体超过 `MAX_REQUEST_BYTES` 时不会转发并返回 413 |
 | `status_code` | `Int32` | 否 | 上游返回的 HTTP 状态码。`502` 通常意味着上游连接失败；WebSocket 成功升级记录为 `101` |
 | `response_headers` | `String` | 否 | 响应头，结构同 `request_headers`。对于 SSE，会包含 `Content-Type: text/event-stream` 等 |
 | `response_body` | `String` | 否 | 响应体内容。对于未截断的 SSE，代理会尝试将 OpenAI Completions、Chat Completions、Responses、Anthropic 或 Gemini 事件组装为单个 JSON，并以该 JSON 替换原始 SSE 字节流；无法识别或组装失败时保留原始 SSE。超过 `MAX_BODY_BYTES` 时不组装，仅保留原始字节流的前 `MAX_BODY_BYTES` 字节 |
@@ -219,7 +220,7 @@ SELECT JSONExtractRaw(request_headers, 'Authorization') FROM proxy_logs LIMIT 5;
 
 ## 设计要点
 
-- 转发前使用 `io.ReadAll` 将请求体完整读入内存，读取或关闭失败时不向上游发送请求；数据库仅记录前 `MAX_BODY_BYTES` 字节，随后通过 `bytes.Reader` 重放完整请求体。因此 `MAX_BODY_BYTES` 只限制日志字段大小，不限制代理读取请求体时的内存占用。
+- 转发前最多读取 `MAX_REQUEST_BYTES + 1` 字节，超限返回 413，读取或关闭失败时不向上游发送请求；数据库仅记录前 `MAX_BODY_BYTES` 字节，随后通过 `bytes.Reader` 重放完整请求体。
 - `httputil.ReverseProxy` + `FlushInterval = -1`，自定义 `ResponseWriter` 同时实现 `Flusher`/`Hijacker`，写入时先转发再缓冲，保证流式实时性。
 - 日志通过非阻塞 channel 投递，队列满或 DB 不健康时直接丢弃（每 30 秒打印 metrics）。
 - DB 健康状态机：写入失败立即标记 unhealthy，后台 ping 恢复后重新启用。

compose.yml

@@ -17,6 +17,7 @@ services:
       UPSTREAM_TLS_INSECURE_SKIP_VERIFY: "${UPSTREAM_TLS_INSECURE_SKIP_VERIFY:-false}"
       CLICKHOUSE_URL: "${CLICKHOUSE_URL:?set CLICKHOUSE_URL with URL-encoded credentials}"
       MAX_BODY_BYTES: "${MAX_BODY_BYTES:-1048576}"
+      MAX_REQUEST_BYTES: "${MAX_REQUEST_BYTES:-16777216}"
       LOG_QUEUE_SIZE: "${LOG_QUEUE_SIZE:-256}"
       LOG_QUEUE_BYTES: "${LOG_QUEUE_BYTES:-67108864}"
       LOG_BATCH_SIZE: "${LOG_BATCH_SIZE:-50}"

config/config.go

@@ -3,6 +3,7 @@ package config
 import (
 	"errors"
 	"fmt"
+	"math"
 	"net/netip"
 	"net/url"
 	"os"
@@ -19,6 +20,7 @@ type Config struct {
 	UpstreamURL                   *url.URL
 	ClickHouseURL                 string
 	MaxBodyBytes                  int64
+	MaxRequestBytes               int64
 	LogQueueSize                  int
 	LogQueueBytes                 int64
 	LogBatchSize                  int
@@ -42,6 +44,10 @@ func Load() (*Config, error) {
 	if err != nil {
 		return nil, err
 	}
+	maxRequestBytes, err := getEnvInt64("MAX_REQUEST_BYTES", 16<<20)
+	if err != nil {
+		return nil, err
+	}
 	logQueueSize, err := getEnvInt("LOG_QUEUE_SIZE", 256)
 	if err != nil {
 		return nil, err
@@ -103,6 +109,7 @@ func Load() (*Config, error) {
 		ListenAddr:                    getEnv("LISTEN_ADDR", ":8080"),
 		ClickHouseURL:                 os.Getenv("CLICKHOUSE_URL"),
 		MaxBodyBytes:                  maxBodyBytes,
+		MaxRequestBytes:               maxRequestBytes,
 		LogQueueSize:                  logQueueSize,
 		LogQueueBytes:                 logQueueBytes,
 		LogBatchSize:                  logBatchSize,
@@ -128,7 +135,7 @@ func Load() (*Config, error) {
 	if err != nil {
 		return nil, fmt.Errorf("invalid UPSTREAM_URL: %w", err)
 	}
-	if u.Scheme == "" || u.Host == "" {
+	if (u.Scheme != "http" && u.Scheme != "https") || u.Host == "" {
 		return nil, fmt.Errorf("invalid UPSTREAM_URL: %q", upstream)
 	}
 	cfg.UpstreamURL = u
@@ -144,6 +151,7 @@ func Load() (*Config, error) {
 		value int64
 	}{
 		{key: "MAX_BODY_BYTES", value: cfg.MaxBodyBytes},
+		{key: "MAX_REQUEST_BYTES", value: cfg.MaxRequestBytes},
 		{key: "LOG_QUEUE_SIZE", value: int64(cfg.LogQueueSize)},
 		{key: "LOG_QUEUE_BYTES", value: cfg.LogQueueBytes},
 		{key: "LOG_BATCH_SIZE", value: int64(cfg.LogBatchSize)},
@@ -162,6 +170,9 @@ func Load() (*Config, error) {
 			return nil, fmt.Errorf("%s must be greater than zero", item.key)
 		}
 	}
+	if cfg.MaxRequestBytes == math.MaxInt64 {
+		return nil, errors.New("MAX_REQUEST_BYTES is too large")
+	}
 	return cfg, nil
 }
 

db/clickhouse.go

@@ -22,6 +22,7 @@ type Pool struct {
 	open              func(*clickhouse.Options) (clickhouse.Conn, error)
 	mu                sync.RWMutex
 	conn              clickhouse.Conn
+	generation        uint64
 	healthy           bool
 	closed            bool
 }
@@ -82,6 +83,7 @@ func (p *Pool) connect(ctx context.Context) error {
 		_ = p.conn.Close()
 	}
 	p.conn = conn
+	p.generation++
 	p.healthy = true
 	p.mu.Unlock()
 	log.Printf("[db] connected and migrated")
@@ -192,11 +194,11 @@ func (p *Pool) watch(ctx context.Context) {
 			return
 		case <-t.C:
 			if p.Healthy() {
-				if conn := p.Get(); conn != nil {
+				if conn, generation := p.GetWithGeneration(); conn != nil {
 					pctx, cancel := context.WithTimeout(ctx, 3*time.Second)
 					if err := conn.Ping(pctx); err != nil {
 						log.Printf("[db] ping failed, marking unhealthy: %v", err)
-						p.MarkUnhealthy()
+						p.MarkUnhealthyGeneration(generation)
 					}
 					cancel()
 				}
@@ -223,6 +225,21 @@ func (p *Pool) MarkUnhealthy() {
 	p.mu.Unlock()
 }
 
+// MarkUnhealthyGeneration only changes health when the failed connection is still current.
+func (p *Pool) MarkUnhealthyGeneration(generation uint64) {
+	p.mu.Lock()
+	if p.generation == generation {
+		p.healthy = false
+	}
+	p.mu.Unlock()
+}
+
+func (p *Pool) Generation() uint64 {
+	p.mu.RLock()
+	defer p.mu.RUnlock()
+	return p.generation
+}
+
 // Get 返回当前连接，可能为 nil。
 func (p *Pool) Get() clickhouse.Conn {
 	p.mu.RLock()
@@ -230,6 +247,12 @@ func (p *Pool) Get() clickhouse.Conn {
 	return p.conn
 }
 
+func (p *Pool) GetWithGeneration() (clickhouse.Conn, uint64) {
+	p.mu.RLock()
+	defer p.mu.RUnlock()
+	return p.conn, p.generation
+}
+
 // Close starts closing the underlying connection and waits within ctx.
 func (p *Pool) Close(ctx context.Context) error {
 	p.mu.Lock()

db/clickhouse_test.go

@@ -214,6 +214,26 @@ func TestClosePreventsConcurrentConnectFromPublishing(t *testing.T) {
 	}
 }
 
+func TestOldGenerationCannotMarkReplacementUnhealthy(t *testing.T) {
+	pool := &Pool{conn: newFakeClickHouseConn(), healthy: true, generation: 1}
+	oldGeneration := pool.Generation()
+
+	pool.mu.Lock()
+	pool.conn = newFakeClickHouseConn()
+	pool.generation++
+	pool.healthy = true
+	pool.mu.Unlock()
+
+	pool.MarkUnhealthyGeneration(oldGeneration)
+	if !pool.Healthy() {
+		t.Fatal("old connection failure marked replacement connection unhealthy")
+	}
+	pool.MarkUnhealthyGeneration(pool.Generation())
+	if pool.Healthy() {
+		t.Fatal("current connection failure did not mark pool unhealthy")
+	}
+}
+
 func assertNonDestructiveMigration(t *testing.T, statements []string) {
 	t.Helper()
 	if len(statements) == 0 || !strings.Contains(strings.ToUpper(statements[0]), "CREATE TABLE IF NOT EXISTS") {

docs/compose/reports/reliability-security-fixes.md

@@ -1,6 +1,6 @@
 ---
 feature: reliability-security-fixes
-status: delivered
+status: complete
 specs:
   - docs/compose/specs/reliability-security-fixes.md
   - docs/compose/specs/2026-07-09-clickhouse-migration.md
@@ -14,7 +14,7 @@ branch: main
 
 ## What Was Built
 
-本轮完成了 token_thief 的可靠性与安全加固。异步日志队列现在安全处理并发 `Submit`/`Stop`、使用条目数和字节双预算、在统一 shutdown deadline 内排空，并区分 ClickHouse 的可安全重试、确定失败和提交结果不明三类写入结果。
+本轮完成了 token_thief 的主要可靠性与安全加固。异步日志队列现在安全处理并发 `Submit`/`Stop`、使用条目数和字节双预算、在统一 shutdown deadline 内排空，并区分 ClickHouse 的可安全重试、确定失败和提交结果不明三类写入结果。
 
 反向代理现在对请求体读取、上游响应复制、普通响应超时、SSE idle timeout、WebSocket 101 元数据和升级连接关闭实施完整性保护。客户端错误文本不再泄露内部信息，转发来源头仅在直接对端属于 `TRUSTED_PROXIES` 时参与客户端 IP 判定。
 
@@ -37,7 +37,7 @@ branch: main
 
 ## Verification
 
-最终验收从工作区根目录运行 `gofmt -l .`、`go test -count=1 ./...`、`go vet ./...` 和 `go build ./`。报告保留可复现命令而不固化会随测试增删失真的测试或包计数；针对 config、queue、ClickHouse、HTTP/SSE/WebSocket 和 Compose 的失败路径均有测试覆盖。
+最终验收从工作区根目录运行 `gofmt -w .`、`go vet ./...`、`go build ./...` 和 `go test -count=1 -timeout 3m ./...`，命令均以 0 退出。另对代理测试执行 20 轮重复验证，对 logger 队列测试执行 20 轮重复验证，均通过。WebSocket hijack 后的普通 HTTP flush 已由 writer 层 guard 阻断，未再出现 server recover panic。
 
 ## Journey Log
 
@@ -47,6 +47,8 @@ branch: main
 - [lesson] 迭代 1：进程内队列只能提供有界 best-effort；跨进程幂等需要持久化 outbox 和下游幂等协议，不能由本地重试可靠模拟。
 - [pivot] 迭代 1：SSE 完成判定限定为真正的 `text/event-stream` 且等待代理正常返回，避免终止标记导致提前记录不完整响应。
 - [lesson] 迭代 4：最终报告记录禁用测试缓存的验证命令而非易过期的精确计数，并同时执行格式、静态检查和构建。
+- [finding] 迭代 5：格式、vet、构建和测试均退出成功，但 passing test 中仍可隐藏由 `net/http` recover 的 handler panic；WebSocket hijack 后不得再刷新普通 HTTP response writer。
+- [fix] 最终验收：增加 `MAX_REQUEST_BYTES` 硬上限、上游连接写 deadline、Queue deadline 后有界清理，以及 hijack 后 flush guard；代理和队列压力测试各连续 20 轮通过。
 
 ## Source Materials
 

logger/queue.go

@@ -8,6 +8,7 @@ import (
 	"sync"
 	"sync/atomic"
 	"time"
+	"unsafe"
 
 	"github.com/ClickHouse/clickhouse-go/v2/lib/driver"
 
@@ -47,10 +48,15 @@ func (p clickHouseBatchPreparer) PrepareBatch(ctx context.Context, query string)
 
 type poolBackend struct {
 	pool       *db.Pool
+	conn       driver.Conn
+	generation uint64
 }
 
 func (b poolBackend) PrepareBatch(ctx context.Context, query string) (Batch, error) {
-	conn := b.pool.Get()
+	conn := b.conn
+	if conn == nil {
+		conn, _ = b.pool.GetWithGeneration()
+	}
 	if conn == nil {
 		return nil, errors.New("pool nil")
 	}
@@ -58,7 +64,14 @@ func (b poolBackend) PrepareBatch(ctx context.Context, query string) (Batch, err
 }
 
 func (b poolBackend) Healthy() bool  { return b.pool.Healthy() }
-func (b poolBackend) MarkUnhealthy() { b.pool.MarkUnhealthy() }
+func (b poolBackend) MarkUnhealthy() { b.pool.MarkUnhealthyGeneration(b.generation) }
+func (b poolBackend) Snapshot() (Backend, uint64) {
+	conn, generation := b.pool.GetWithGeneration()
+	return poolBackend{pool: b.pool, conn: conn, generation: generation}, generation
+}
+func (b poolBackend) MarkUnhealthyGeneration(generation uint64) {
+	b.pool.MarkUnhealthyGeneration(generation)
+}
 
 type Stats struct {
 	Enqueued  uint64
@@ -153,8 +166,8 @@ func EstimatedBytes(e *LogEntry) int64 {
 	if e == nil {
 		return 0
 	}
-	return int64(len(e.RequestID) + len(e.Method) + len(e.Path) + len(e.Query) + len(e.ClientIP) + len(e.Error) +
-		len(e.RequestHeaders) + len(e.RequestBody) + len(e.ResponseHeaders) + len(e.ResponseBody))
+	return int64(unsafe.Sizeof(*e)+unsafe.Sizeof(queuedEntry{})) + int64(len(e.RequestID)+len(e.Method)+len(e.Path)+len(e.Query)+len(e.ClientIP)+len(e.Error)+
+		len(e.RequestHeaders)+len(e.RequestBody)+len(e.ResponseHeaders)+len(e.ResponseBody))
 }
 
 func (q *Queue) Stats() Stats {
@@ -195,7 +208,6 @@ func (q *Queue) Stop(contexts ...context.Context) error {
 	if len(contexts) > 0 && contexts[0] != nil {
 		ctx = contexts[0]
 	}
-
 	q.mu.Lock()
 	shutdownOwner := false
 	if !q.stopped {
@@ -226,7 +238,12 @@ func (q *Queue) Stop(contexts ...context.Context) error {
 	case <-ctx.Done():
 		if shutdownOwner && workCancel != nil {
 			workCancel()
-			<-done
+			timer := time.NewTimer(100 * time.Millisecond)
+			defer timer.Stop()
+			select {
+			case <-done:
+			case <-timer.C:
+			}
 		}
 		return ctx.Err()
 	}
@@ -238,8 +255,7 @@ func (q *Queue) Submit(e *LogEntry) {
 		q.dropped.Add(1)
 		return
 	}
-	entry := cloneLogEntry(e)
-	size := EstimatedBytes(entry)
+	size := EstimatedBytes(e)
 	if !q.mu.TryRLock() {
 		q.dropped.Add(1)
 		return
@@ -249,6 +265,7 @@ func (q *Queue) Submit(e *LogEntry) {
 		q.dropped.Add(1)
 		return
 	}
+	entry := cloneLogEntry(e)
 	select {
 	case q.ch <- queuedEntry{entry: entry, size: size}:
 		q.enq.Add(1)
@@ -367,16 +384,19 @@ func (q *Queue) flush(ctx context.Context, entries []*LogEntry) {
 
 	retry := entries
 	var lastErr error
+	var failedGeneration uint64
 	for attempt := 1; attempt <= maxAttempts && len(retry) > 0; attempt++ {
 		if err := ctx.Err(); err != nil {
 			lastErr = err
 			break
 		}
-		result := Flush(ctx, q.backend, retry)
+		attemptBackend, generation := backendSnapshot(q.backend)
+		failedGeneration = generation
+		result := Flush(ctx, attemptBackend, retry)
 		q.failed.Add(uint64(result.Failed))
 		q.ambiguous.Add(uint64(result.Ambiguous))
 		if result.Ambiguous > 0 {
-			q.backend.MarkUnhealthy()
+			markBackendUnhealthy(q.backend, failedGeneration)
 		}
 		lastErr = result.Err
 		retry = result.Retry
@@ -391,11 +411,31 @@ func (q *Queue) flush(ctx context.Context, entries []*LogEntry) {
 
 	if len(retry) > 0 {
 		q.failed.Add(uint64(len(retry)))
-		q.backend.MarkUnhealthy()
+		markBackendUnhealthy(q.backend, failedGeneration)
 		log.Printf("[logger] giving up %d retry-safe rows after %d attempts: %v", len(retry), maxAttempts, lastErr)
 	}
 }
 
+type generationBackend interface {
+	Snapshot() (Backend, uint64)
+	MarkUnhealthyGeneration(uint64)
+}
+
+func backendSnapshot(backend Backend) (Backend, uint64) {
+	if versioned, ok := backend.(generationBackend); ok {
+		return versioned.Snapshot()
+	}
+	return backend, 0
+}
+
+func markBackendUnhealthy(backend Backend, generation uint64) {
+	if versioned, ok := backend.(generationBackend); ok {
+		versioned.MarkUnhealthyGeneration(generation)
+		return
+	}
+	backend.MarkUnhealthy()
+}
+
 func waitBackoff(ctx context.Context, attempt int) bool {
 	wait := BaseBackoff
 	for i := 1; i < attempt; i++ {

main.go

@@ -30,7 +30,7 @@ type shutdownStep struct {
 	run  func(context.Context) error
 }
 
-func waitForShutdown(stop <-chan os.Signal, serverErr <-chan error) error {
+func waitForShutdown(stop <-chan struct{}, serverErr <-chan error) error {
 	select {
 	case <-stop:
 		log.Printf("[main] shutdown signal received")
@@ -83,11 +83,9 @@ func run() error {
 	}
 	log.Printf("[main] filter mode=%s patterns=%d", filter.Mode, len(filter.Patterns))
 
-	stop := make(chan os.Signal, 1)
-	signal.Notify(stop, syscall.SIGINT, syscall.SIGTERM)
-	defer signal.Stop(stop)
-
-	rootCtx, cancel := context.WithCancel(context.Background())
+	rootCtx, stopSignals := signal.NotifyContext(context.Background(), syscall.SIGINT, syscall.SIGTERM)
+	defer stopSignals()
+	rootCtx, cancel := context.WithCancel(rootCtx)
 	defer cancel()
 
 	pool := db.NewPool(rootCtx, cfg.ClickHouseURL, cfg.DBReconnectInterval)
@@ -101,6 +99,7 @@ func run() error {
 		SSEIdleTimeout:                cfg.UpstreamStreamIdleTimeout,
 		UpstreamTLSInsecureSkipVerify: cfg.UpstreamTLSInsecureSkipVerify,
 		TrustedProxies:                cfg.TrustedProxies,
+		MaxRequestBytes:               cfg.MaxRequestBytes,
 	})
 
 	srv := &http.Server{
@@ -118,7 +117,7 @@ func run() error {
 		serverErr <- srv.ListenAndServe()
 	}()
 
-	runErr := waitForShutdown(stop, serverErr)
+	runErr := waitForShutdown(rootCtx.Done(), serverErr)
 
 	// HTTP、升级连接、日志排空和数据库共享同一个关闭总预算。
 	shutdownCtx, shutdownCancel := context.WithTimeout(context.Background(), 30*time.Second)

main_test.go

@@ -7,7 +7,6 @@ import (
 	"go/parser"
 	"go/token"
 	"net/http"
-	"os"
 	"reflect"
 	"strings"
 	"testing"
@@ -49,13 +48,13 @@ func TestRunRegistersSignalsBeforeStartingResources(t *testing.T) {
 		}
 		name := owner.Name + "." + selector.Sel.Name
 		switch name {
-		case "signal.Notify", "db.NewPool", "queue.Start", "srv.ListenAndServe":
+		case "signal.NotifyContext", "db.NewPool", "queue.Start", "srv.ListenAndServe":
 			positions[name] = call.Pos()
 		}
 		return true
 	})
 
-	notifyPos, ok := positions["signal.Notify"]
+	notifyPos, ok := positions["signal.NotifyContext"]
 	if !ok {
 		t.Fatal("run does not register for shutdown signals")
 	}
@@ -65,7 +64,7 @@ func TestRunRegistersSignalsBeforeStartingResources(t *testing.T) {
 			t.Fatalf("run does not call %s", start)
 		}
 		if notifyPos >= startPos {
-			t.Errorf("signal.Notify at line %d must precede %s at line %d",
+			t.Errorf("signal.NotifyContext at line %d must precede %s at line %d",
 				fset.Position(notifyPos).Line, start, fset.Position(startPos).Line)
 		}
 	}
@@ -76,7 +75,7 @@ func TestWaitForShutdownReturnsListenerError(t *testing.T) {
 	serverErr := make(chan error, 1)
 	serverErr <- listenErr
 
-	err := waitForShutdown(make(chan os.Signal), serverErr)
+	err := waitForShutdown(make(chan struct{}), serverErr)
 	if !errors.Is(err, listenErr) || !strings.Contains(err.Error(), "server") {
 		t.Fatalf("waitForShutdown() error = %v, want wrapped listener error", err)
 	}
@@ -84,8 +83,8 @@ func TestWaitForShutdownReturnsListenerError(t *testing.T) {
 
 func TestWaitForShutdownAcceptsSignalAndServerClosed(t *testing.T) {
 	t.Run("signal", func(t *testing.T) {
-		stop := make(chan os.Signal, 1)
-		stop <- os.Interrupt
+		stop := make(chan struct{})
+		close(stop)
 		if err := waitForShutdown(stop, make(chan error)); err != nil {
 			t.Fatalf("waitForShutdown() error = %v, want nil", err)
 		}
@@ -94,7 +93,7 @@ func TestWaitForShutdownAcceptsSignalAndServerClosed(t *testing.T) {
 	t.Run("server closed", func(t *testing.T) {
 		serverErr := make(chan error, 1)
 		serverErr <- http.ErrServerClosed
-		if err := waitForShutdown(make(chan os.Signal), serverErr); err != nil {
+		if err := waitForShutdown(make(chan struct{}), serverErr); err != nil {
 			t.Fatalf("waitForShutdown() error = %v, want nil", err)
 		}
 	})

proxy/capture.go

@@ -5,21 +5,30 @@ import (
 	"crypto/rand"
 	"encoding/hex"
 	"encoding/json"
+	"errors"
 	"io"
 	"net/http"
 	"net/netip"
 	"strings"
 )
 
+var errRequestBodyTooLarge = errors.New("request body too large")
+
 // readRequestBody 在转发前完整读取请求体，确保读取失败时不会向上游发送损坏请求。
-func readRequestBody(r *http.Request, max int64) (captured []byte, truncated bool, err error) {
+func readRequestBody(r *http.Request, max, requestLimit int64) (captured []byte, truncated bool, err error) {
 	if r.Body == nil || r.Body == http.NoBody {
 		return nil, false, nil
 	}
-	body, err := io.ReadAll(r.Body)
+	if r.ContentLength > requestLimit {
+		return nil, false, errRequestBodyTooLarge
+	}
+	body, err := io.ReadAll(io.LimitReader(r.Body, requestLimit+1))
 	if err != nil {
 		return nil, false, err
 	}
+	if int64(len(body)) > requestLimit {
+		return nil, false, errRequestBodyTooLarge
+	}
 	if err := r.Body.Close(); err != nil {
 		return nil, false, err
 	}

proxy/proxy.go

@@ -3,6 +3,7 @@ package proxy
 import (
 	"context"
 	"crypto/tls"
+	"errors"
 	"log"
 	"net"
 	"net/http"
@@ -29,6 +30,7 @@ type Handler struct {
 	filter          *config.Filter
 	queue           LogSubmitter
 	maxBodyBytes    int64
+	maxRequestBytes int64
 	trusted         []netip.Prefix
 	connMu          sync.Mutex
 	conns           map[net.Conn]struct{}
@@ -44,6 +46,7 @@ type Options struct {
 	SSEIdleTimeout                time.Duration
 	UpstreamTLSInsecureSkipVerify bool
 	TrustedProxies                []netip.Prefix
+	MaxRequestBytes               int64
 }
 
 // requestState 通过 context 在 ErrorHandler / ModifyResponse / 主 handler 之间共享状态。
@@ -77,12 +80,22 @@ func New(upstream *url.URL, filter *config.Filter, queue LogSubmitter, maxBody i
 }
 
 func NewWithOptions(upstream *url.URL, filter *config.Filter, queue LogSubmitter, maxBody int64, opts Options) *Handler {
+	if opts.MaxRequestBytes <= 0 {
+		opts.MaxRequestBytes = 16 << 20
+	}
 	rp := httputil.NewSingleHostReverseProxy(upstream)
 	rp.FlushInterval = -1 // 让流式 chunk 立即转发
 	if opts.UpstreamTimeout > 0 || opts.UpstreamTLSInsecureSkipVerify {
 		transport := http.DefaultTransport.(*http.Transport).Clone()
 		if opts.UpstreamTimeout > 0 {
-			transport.DialContext = (&net.Dialer{Timeout: opts.UpstreamTimeout, KeepAlive: 30 * time.Second}).DialContext
+			dialer := &net.Dialer{Timeout: opts.UpstreamTimeout, KeepAlive: 30 * time.Second}
+			transport.DialContext = func(ctx context.Context, network, address string) (net.Conn, error) {
+				conn, err := dialer.DialContext(ctx, network, address)
+				if err != nil {
+					return nil, err
+				}
+				return &writeTimeoutConn{Conn: conn, timeout: opts.UpstreamTimeout}, nil
+			}
 			transport.ResponseHeaderTimeout = opts.UpstreamTimeout
 			transport.TLSHandshakeTimeout = opts.UpstreamTimeout
 		}
@@ -134,6 +147,7 @@ func NewWithOptions(upstream *url.URL, filter *config.Filter, queue LogSubmitter
 		filter:          filter,
 		queue:           queue,
 		maxBodyBytes:    maxBody,
+		maxRequestBytes: opts.MaxRequestBytes,
 		trusted:         append([]netip.Prefix(nil), opts.TrustedProxies...),
 		conns:           make(map[net.Conn]struct{}),
 		connChanged:     make(chan struct{}),
@@ -149,11 +163,15 @@ func (h *Handler) ServeHTTP(w http.ResponseWriter, r *http.Request) {
 	}
 
 	shouldLog := h.filter.ShouldLog(r.URL.Path)
-	reqBody, reqTruncated, err := readRequestBody(r, h.maxBodyBytes)
+	reqBody, reqTruncated, err := readRequestBody(r, h.maxBodyBytes, h.maxRequestBytes)
 	if err != nil && !shouldLog {
 		log.Printf("[proxy] read request body failed: %v", err)
 		_ = r.Body.Close()
+		if errors.Is(err, errRequestBodyTooLarge) {
+			http.Error(w, "request body too large", http.StatusRequestEntityTooLarge)
+		} else {
 			http.Error(w, "bad request", http.StatusBadRequest)
+		}
 		return
 	}
 	if !shouldLog {
@@ -176,7 +194,7 @@ func (h *Handler) ServeHTTP(w http.ResponseWriter, r *http.Request) {
 		_ = r.Body.Close()
 		s := "read request body: " + err.Error()
 		st.lastErr.Store(&s)
-		h.serveRequestBodyError(cw, r, st, started, reqID)
+		h.serveRequestBodyError(cw, r, st, started, reqID, errors.Is(err, errRequestBodyTooLarge))
 		return
 	}
 
@@ -230,9 +248,15 @@ func (h *Handler) ServeHTTP(w http.ResponseWriter, r *http.Request) {
 	}
 }
 
-func (h *Handler) serveRequestBodyError(cw *captureWriter, r *http.Request, st *requestState, started time.Time, requestID string) {
+func (h *Handler) serveRequestBodyError(cw *captureWriter, r *http.Request, st *requestState, started time.Time, requestID string, tooLarge bool) {
 	cw.Header().Set("X-Request-Id", requestID)
-	http.Error(cw, "bad request", http.StatusBadRequest)
+	status := http.StatusBadRequest
+	message := "bad request"
+	if tooLarge {
+		status = http.StatusRequestEntityTooLarge
+		message = "request body too large"
+	}
+	http.Error(cw, message, status)
 	if !cw.Complete() {
 		return
 	}

proxy/sse.go

@@ -555,8 +555,13 @@ type geminiCandidate struct {
 	raw   map[string]json.RawMessage
 	index int
 	role  string
+	parts []*geminiPart
+}
+
+type geminiPart struct {
+	raw  map[string]json.RawMessage
 	text strings.Builder
-	partRaw map[string]json.RawMessage
+	kind string
 }
 
 func assembleGeminiSSE(payloads []string) ([]byte, bool) {
@@ -588,15 +593,31 @@ func assembleGeminiSSE(payloads []string) ([]byte, bool) {
 				assembled.role = candidate.Content.Role
 			}
 			if len(candidate.Content.Parts) > 0 {
-				for _, part := range candidate.Content.Parts {
-					assembled.text.WriteString(part.Text)
-				}
-				if len(assembled.partRaw) == 0 {
 				var contentRaw struct {
 					Parts []map[string]json.RawMessage `json:"parts"`
 				}
-					if rawContent, ok := candidate.Raw["content"]; ok && json.Unmarshal(rawContent, &contentRaw) == nil && len(contentRaw.Parts) > 0 {
-						assembled.partRaw = cloneRawMap(contentRaw.Parts[0])
+				if rawContent, ok := candidate.Raw["content"]; ok {
+					_ = json.Unmarshal(rawContent, &contentRaw)
+				}
+				mergeByIndex := len(assembled.parts) == len(candidate.Content.Parts)
+				if mergeByIndex {
+					for i := range candidate.Content.Parts {
+						if i >= len(contentRaw.Parts) || assembled.parts[i].kind != geminiPartKind(contentRaw.Parts[i]) {
+							mergeByIndex = false
+							break
+						}
+					}
+				}
+				for i, part := range candidate.Content.Parts {
+					target := i
+					if !mergeByIndex {
+						target = len(assembled.parts)
+						assembled.parts = append(assembled.parts, &geminiPart{})
+					}
+					assembled.parts[target].text.WriteString(part.Text)
+					if i < len(contentRaw.Parts) {
+						assembled.parts[target].kind = geminiPartKind(contentRaw.Parts[i])
+						assembled.parts[target].raw = mergeRawMap(assembled.parts[target].raw, contentRaw.Parts[i])
 					}
 				}
 			}
@@ -614,12 +635,15 @@ func assembleGeminiSSE(payloads []string) ([]byte, bool) {
 		candidate := candidates[index]
 		candidateRaw := cloneRawMap(candidate.raw)
 		candidateRaw["index"] = mustJSON(candidate.index)
-		contentRaw := map[string]any{"role": candidate.role, "parts": []any{map[string]any{"text": candidate.text.String()}}}
-		if len(candidate.partRaw) > 0 {
-			partRaw := cloneRawMap(candidate.partRaw)
-			partRaw["text"] = mustJSON(candidate.text.String())
-			contentRaw["parts"] = []any{rawMapToMap(partRaw)}
+		parts := make([]any, 0, len(candidate.parts))
+		for _, part := range candidate.parts {
+			partRaw := cloneRawMap(part.raw)
+			if part.text.Len() > 0 || len(partRaw) == 0 {
+				partRaw["text"] = mustJSON(part.text.String())
 			}
+			parts = append(parts, rawMapToMap(partRaw))
+		}
+		contentRaw := map[string]any{"role": candidate.role, "parts": parts}
 		candidateRaw["content"] = mustJSON(contentRaw)
 		assembled["candidates"] = append(assembled["candidates"].([]any), rawMapToMap(candidateRaw))
 	}
@@ -631,6 +655,28 @@ func assembleGeminiSSE(payloads []string) ([]byte, bool) {
 	return data, err == nil
 }
 
+func geminiPartKind(part map[string]json.RawMessage) string {
+	for _, key := range []string{"functionCall", "functionResponse", "inlineData", "fileData", "executableCode", "codeExecutionResult", "text"} {
+		if _, ok := part[key]; ok {
+			return key
+		}
+	}
+	return "unknown"
+}
+
+func mergeRawMap(dst, src map[string]json.RawMessage) map[string]json.RawMessage {
+	if dst == nil {
+		dst = make(map[string]json.RawMessage, len(src))
+	}
+	for key, value := range src {
+		if key == "text" {
+			continue
+		}
+		dst[key] = append(json.RawMessage(nil), value...)
+	}
+	return dst
+}
+
 func cloneRawMap(in map[string]json.RawMessage) map[string]json.RawMessage {
 	out := make(map[string]json.RawMessage, len(in))
 	for k, v := range in {

proxy/sse_event.go

@@ -8,12 +8,29 @@ import (
 
 type sseEventTracker struct {
 	buf           []byte
+	limit         int
+	overflow      bool
 	recognized    bool
 	terminal      bool
 	afterTerminal bool
 }
 
+func newSSEEventTracker(limit int64) sseEventTracker {
+	if limit > int64(^uint(0)>>1) {
+		limit = int64(^uint(0) >> 1)
+	}
+	return sseEventTracker{limit: int(limit)}
+}
+
 func (t *sseEventTracker) Write(p []byte) {
+	if t.overflow {
+		return
+	}
+	if t.limit > 0 && len(p) > t.limit-len(t.buf) {
+		t.buf = nil
+		t.overflow = true
+		return
+	}
 	t.buf = append(t.buf, p...)
 	for {
 		end, separator := completeSSEEvent(t.buf)
@@ -34,7 +51,7 @@ func (t *sseEventTracker) Write(p []byte) {
 }
 
 func (t *sseEventTracker) Complete() bool {
-	return len(t.buf) == 0 && !t.afterTerminal && (!t.recognized || t.terminal)
+	return !t.overflow && len(t.buf) == 0 && !t.afterTerminal && (!t.recognized || t.terminal)
 }
 
 func completeSSEEvent(buf []byte) (int, int) {

proxy/sse_event_test.go

@@ -19,3 +19,15 @@ func TestSSEEventTrackerAcceptsCROnlyEventBoundary(t *testing.T) {
 		t.Fatal("terminal SSE event with CR-only boundary should be complete")
 	}
 }
+
+func TestSSEEventTrackerStopsBufferingOverLimit(t *testing.T) {
+	tracker := newSSEEventTracker(8)
+	tracker.Write([]byte("data: 123456789"))
+
+	if tracker.Complete() {
+		t.Fatal("overflowed SSE tracker must not report a complete stream")
+	}
+	if len(tracker.buf) != 0 || !tracker.overflow {
+		t.Fatalf("overflow state=%v buffered=%d, want overflow with released buffer", tracker.overflow, len(tracker.buf))
+	}
+}

proxy/timeout.go

@@ -2,12 +2,30 @@ package proxy
 
 import (
 	"io"
+	"net"
 	"net/http"
 	"sync"
 	"sync/atomic"
 	"time"
 )
 
+type writeTimeoutConn struct {
+	net.Conn
+	timeout time.Duration
+}
+
+func (c *writeTimeoutConn) Write(p []byte) (int, error) {
+	if err := c.Conn.SetWriteDeadline(time.Now().Add(c.timeout)); err != nil {
+		return 0, err
+	}
+	n, err := c.Conn.Write(p)
+	clearErr := c.Conn.SetWriteDeadline(time.Time{})
+	if err == nil {
+		err = clearErr
+	}
+	return n, err
+}
+
 type trackingBody struct {
 	io.ReadCloser
 	failed *atomic.Bool

proxy/writer.go

@@ -26,7 +26,7 @@ type captureWriter struct {
 }
 
 func newCaptureWriter(w http.ResponseWriter, max int64) *captureWriter {
-	return &captureWriter{ResponseWriter: w, max: max, status: http.StatusOK}
+	return &captureWriter{ResponseWriter: w, max: max, status: http.StatusOK, sse: newSSEEventTracker(max)}
 }
 
 func (c *captureWriter) WriteHeader(code int) {
@@ -63,7 +63,7 @@ func (c *captureWriter) Write(p []byte) (int, error) {
 			c.truncated = true
 		}
 		c.written += int64(n)
-		if isStreamResponse(c.Header()) {
+		if c.max > 0 && isStreamResponse(c.Header()) {
 			c.sse.Write(p[:n])
 		}
 		c.flush()
@@ -78,7 +78,7 @@ func (c *captureWriter) Flush() {
 // ConfirmDelivery establishes an observable delivery boundary for responses
 // whose headers were not followed by a body write.
 func (c *captureWriter) ConfirmDelivery() {
-	if c.written != 0 {
+	if c.hijacked || c.written != 0 {
 		return
 	}
 	if _, ok := c.ResponseWriter.(http.Flusher); !ok {
@@ -89,6 +89,9 @@ func (c *captureWriter) ConfirmDelivery() {
 }
 
 func (c *captureWriter) flush() {
+	if c.hijacked {
+		return
+	}
 	if _, ok := c.ResponseWriter.(http.Flusher); !ok {
 		return
 	}

tests/config/config_test.go

@@ -1,7 +1,9 @@
 package config_test
 
 import (
+	"math"
 	"net/netip"
+	"strconv"
 	"strings"
 	"testing"
 	"time"
@@ -50,6 +52,9 @@ func TestLoadQueueDefaults(t *testing.T) {
 	if cfg.MaxBodyBytes != 1<<20 {
 		t.Fatalf("MaxBodyBytes=%d want %d", cfg.MaxBodyBytes, 1<<20)
 	}
+	if cfg.MaxRequestBytes != 16<<20 {
+		t.Fatalf("MaxRequestBytes=%d want %d", cfg.MaxRequestBytes, 16<<20)
+	}
 	if cfg.LogQueueSize != 256 {
 		t.Fatalf("LogQueueSize=%d want 256", cfg.LogQueueSize)
 	}
@@ -242,6 +247,23 @@ func TestLoadRequiresClickHouseURL(t *testing.T) {
 	}
 }
 
+func TestLoadRejectsUnsupportedUpstreamScheme(t *testing.T) {
+	t.Setenv("UPSTREAM_URL", "ftp://example.com")
+	t.Setenv("CLICKHOUSE_URL", "clickhouse://user:pass@localhost:9000/db")
+	if _, err := config.Load(); err == nil {
+		t.Fatal("Load succeeded with unsupported UPSTREAM_URL scheme")
+	}
+}
+
+func TestLoadRejectsMaxRequestBytesOverflowBoundary(t *testing.T) {
+	t.Setenv("UPSTREAM_URL", "https://example.com")
+	t.Setenv("CLICKHOUSE_URL", "clickhouse://user:pass@localhost:9000/db")
+	t.Setenv("MAX_REQUEST_BYTES", strconv.FormatInt(math.MaxInt64, 10))
+	if _, err := config.Load(); err == nil || !strings.Contains(err.Error(), "MAX_REQUEST_BYTES") {
+		t.Fatalf("Load error=%v, want MAX_REQUEST_BYTES rejection", err)
+	}
+}
+
 func TestLoadValidatesClickHouseURL(t *testing.T) {
 	tests := []struct {
 		name string

tests/logger/queue_test.go

@@ -244,8 +244,8 @@ func TestEstimatedBytesCoversStringsAndByteSlices(t *testing.T) {
 		RequestHeaders: []byte("7777777"), RequestBody: []byte("88888888"),
 		ResponseHeaders: []byte("999999999"), ResponseBody: []byte("0000000000"),
 	}
-	if got, want := logger.EstimatedBytes(entry), int64(55); got != want {
-		t.Fatalf("EstimatedBytes=%d want %d", got, want)
+	if got := logger.EstimatedBytes(entry); got <= 55 {
+		t.Fatalf("EstimatedBytes=%d must include fixed entry overhead", got)
 	}
 }
 
@@ -485,6 +485,7 @@ func TestConcurrentStopCannotCancelFirstStopDrain(t *testing.T) {
 	firstCtx, cancelFirst := context.WithTimeout(context.Background(), time.Second)
 	defer cancelFirst()
 	go func() { firstResult <- q.Stop(firstCtx) }()
+	time.Sleep(10 * time.Millisecond)
 	for q.Stats().Dropped == 0 {
 		q.Submit(&logger.LogEntry{RequestID: "stop-probe"})
 	}

tests/proxy/robustness_test.go

@@ -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) {

tests/proxy/sse_capture_test.go

@@ -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")