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

docs: keep verification report reproducible
fix: prevent reconnect after pool close
2026-07-11 17:23:32 +08:00 · 2026-07-10 19:51:51 +08:00 · 2026-07-10 19:48:06 +08:00 · 2026-07-10 19:45:53 +08:00 · 2026-07-10 19:43:41 +08:00 · 2026-07-10 19:40:03 +08:00
22 changed files with 824 additions and 100 deletions
@@ -11,6 +11,8 @@ CLICKHOUSE_URL=

 # 单条请求/响应 body 最大记录字节数
 MAX_BODY_BYTES=1048576
+# 单个代理请求体硬上限，超出时返回 413
+MAX_REQUEST_BYTES=16777216

 # 异步日志队列容量
 LOG_QUEUE_SIZE=256
@@ -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 恢复后重新启用。
@@ -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}"
@@ -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
 }

@@ -2,6 +2,7 @@ package db

 import (
 	"context"
+	"errors"
 	"fmt"
 	"log"
 	"net"
@@ -21,9 +22,13 @@ type Pool struct {
 	open              func(*clickhouse.Options) (clickhouse.Conn, error)
 	mu                sync.RWMutex
 	conn              clickhouse.Conn
+	generation        uint64
 	healthy           bool
+	closed            bool
 }

+var errPoolClosed = errors.New("db pool is closed")
+
 // NewPool 创建 Pool。即使首次连接失败也返回非 nil 实例，后台会持续重试。
 func NewPool(ctx context.Context, dsn string, reconnectInterval time.Duration) *Pool {
 	p := &Pool{dsn: dsn, reconnectInterval: reconnectInterval}
@@ -35,6 +40,13 @@ func NewPool(ctx context.Context, dsn string, reconnectInterval time.Duration) *
 }

 func (p *Pool) connect(ctx context.Context) error {
+	p.mu.RLock()
+	closed := p.closed
+	p.mu.RUnlock()
+	if closed {
+		return errPoolClosed
+	}
+
 	cctx, cancel := context.WithTimeout(ctx, 5*time.Second)
 	defer cancel()
 	opts, err := ClickHouseOptions(p.dsn)
@@ -62,10 +74,16 @@ func (p *Pool) connect(ctx context.Context) error {
 		return err
 	}
 	p.mu.Lock()
+	if p.closed {
+		p.mu.Unlock()
+		_ = conn.Close()
+		return errPoolClosed
+	}
 	if p.conn != nil {
 		_ = p.conn.Close()
 	}
 	p.conn = conn
+	p.generation++
 	p.healthy = true
 	p.mu.Unlock()
 	log.Printf("[db] connected and migrated")
@@ -176,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()
 				}
@@ -207,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()
@@ -214,9 +247,16 @@ 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()
+	p.closed = true
 	conn := p.conn
 	p.conn = nil
 	p.healthy = false
@@ -183,6 +183,57 @@ func TestInitialConnectWithIncompatibleSchemaStaysUnhealthy(t *testing.T) {
 	assertNonDestructiveMigration(t, candidate.statements)
 }

+func TestClosePreventsConcurrentConnectFromPublishing(t *testing.T) {
+	old := newFakeClickHouseConn()
+	pingStarted := make(chan struct{})
+	pingRelease := make(chan struct{})
+	candidate := newFakeClickHouseConn().withBlockedPing(pingStarted, pingRelease)
+	pool := &Pool{dsn: "localhost:9000", conn: old, healthy: true}
+	pool.open = func(*clickhouse.Options) (clickhouse.Conn, error) { return candidate, nil }
+
+	connectDone := make(chan error, 1)
+	go func() { connectDone <- pool.connect(context.Background()) }()
+	<-pingStarted
+
+	if err := pool.Close(context.Background()); err != nil {
+		t.Fatalf("Close: %v", err)
+	}
+	close(pingRelease)
+	if err := <-connectDone; err == nil {
+		t.Fatal("connect succeeded after pool was closed")
+	}
+
+	if pool.Get() != nil || pool.Healthy() {
+		t.Fatalf("closed pool connection = %T, healthy = %v; want nil, false", pool.Get(), pool.Healthy())
+	}
+	if old.closeCalls != 1 {
+		t.Fatalf("old connection close calls = %d, want 1", old.closeCalls)
+	}
+	if candidate.closeCalls != 1 {
+		t.Fatalf("candidate close calls = %d, want 1", candidate.closeCalls)
+	}
+}
+
+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") {
@@ -197,12 +248,14 @@ func assertNonDestructiveMigration(t *testing.T, statements []string) {
 }

 type fakeClickHouseConn struct {
-	rows       *fakeRows
-	row        driver.Row
-	execErr    error
-	queryErr   error
-	statements []string
-	closeCalls int
+	rows        *fakeRows
+	row         driver.Row
+	execErr     error
+	queryErr    error
+	statements  []string
+	closeCalls  int
+	pingStarted chan struct{}
+	pingRelease chan struct{}
 }

 func newFakeClickHouseConn() *fakeClickHouseConn {
@@ -229,6 +282,11 @@ func (c *fakeClickHouseConn) withColumnType(index int, typ string) *fakeClickHou
 	c.rows.values[index][1] = typ
 	return c
 }
+func (c *fakeClickHouseConn) withBlockedPing(started, release chan struct{}) *fakeClickHouseConn {
+	c.pingStarted = started
+	c.pingRelease = release
+	return c
+}

 func (c *fakeClickHouseConn) Contributors() []string                            { return nil }
 func (c *fakeClickHouseConn) ServerVersion() (*driver.ServerVersion, error)     { return nil, nil }
@@ -249,9 +307,15 @@ func (c *fakeClickHouseConn) Exec(_ context.Context, query string, _ ...any) err
 	return c.execErr
 }
 func (c *fakeClickHouseConn) AsyncInsert(context.Context, string, bool, ...any) error { return nil }
-func (c *fakeClickHouseConn) Ping(context.Context) error                              { return nil }
-func (c *fakeClickHouseConn) Stats() driver.Stats                                     { return driver.Stats{} }
-func (c *fakeClickHouseConn) Close() error                                            { c.closeCalls++; return nil }
+func (c *fakeClickHouseConn) Ping(context.Context) error {
+	if c.pingStarted != nil {
+		close(c.pingStarted)
+		<-c.pingRelease
+	}
+	return nil
+}
+func (c *fakeClickHouseConn) Stats() driver.Stats { return driver.Stats{} }
+func (c *fakeClickHouseConn) Close() error        { c.closeCalls++; return nil }

 type fakeRows struct {
 	values    [][]string
@@ -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

-迭代 1 验证全部通过：`gofmt -l .` 无输出，`go test -json ./...` 共 118 个测试通过、0 失败、5 个测试包通过，`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

@@ -46,6 +46,9 @@ branch: main
 - [lesson] 迭代 1：网络写入的 `Send` 错误无法证明提交与否；最小安全策略是不自动重试、显式统计 ambiguous，并将连接标记为不健康。
 - [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

@@ -8,6 +8,7 @@ import (
 	"sync"
 	"sync/atomic"
 	"time"
+	"unsafe"

 	"github.com/ClickHouse/clickhouse-go/v2/lib/driver"

@@ -46,11 +47,16 @@ func (p clickHouseBatchPreparer) PrepareBatch(ctx context.Context, query string)
 }

 type poolBackend struct {
-	pool *db.Pool
+	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,9 +208,10 @@ 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 {
+		shutdownOwner = true
 		q.stopped = true
 		q.shutdownCtx = ctx
 		close(q.ch)
@@ -222,8 +236,14 @@ func (q *Queue) Stop(contexts ...context.Context) error {
 	case <-done:
 		return nil
 	case <-ctx.Done():
-		if workCancel != nil {
+		if shutdownOwner && workCancel != nil {
 			workCancel()
+			timer := time.NewTimer(100 * time.Millisecond)
+			defer timer.Stop()
+			select {
+			case <-done:
+			case <-timer.C:
+			}
 		}
 		return ctx.Err()
 	}
@@ -235,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
@@ -246,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)
@@ -364,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
@@ -388,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++ {
@@ -428,12 +471,20 @@ func Flush(ctx context.Context, conn BatchPreparer, entries []*LogEntry) FlushRe
 	}

 	combined := FlushResult{}
-	for _, entry := range entries {
+	for i, entry := range entries {
 		single := flushOnce(ctx, conn, []*LogEntry{entry}).public([]*LogEntry{entry})
 		combined.Retry = append(combined.Retry, single.Retry...)
 		combined.Failed += single.Failed
 		combined.Ambiguous += single.Ambiguous
 		combined.Err = errors.Join(combined.Err, single.Err)
+		if single.Ambiguous > 0 {
+			if backend, ok := conn.(interface{ MarkUnhealthy() }); ok {
+				backend.MarkUnhealthy()
+			}
+			combined.Failed += len(combined.Retry) + len(entries) - i - 1
+			combined.Retry = nil
+			break
+		}
 	}
 	return combined
 }
@@ -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,7 +83,9 @@ func run() error {
 	}
 	log.Printf("[main] filter mode=%s patterns=%d", filter.Mode, len(filter.Patterns))

-	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)
@@ -97,6 +99,7 @@ func run() error {
 		SSEIdleTimeout:                cfg.UpstreamStreamIdleTimeout,
 		UpstreamTLSInsecureSkipVerify: cfg.UpstreamTLSInsecureSkipVerify,
 		TrustedProxies:                cfg.TrustedProxies,
+		MaxRequestBytes:               cfg.MaxRequestBytes,
 	})

 	srv := &http.Server{
@@ -114,10 +117,7 @@ func run() error {
 		serverErr <- srv.ListenAndServe()
 	}()

-	stop := make(chan os.Signal, 1)
-	signal.Notify(stop, syscall.SIGINT, syscall.SIGTERM)
-	runErr := waitForShutdown(stop, serverErr)
-	signal.Stop(stop)
+	runErr := waitForShutdown(rootCtx.Done(), serverErr)

 	// HTTP、升级连接、日志排空和数据库共享同一个关闭总预算。
 	shutdownCtx, shutdownCancel := context.WithTimeout(context.Background(), 30*time.Second)
@@ -3,20 +3,79 @@ package main
 import (
 	"context"
 	"errors"
+	"go/ast"
+	"go/parser"
+	"go/token"
 	"net/http"
-	"os"
 	"reflect"
 	"strings"
 	"testing"
 	"time"
 )

+func TestRunRegistersSignalsBeforeStartingResources(t *testing.T) {
+	fset := token.NewFileSet()
+	file, err := parser.ParseFile(fset, "main.go", nil, 0)
+	if err != nil {
+		t.Fatalf("parse main.go: %v", err)
+	}
+
+	var runBody *ast.BlockStmt
+	for _, declaration := range file.Decls {
+		function, ok := declaration.(*ast.FuncDecl)
+		if ok && function.Name.Name == "run" {
+			runBody = function.Body
+			break
+		}
+	}
+	if runBody == nil {
+		t.Fatal("main.go does not define run")
+	}
+
+	positions := make(map[string]token.Pos)
+	ast.Inspect(runBody, func(node ast.Node) bool {
+		call, ok := node.(*ast.CallExpr)
+		if !ok {
+			return true
+		}
+		selector, ok := call.Fun.(*ast.SelectorExpr)
+		if !ok {
+			return true
+		}
+		owner, ok := selector.X.(*ast.Ident)
+		if !ok {
+			return true
+		}
+		name := owner.Name + "." + selector.Sel.Name
+		switch name {
+		case "signal.NotifyContext", "db.NewPool", "queue.Start", "srv.ListenAndServe":
+			positions[name] = call.Pos()
+		}
+		return true
+	})
+
+	notifyPos, ok := positions["signal.NotifyContext"]
+	if !ok {
+		t.Fatal("run does not register for shutdown signals")
+	}
+	for _, start := range []string{"db.NewPool", "queue.Start", "srv.ListenAndServe"} {
+		startPos, ok := positions[start]
+		if !ok {
+			t.Fatalf("run does not call %s", start)
+		}
+		if notifyPos >= startPos {
+			t.Errorf("signal.NotifyContext at line %d must precede %s at line %d",
+				fset.Position(notifyPos).Line, start, fset.Position(startPos).Line)
+		}
+	}
+}
+
 func TestWaitForShutdownReturnsListenerError(t *testing.T) {
 	listenErr := errors.New("listen failed")
 	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)
 	}
@@ -24,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)
 		}
@@ -34,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)
 		}
 	})
@@ -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) {
-	if r.Body == nil || r.ContentLength == 0 {
+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
 	}
@@ -3,6 +3,7 @@ package proxy
 import (
 	"context"
 	"crypto/tls"
+	"errors"
 	"log"
 	"net"
 	"net/http"
@@ -25,14 +26,17 @@ type LogSubmitter interface {

 // Handler 构造反代 HTTP handler。
 type Handler struct {
-	rp           *httputil.ReverseProxy
-	filter       *config.Filter
-	queue        LogSubmitter
-	maxBodyBytes int64
-	trusted      []netip.Prefix
-	connMu       sync.Mutex
-	conns        map[net.Conn]struct{}
-	connChanged  chan struct{}
+	rp              *httputil.ReverseProxy
+	filter          *config.Filter
+	queue           LogSubmitter
+	maxBodyBytes    int64
+	maxRequestBytes int64
+	trusted         []netip.Prefix
+	connMu          sync.Mutex
+	conns           map[net.Conn]struct{}
+	connChanged     chan struct{}
+	closing         bool
+	closed          bool
 }

 // Options 控制反代连接上游时的网络行为。
@@ -42,6 +46,7 @@ type Options struct {
 	SSEIdleTimeout                time.Duration
 	UpstreamTLSInsecureSkipVerify bool
 	TrustedProxies                []netip.Prefix
+	MaxRequestBytes               int64
 }

 // requestState 通过 context 在 ErrorHandler / ModifyResponse / 主 handler 之间共享状态。
@@ -75,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
 		}
@@ -128,13 +143,14 @@ func NewWithOptions(upstream *url.URL, filter *config.Filter, queue LogSubmitter
 	}

 	return &Handler{
-		rp:           rp,
-		filter:       filter,
-		queue:        queue,
-		maxBodyBytes: maxBody,
-		trusted:      append([]netip.Prefix(nil), opts.TrustedProxies...),
-		conns:        make(map[net.Conn]struct{}),
-		connChanged:  make(chan struct{}),
+		rp:              rp,
+		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{}),
 	}
 }

@@ -147,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()
-		http.Error(w, "bad request", http.StatusBadRequest)
+		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 {
@@ -174,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
 	}

@@ -221,15 +241,22 @@ func (h *Handler) ServeHTTP(w http.ResponseWriter, r *http.Request) {
 	h.rp.ServeHTTP(cw, r)

 	finished := time.Now()
+	cw.ConfirmDelivery()
 	responseComplete := !isStreamResponse(cw.Header()) || cw.SSEComplete()
 	if cw.Complete() && responseComplete && !st.readFailed.Load() && r.Context().Err() == nil {
 		submitOnce.Do(func() { submit(finished) })
 	}
 }

-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
 	}
@@ -252,7 +279,9 @@ func (h *Handler) serveRequestBodyError(cw *captureWriter, r *http.Request, st *
 func (h *Handler) Shutdown(ctx context.Context) error {
 	for {
 		h.connMu.Lock()
+		h.closing = true
 		if len(h.conns) == 0 {
+			h.closed = true
 			h.connMu.Unlock()
 			return nil
 		}
@@ -283,6 +312,11 @@ func (h *Handler) trackConn(conn net.Conn) net.Conn {
 		h.connMu.Unlock()
 	}
 	h.connMu.Lock()
+	if h.closing || h.closed {
+		h.connMu.Unlock()
+		_ = conn.Close()
+		return conn
+	}
 	h.conns[tracked] = struct{}{}
 	close(h.connChanged)
 	h.connChanged = make(chan struct{})
@@ -552,11 +552,16 @@ func (g *geminiChunk) UnmarshalJSON(data []byte) error {
 }

 type geminiCandidate struct {
-	raw     map[string]json.RawMessage
-	index   int
-	role    string
-	text    strings.Builder
-	partRaw map[string]json.RawMessage
+	raw   map[string]json.RawMessage
+	index int
+	role  string
+	parts []*geminiPart
+}
+
+type geminiPart struct {
+	raw  map[string]json.RawMessage
+	text strings.Builder
+	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)
+				var contentRaw struct {
+					Parts []map[string]json.RawMessage `json:"parts"`
 				}
-				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)
+				}
+				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
+						}
 					}
-					if rawContent, ok := candidate.Raw["content"]; ok && json.Unmarshal(rawContent, &contentRaw) == nil && len(contentRaw.Parts) > 0 {
-						assembled.partRaw = cloneRawMap(contentRaw.Parts[0])
+				}
+				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 {
@@ -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) {
@@ -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))
+	}
+}
@@ -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
@@ -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()
@@ -75,7 +75,23 @@ func (c *captureWriter) Flush() {
 	c.flush()
 }

+// ConfirmDelivery establishes an observable delivery boundary for responses
+// whose headers were not followed by a body write.
+func (c *captureWriter) ConfirmDelivery() {
+	if c.hijacked || c.written != 0 {
+		return
+	}
+	if _, ok := c.ResponseWriter.(http.Flusher); !ok {
+		c.writeFailed = true
+		return
+	}
+	c.flush()
+}
+
 func (c *captureWriter) flush() {
+	if c.hijacked {
+		return
+	}
 	if _, ok := c.ResponseWriter.(http.Flusher); !ok {
 		return
 	}
@@ -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
@@ -38,12 +38,15 @@ func (f *fakeBatch) Send() error  { return f.script.sendErr }
 func (f *fakeBatch) Abort() error { f.abortCalls++; return nil }

 type fakePool struct {
-	calls   int
-	queries []string
-	batches []*fakeBatch
-	scripts []batchScript
+	calls     int
+	queries   []string
+	batches   []*fakeBatch
+	scripts   []batchScript
+	unhealthy bool
 }

+func (p *fakePool) MarkUnhealthy() { p.unhealthy = true }
+
 type queueBackend struct {
 	prepare func(context.Context, string) (logger.Batch, error)
 	healthy atomic.Bool
@@ -215,14 +218,34 @@ func TestFlushAppendRecoveryCountsAmbiguousSingleSend(t *testing.T) {
 	}
 }

+func TestFlushAppendRecoveryStopsAfterFirstAmbiguousSend(t *testing.T) {
+	p := &fakePool{scripts: []batchScript{
+		{appendAt: 1, appendErr: errors.New("bad batch")},
+		{sendErr: errors.New("ack lost")},
+		{},
+	}}
+	entries := []*logger.LogEntry{{RequestID: "a"}, {RequestID: "b"}, {RequestID: "c"}}
+
+	result := logger.Flush(context.Background(), p, entries)
+	if result.Failed != 3 || result.Ambiguous != 1 || len(result.Retry) != 0 {
+		t.Fatalf("unexpected result: %+v", result)
+	}
+	if p.calls != 2 {
+		t.Fatalf("PrepareBatch calls=%d want 2; entries after ambiguous Send must not be attempted", p.calls)
+	}
+	if !p.unhealthy {
+		t.Fatal("ambiguous singleton Send did not immediately mark backend unhealthy")
+	}
+}
+
 func TestEstimatedBytesCoversStringsAndByteSlices(t *testing.T) {
 	entry := &logger.LogEntry{
 		RequestID: "1", Method: "22", Path: "333", Query: "4444", ClientIP: "55555", Error: "666666",
 		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)
 	}
 }

@@ -393,7 +416,7 @@ func TestQueueByteBudgetTracksOwnedEntrySnapshot(t *testing.T) {
 	}
 }

-func TestQueueCanceledStopEventuallyReleasesAllBudget(t *testing.T) {
+func TestQueueCanceledStopReleasesAllBudgetBeforeReturning(t *testing.T) {
 	entry := &logger.LogEntry{RequestID: "queued"}
 	q := logger.NewQueue(nil, 32, 32, 1, time.Hour, 32*logger.EstimatedBytes(entry))
 	q.Start(context.Background())
@@ -405,10 +428,6 @@ func TestQueueCanceledStopEventuallyReleasesAllBudget(t *testing.T) {
 	cancel()
 	_ = q.Stop(ctx)

-	deadline := time.Now().Add(time.Second)
-	for q.Stats().Bytes != 0 && time.Now().Before(deadline) {
-		time.Sleep(time.Millisecond)
-	}
 	if stats := q.Stats(); stats.Bytes != 0 || stats.Failed != 32 {
 		t.Fatalf("stats=%+v", stats)
 	}
@@ -443,6 +462,49 @@ func TestQueueStopAndSubmitAreConcurrentAndRepeatSafe(t *testing.T) {
 	}
 }

+func TestConcurrentStopCannotCancelFirstStopDrain(t *testing.T) {
+	const initialEntries = 32
+	firstPrepareStarted := make(chan struct{})
+	releaseFirstPrepare := make(chan struct{})
+	var prepareCalls atomic.Int32
+	backend := newQueueBackend(func(context.Context, string) (logger.Batch, error) {
+		if prepareCalls.Add(1) == 1 {
+			close(firstPrepareStarted)
+			<-releaseFirstPrepare
+		}
+		return &fakeBatch{}, nil
+	})
+	q := logger.NewQueueWithBackend(backend, 1024, 1, 1, time.Hour)
+	q.Start(context.Background())
+	for i := 0; i < initialEntries; i++ {
+		q.Submit(&logger.LogEntry{RequestID: "queued"})
+	}
+	<-firstPrepareStarted
+
+	firstResult := make(chan error, 1)
+	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"})
+	}
+	enqueued := q.Stats().Enqueued
+
+	secondCtx, cancelSecond := context.WithCancel(context.Background())
+	cancelSecond()
+	if err := q.Stop(secondCtx); !errors.Is(err, context.Canceled) {
+		t.Fatalf("second Stop error=%v want canceled", err)
+	}
+	close(releaseFirstPrepare)
+	if err := <-firstResult; err != nil {
+		t.Fatalf("first Stop: %v", err)
+	}
+	if calls := uint64(backend.calls.Load()); calls != enqueued {
+		t.Fatalf("PrepareBatch calls=%d want %d; second Stop interrupted drain", calls, enqueued)
+	}
+}
+
 func TestStopIsIndependentFromStartContext(t *testing.T) {
 	root, cancelRoot := context.WithCancel(context.Background())
 	q := logger.NewQueue(nil, 4, 4, 1, time.Hour, 1024)
@@ -498,10 +560,6 @@ func TestQueueStopDeadlineCancelsBlockedSend(t *testing.T) {
 	if err := q.Stop(ctx); !errors.Is(err, context.DeadlineExceeded) {
 		t.Fatalf("Stop error=%v want deadline exceeded", err)
 	}
-	deadline := time.Now().Add(250 * time.Millisecond)
-	for q.Stats().Bytes != 0 && time.Now().Before(deadline) {
-		time.Sleep(time.Millisecond)
-	}
 	if stats := q.Stats(); stats.Bytes != 0 || stats.Ambiguous != 1 || stats.Failed != 1 {
 		t.Fatalf("stats after canceled Send=%+v", stats)
 	}
@@ -67,6 +67,14 @@ func (*failingFlushResponseWriter) FlushError() error {
 	return errors.New("flush failed")
 }

+type unflushableResponseWriter struct{ header http.Header }
+
+func (w *unflushableResponseWriter) Header() http.Header { return w.header }
+func (*unflushableResponseWriter) WriteHeader(int)       {}
+func (*unflushableResponseWriter) Write(p []byte) (int, error) {
+	return len(p), nil
+}
+
 func newTestHandler(t *testing.T, upstream *url.URL, sub *captureSubmitter, opts proxy.Options) *proxy.Handler {
 	t.Helper()
 	filter, err := config.NewFilter(config.FilterDisabled, nil)
@@ -100,6 +108,62 @@ 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) {
+		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 = 0
+	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) {
@@ -196,6 +260,70 @@ func TestResponseFlushFailurePreventsCommit(t *testing.T) {
 	}
 }

+func TestHeaderOnlyResponseFlushFailurePreventsCommit(t *testing.T) {
+	tests := []struct {
+		name   string
+		method string
+		status int
+	}{
+		{name: "head", method: http.MethodHead, status: http.StatusOK},
+		{name: "no content", method: http.MethodGet, status: http.StatusNoContent},
+		{name: "not modified", method: http.MethodGet, status: http.StatusNotModified},
+	}
+	for _, tc := range tests {
+		t.Run(tc.name, func(t *testing.T) {
+			upstream := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+				w.WriteHeader(tc.status)
+			}))
+			defer upstream.Close()
+			u, _ := url.Parse(upstream.URL)
+
+			sub := &captureSubmitter{}
+			h := newTestHandler(t, u, sub, proxy.Options{})
+			w := &failingFlushResponseWriter{header: make(http.Header)}
+			h.ServeHTTP(w, httptest.NewRequest(tc.method, "http://proxy.test/v1/test", nil))
+			if sub.Len() != 0 {
+				t.Fatal("response with undelivered headers was committed")
+			}
+		})
+	}
+}
+
+func TestHeaderOnlyResponseSuccessfulFlushCommits(t *testing.T) {
+	for _, status := range []int{http.StatusNoContent, http.StatusNotModified} {
+		t.Run(http.StatusText(status), func(t *testing.T) {
+			upstream := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+				w.WriteHeader(status)
+			}))
+			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() != 1 {
+				t.Fatalf("entries=%d, want successfully delivered response committed", sub.Len())
+			}
+		})
+	}
+}
+
+func TestHeaderOnlyResponseWithoutDeliveryBoundaryDoesNotCommit(t *testing.T) {
+	upstream := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		w.WriteHeader(http.StatusNoContent)
+	}))
+	defer upstream.Close()
+	u, _ := url.Parse(upstream.URL)
+
+	sub := &captureSubmitter{}
+	h := newTestHandler(t, u, sub, proxy.Options{})
+	w := &unflushableResponseWriter{header: make(http.Header)}
+	h.ServeHTTP(w, httptest.NewRequest(http.MethodGet, "http://proxy.test/v1/test", nil))
+	if sub.Len() != 0 {
+		t.Fatal("header-only response without a delivery boundary was committed")
+	}
+}
+
 func TestTrustedProxyClientIP(t *testing.T) {
 	upstream := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
 		_, _ = io.WriteString(w, "ok")
@@ -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")

@@ -11,6 +11,7 @@ import (
 	"net/http/httptest"
 	"net/url"
 	"strings"
+	"sync"
 	"testing"
 	"time"

@@ -45,6 +46,30 @@ type failingWriteConn struct{ net.Conn }

 func (failingWriteConn) Write([]byte) (int, error) { return 0, errors.New("handshake write failed") }

+type closeNotifyConn struct {
+	net.Conn
+	closed chan struct{}
+	once   sync.Once
+}
+
+func (c *closeNotifyConn) Close() error {
+	err := c.Conn.Close()
+	c.once.Do(func() { close(c.closed) })
+	return err
+}
+
+type blockingHijackResponseWriter struct {
+	*hijackResponseWriter
+	hijacked chan struct{}
+	release  chan struct{}
+}
+
+func (w *blockingHijackResponseWriter) Hijack() (net.Conn, *bufio.ReadWriter, error) {
+	close(w.hijacked)
+	<-w.release
+	return w.hijackResponseWriter.Hijack()
+}
+
 // fakeUpstream 模拟一个最简 WebSocket 升级：
 // 收到 GET + Upgrade: websocket 后回 101，然后做字节回声直到对端关闭。
 func fakeUpstream(t *testing.T) *httptest.Server {
@@ -161,6 +186,61 @@ func TestWebSocketHandshakeCapturedAndShutdownClosesConnection(t *testing.T) {
 	}
 }

+func TestShutdownClosesConnectionHijackedBeforeRegistration(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)
+	downstream, peer := net.Pipe()
+	defer peer.Close()
+	closed := make(chan struct{})
+	conn := &closeNotifyConn{Conn: downstream, closed: closed}
+	w := &blockingHijackResponseWriter{
+		hijackResponseWriter: &hijackResponseWriter{header: make(http.Header), conn: conn},
+		hijacked:             make(chan struct{}),
+		release:              make(chan struct{}),
+	}
+	req := httptest.NewRequest(http.MethodGet, "http://proxy.test/v1/realtime", nil)
+	req.Header.Set("Upgrade", "websocket")
+	req.Header.Set("Connection", "Upgrade")
+	serveDone := make(chan struct{})
+	var releaseOnce sync.Once
+	releaseHijack := func() { releaseOnce.Do(func() { close(w.release) }) }
+	go func() {
+		h.ServeHTTP(w, req)
+		close(serveDone)
+	}()
+	defer func() {
+		releaseHijack()
+		ctx, cancel := context.WithTimeout(context.Background(), time.Second)
+		defer cancel()
+		_ = h.Shutdown(ctx)
+		<-serveDone
+	}()
+
+	select {
+	case <-w.hijacked:
+	case <-time.After(time.Second):
+		t.Fatal("downstream connection was not hijacked")
+	}
+	ctx, cancel := context.WithTimeout(context.Background(), time.Second)
+	defer cancel()
+	if err := h.Shutdown(ctx); err != nil {
+		t.Fatal(err)
+	}
+	releaseHijack()
+
+	select {
+	case <-closed:
+	case <-time.After(100 * time.Millisecond):
+		t.Fatal("connection hijacked during Shutdown remained open")
+	}
+}
+
 func TestWebSocketHandshakeWriteFailureIsNotCaptured(t *testing.T) {
 	upstream := fakeUpstream(t)
 	defer upstream.Close()
Author	SHA1	Message	Date
m1saka	9e6e7ca3c7	修复代理与日志链路的可靠性问题	2026-07-11 17:23:32 +08:00
MiMoCode	d39f5a1048	docs: keep verification report reproducible	2026-07-10 19:51:51 +08:00
MiMoCode	b7442ab514	fix: prevent reconnect after pool close	2026-07-10 19:48:06 +08:00
MiMoCode	fcb4a87be8	fix: register shutdown signals before startup	2026-07-10 19:45:53 +08:00
MiMoCode	f8abacf2e5	fix: close hijacks racing with shutdown	2026-07-10 19:43:41 +08:00
MiMoCode	62675e400b	fix: confirm header-only response delivery	2026-07-10 19:40:03 +08:00
MiMoCode	32988b6367	fix: pre-read bodies with unknown length	2026-07-10 19:36:49 +08:00
MiMoCode	f27382c33a	fix: stop row isolation after ambiguous send	2026-07-10 19:34:56 +08:00
MiMoCode	3717632f27	fix: finish queue cleanup before stop returns	2026-07-10 19:28:03 +08:00
MiMoCode	0520f75cdf	fix: preserve first stop shutdown ownership	2026-07-10 19:25:44 +08:00