修复日志持久化可靠性与数据库安全默认值

.env.example

@@ -11,6 +11,8 @@ CLICKHOUSE_URL=
 
 # 单条请求/响应 body 最大记录字节数
 MAX_BODY_BYTES=1048576
+# 单个代理请求体硬上限，超出时返回 413
+MAX_REQUEST_BYTES=16777216
 
 # 异步日志队列容量
 LOG_QUEUE_SIZE=256
@@ -45,13 +47,18 @@ UPSTREAM_STREAM_IDLE_TIMEOUT=2m
 TRUSTED_PROXIES=
 
 # ===== Docker Compose 可选配置 =====
-# 对外暴露的代理端口（容器内固定 :8080）
+# tokenthief 的宿主机监听 IP 和端口（容器内固定 :8080）
+TOKENTHIEF_LISTEN_IP=0.0.0.0
 LISTEN_PORT=8080
 
 # Compose 内置 ClickHouse 配置
 CLICKHOUSE_USER=tokenthief
 CLICKHOUSE_PASSWORD=
 CLICKHOUSE_DB=tokenthief
+# ClickHouse 的宿主机监听 IP，以及 HTTP 和原生客户端端口
+CLICKHOUSE_LISTEN_IP=127.0.0.1
+CLICKHOUSE_HTTP_PORT=8123
+CLICKHOUSE_NATIVE_PORT=9000
 
 # 容器时区
 TZ=Asia/Shanghai

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

@@ -10,13 +10,14 @@ services:
       thief_clickhouse:
         condition: service_healthy
     ports:
-      - "${LISTEN_PORT:-8080}:8080"
+      - "${TOKENTHIEF_LISTEN_IP:-0.0.0.0}:${LISTEN_PORT:-8080}:8080"
     environment:
       LISTEN_ADDR: ":8080"
       UPSTREAM_URL: "${UPSTREAM_URL:?set UPSTREAM_URL}"
       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}"
@@ -38,9 +39,12 @@ services:
       - tokenthief
 
   thief_clickhouse:
-    image: clickhouse/clickhouse-server:25.3.3.42-alpine
+    image: clickhouse/clickhouse-server:26.3
     container_name: tokenthief-clickhouse
     restart: unless-stopped
+    ports:
+      - "${CLICKHOUSE_LISTEN_IP:-127.0.0.1}:${CLICKHOUSE_HTTP_PORT:-8123}:8123"
+      - "${CLICKHOUSE_LISTEN_IP:-127.0.0.1}:${CLICKHOUSE_NATIVE_PORT:-9000}:9000"
     environment:
       CLICKHOUSE_USER: "${CLICKHOUSE_USER:-tokenthief}"
       CLICKHOUSE_PASSWORD: "${CLICKHOUSE_PASSWORD:?set a strong CLICKHOUSE_PASSWORD}"

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

@@ -2,6 +2,7 @@ package db
 
 import (
 	"context"
+	"errors"
 	"fmt"
 	"log"
 	"net"
@@ -21,9 +22,23 @@ type Pool struct {
 	open              func(*clickhouse.Options) (clickhouse.Conn, error)
 	mu                sync.RWMutex
 	conn              clickhouse.Conn
+	generation        uint64
 	healthy           bool
+	closed            bool
+	leases            map[uint64]*connectionLease
+	leaseChanged      chan struct{}
+	closing           int
 }
 
+type connectionLease struct {
+	conn    clickhouse.Conn
+	active  int
+	retired 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 +50,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,12 +84,21 @@ func (p *Pool) connect(ctx context.Context) error {
 		return err
 	}
 	p.mu.Lock()
-	if p.conn != nil {
+	if p.closed {
-		_ = p.conn.Close()
+		p.mu.Unlock()
+		_ = conn.Close()
+		return errPoolClosed
 	}
+	old := p.retireCurrentLocked()
 	p.conn = conn
+	p.generation++
+	p.ensureLeaseLocked(p.generation, conn)
 	p.healthy = true
 	p.mu.Unlock()
+	if old != nil {
+		_ = old.Close()
+		p.finishClosing()
+	}
 	log.Printf("[db] connected and migrated")
 	return nil
 }
@@ -176,13 +207,14 @@ func (p *Pool) watch(ctx context.Context) {
 			return
 		case <-t.C:
 			if p.Healthy() {
-				if conn := p.Get(); conn != nil {
+				if conn, generation, release := p.Acquire(); 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()
+					release()
 				}
 				continue
 			}
@@ -207,6 +239,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,23 +261,132 @@ 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
+}
+
+// Acquire pins the current connection until release is called.
+func (p *Pool) Acquire() (clickhouse.Conn, uint64, func()) {
+	p.mu.Lock()
+	if p.closed || p.conn == nil {
+		p.mu.Unlock()
+		return nil, p.generation, nil
+	}
+	generation := p.generation
+	lease := p.ensureLeaseLocked(generation, p.conn)
+	lease.active++
+	p.mu.Unlock()
+
+	var once sync.Once
+	return lease.conn, generation, func() {
+		once.Do(func() { p.release(generation) })
+	}
+}
+
+func (p *Pool) ensureLeaseLocked(generation uint64, conn clickhouse.Conn) *connectionLease {
+	if p.leases == nil {
+		p.leases = make(map[uint64]*connectionLease)
+	}
+	lease := p.leases[generation]
+	if lease == nil {
+		lease = &connectionLease{conn: conn}
+		p.leases[generation] = lease
+	}
+	return lease
+}
+
+func (p *Pool) notifyLeaseChangedLocked() {
+	if p.leaseChanged != nil {
+		close(p.leaseChanged)
+	}
+	p.leaseChanged = make(chan struct{})
+}
+
+func (p *Pool) retireCurrentLocked() clickhouse.Conn {
+	if p.conn == nil {
+		return nil
+	}
+	lease := p.ensureLeaseLocked(p.generation, p.conn)
+	lease.retired = true
+	if lease.active != 0 || lease.closed {
+		return nil
+	}
+	lease.closed = true
+	delete(p.leases, p.generation)
+	p.closing++
+	return lease.conn
+}
+
+func (p *Pool) finishClosing() {
+	p.mu.Lock()
+	p.closing--
+	p.notifyLeaseChangedLocked()
+	p.mu.Unlock()
+}
+
+func (p *Pool) release(generation uint64) {
+	p.mu.Lock()
+	lease := p.leases[generation]
+	if lease == nil {
+		p.mu.Unlock()
+		return
+	}
+	lease.active--
+	var conn clickhouse.Conn
+	if lease.active == 0 && lease.retired && !lease.closed {
+		lease.closed = true
+		conn = lease.conn
+		p.closing++
+	}
+	p.mu.Unlock()
+	if conn != nil {
+		_ = conn.Close()
+		p.mu.Lock()
+		delete(p.leases, generation)
+		p.closing--
+		p.notifyLeaseChangedLocked()
+		p.mu.Unlock()
+	}
+}
+
 // Close starts closing the underlying connection and waits within ctx.
 func (p *Pool) Close(ctx context.Context) error {
 	p.mu.Lock()
-	conn := p.conn
+	p.closed = true
+	conn := p.retireCurrentLocked()
 	p.conn = nil
 	p.healthy = false
 	p.mu.Unlock()
-	if conn == nil {
+	var closeErr error
-		return nil
+	if conn != nil {
+		done := make(chan error, 1)
+		go func() {
+			done <- conn.Close()
+			p.finishClosing()
+		}()
+		select {
+		case closeErr = <-done:
+		case <-ctx.Done():
+			return ctx.Err()
+		}
 	}
-
+	for {
-	done := make(chan error, 1)
+		p.mu.Lock()
-	go func() { done <- conn.Close() }()
+		if len(p.leases) == 0 && p.closing == 0 {
-	select {
+			p.mu.Unlock()
-	case err := <-done:
+			return closeErr
-		return err
+		}
-	case <-ctx.Done():
+		if p.leaseChanged == nil {
-		return ctx.Err()
+			p.leaseChanged = make(chan struct{})
+		}
+		changed := p.leaseChanged
+		p.mu.Unlock()
+		select {
+		case <-changed:
+		case <-ctx.Done():
+			return ctx.Err()
+		}
 	}
 }

db/clickhouse_test.go

@@ -6,6 +6,7 @@ import (
 	"reflect"
 	"strings"
 	"testing"
+	"time"
 
 	"github.com/ClickHouse/clickhouse-go/v2"
 	"github.com/ClickHouse/clickhouse-go/v2/lib/driver"
@@ -133,6 +134,29 @@ func TestValidateProxyLogsSchema(t *testing.T) {
 	}
 }
 
+func TestMissingProxyLogColumnsReturnsAdditiveDDL(t *testing.T) {
+	columns := append([]schemaColumn(nil), proxyLogsColumns...)
+	columns = append(columns[:7], columns[8:]...)
+
+	statements, err := missingProxyLogColumns(columns)
+	if err != nil {
+		t.Fatalf("missingProxyLogColumns: %v", err)
+	}
+	want := []string{"ALTER TABLE proxy_logs ADD COLUMN IF NOT EXISTS request_truncated Bool DEFAULT false"}
+	if !reflect.DeepEqual(statements, want) {
+		t.Fatalf("statements=%q want %q", statements, want)
+	}
+}
+
+func TestValidateProxyLogsSchemaAllowsExtraColumns(t *testing.T) {
+	columns := append([]schemaColumn(nil), proxyLogsColumns...)
+	columns = append(columns, schemaColumn{"deployment", "String"})
+	table := schemaTable{"MergeTree", "toYYYYMM(started_at)", "started_at, request_id"}
+	if err := validateProxyLogsSchema(columns, table); err != nil {
+		t.Fatalf("validateProxyLogsSchema: %v", err)
+	}
+}
+
 func TestConnectRejectsUnhealthySchemaWithoutReplacingHealthyConnection(t *testing.T) {
 	tests := []struct {
 		name      string
@@ -183,6 +207,154 @@ func TestInitialConnectWithIncompatibleSchemaStaysUnhealthy(t *testing.T) {
 	assertNonDestructiveMigration(t, candidate.statements)
 }
 
+func TestMigrationDoesNotAlterIncompatibleTable(t *testing.T) {
+	values := schemaValues()
+	values = append(values[:7], values[8:]...)
+	candidate := newFakeClickHouseConn().withRows(&fakeRows{values: values, scanErrAt: -1})
+	candidate.row = fakeRow{values: []string{"ReplacingMergeTree", "toYYYYMM(started_at)", "started_at, request_id"}}
+	pool := &Pool{dsn: "localhost:9000"}
+	pool.open = func(*clickhouse.Options) (clickhouse.Conn, error) { return candidate, nil }
+
+	if err := pool.connect(context.Background()); err == nil {
+		t.Fatal("connect succeeded with incompatible table")
+	}
+	for _, statement := range candidate.statements {
+		if strings.Contains(strings.ToUpper(statement), "ALTER TABLE") {
+			t.Fatalf("migration altered incompatible table: %q", statement)
+		}
+	}
+}
+
+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 TestConnectionLeaseDefersReplacementCloseUntilRelease(t *testing.T) {
+	old := newFakeClickHouseConn()
+	candidate := newFakeClickHouseConn()
+	pool := &Pool{dsn: "localhost:9000", conn: old, healthy: true, generation: 1}
+	pool.open = func(*clickhouse.Options) (clickhouse.Conn, error) { return candidate, nil }
+
+	conn, generation, release := pool.Acquire()
+	if conn != old || generation != 1 || release == nil {
+		t.Fatalf("Acquire=(%T,%d,%v), want old,1,release", conn, generation, release != nil)
+	}
+	if err := pool.connect(context.Background()); err != nil {
+		t.Fatalf("connect: %v", err)
+	}
+	if old.closeCalls != 0 {
+		t.Fatalf("leased old connection closed %d times", old.closeCalls)
+	}
+	release()
+	if old.closeCalls != 1 {
+		t.Fatalf("old connection close calls=%d want 1 after release", old.closeCalls)
+	}
+	release()
+	if old.closeCalls != 1 {
+		t.Fatalf("release was not idempotent: close calls=%d", old.closeCalls)
+	}
+}
+
+func TestCloseWaitsForConnectionLease(t *testing.T) {
+	conn := newFakeClickHouseConn()
+	pool := &Pool{conn: conn, healthy: true, generation: 1}
+	_, _, release := pool.Acquire()
+
+	ctx, cancel := context.WithCancel(context.Background())
+	cancel()
+	if err := pool.Close(ctx); !errors.Is(err, context.Canceled) {
+		t.Fatalf("Close error=%v want canceled", err)
+	}
+	if conn.closeCalls != 0 {
+		t.Fatalf("leased connection closed %d times", conn.closeCalls)
+	}
+	if acquired, _, _ := pool.Acquire(); acquired != nil {
+		t.Fatal("Acquire succeeded after Close")
+	}
+	release()
+	if conn.closeCalls != 1 {
+		t.Fatalf("connection close calls=%d want 1 after release", conn.closeCalls)
+	}
+}
+
+func TestCloseWaitsForRetiredConnectionLease(t *testing.T) {
+	old := newFakeClickHouseConn()
+	current := newFakeClickHouseConn()
+	pool := &Pool{conn: old, healthy: true, generation: 1}
+	_, _, release := pool.Acquire()
+	pool.mu.Lock()
+	pool.retireCurrentLocked()
+	pool.conn = current
+	pool.generation++
+	pool.ensureLeaseLocked(pool.generation, current)
+	pool.mu.Unlock()
+
+	closeDone := make(chan error, 1)
+	go func() { closeDone <- pool.Close(context.Background()) }()
+	select {
+	case err := <-closeDone:
+		t.Fatalf("Close returned before retired lease release: %v", err)
+	case <-time.After(25 * time.Millisecond):
+	}
+	if old.closeCalls != 0 {
+		t.Fatalf("retired leased connection closed %d times", old.closeCalls)
+	}
+	release()
+	if err := <-closeDone; err != nil {
+		t.Fatalf("Close: %v", err)
+	}
+	if old.closeCalls != 1 {
+		t.Fatalf("retired connection close calls=%d want 1", old.closeCalls)
+	}
+}
+
 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 +369,14 @@ func assertNonDestructiveMigration(t *testing.T, statements []string) {
 }
 
 type fakeClickHouseConn struct {
-	rows       *fakeRows
+	rows        *fakeRows
-	row        driver.Row
+	row         driver.Row
-	execErr    error
+	execErr     error
-	queryErr   error
+	queryErr    error
-	statements []string
+	statements  []string
-	closeCalls int
+	closeCalls  int
+	pingStarted chan struct{}
+	pingRelease chan struct{}
 }
 
 func newFakeClickHouseConn() *fakeClickHouseConn {
@@ -229,6 +403,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 +428,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) Ping(context.Context) error {
-func (c *fakeClickHouseConn) Stats() driver.Stats                                     { return driver.Stats{} }
+	if c.pingStarted != nil {
-func (c *fakeClickHouseConn) Close() error                                            { c.closeCalls++; return 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

db/migrate.go

@@ -38,43 +38,81 @@ func migrate(ctx context.Context, conn clickhouse.Conn) error {
 		return fmt.Errorf("create proxy_logs: %w", err)
 	}
 
-	rows, err := conn.Query(ctx, `
+	columns, err := queryProxyLogColumns(ctx, conn)
-SELECT name, type
-FROM system.columns
-WHERE database = currentDatabase() AND table = 'proxy_logs'
-ORDER BY position`)
 	if err != nil {
-		return fmt.Errorf("query proxy_logs columns: %w", err)
+		return err
 	}
-	var columns []schemaColumn
+	table, err := queryProxyLogsTable(ctx, conn)
-	for rows.Next() {
+	if err != nil {
-		var column schemaColumn
+		return err
-		if err := rows.Scan(&column.name, &column.typ); err != nil {
+	}
-			rows.Close()
+	if err := validateProxyLogsTable(table); err != nil {
-			return fmt.Errorf("scan proxy_logs columns: %w", err)
+		return fmt.Errorf("incompatible proxy_logs schema: %w", err)
+	}
+	statements, err := missingProxyLogColumns(columns)
+	if err != nil {
+		return fmt.Errorf("plan proxy_logs migration: %w", err)
+	}
+	for _, statement := range statements {
+		if err := conn.Exec(ctx, statement); err != nil {
+			return fmt.Errorf("alter proxy_logs: %w", err)
 		}
-		columns = append(columns, column)
 	}
-	if err := rows.Err(); err != nil {
+	if len(statements) > 0 {
-		rows.Close()
+		columns, err = queryProxyLogColumns(ctx, conn)
-		return fmt.Errorf("read proxy_logs columns: %w", err)
+		if err != nil {
+			return err
+		}
 	}
-	rows.Close()
 
+	table, err = queryProxyLogsTable(ctx, conn)
+	if err != nil {
+		return err
+	}
+	if err := validateProxyLogsSchema(columns, table); err != nil {
+		return fmt.Errorf("incompatible proxy_logs schema: %w", err)
+	}
+	return nil
+}
+
+func queryProxyLogsTable(ctx context.Context, conn clickhouse.Conn) (schemaTable, error) {
 	var table schemaTable
-	err = conn.QueryRow(ctx, `
+	err := conn.QueryRow(ctx, `
 SELECT engine, partition_key, sorting_key
 FROM system.tables
 WHERE database = currentDatabase() AND name = 'proxy_logs'`).Scan(
 		&table.engine, &table.partitionKey, &table.sortingKey,
 	)
 	if err != nil {
-		return fmt.Errorf("query proxy_logs table: %w", err)
+		return schemaTable{}, fmt.Errorf("query proxy_logs table: %w", err)
 	}
-	if err := validateProxyLogsSchema(columns, table); err != nil {
+	return table, nil
-		return fmt.Errorf("incompatible proxy_logs schema: %w", err)
+}
+
+func queryProxyLogColumns(ctx context.Context, conn clickhouse.Conn) ([]schemaColumn, error) {
+	rows, err := conn.Query(ctx, `
+SELECT name, type
+FROM system.columns
+WHERE database = currentDatabase() AND table = 'proxy_logs'
+ORDER BY position`)
+	if err != nil {
+		return nil, fmt.Errorf("query proxy_logs columns: %w", err)
 	}
-	return nil
+	var columns []schemaColumn
+	for rows.Next() {
+		var column schemaColumn
+		if err := rows.Scan(&column.name, &column.typ); err != nil {
+			rows.Close()
+			return nil, fmt.Errorf("scan proxy_logs columns: %w", err)
+		}
+		columns = append(columns, column)
+	}
+	if err := rows.Err(); err != nil {
+		rows.Close()
+		return nil, fmt.Errorf("read proxy_logs columns: %w", err)
+	}
+	rows.Close()
+	return columns, nil
 }
 
 type schemaColumn struct {
@@ -108,15 +146,58 @@ var proxyLogsColumns = []schemaColumn{
 	{"error", "String"},
 }
 
-func validateProxyLogsSchema(columns []schemaColumn, table schemaTable) error {
+var proxyLogColumnDDL = map[string]string{
-	if len(columns) != len(proxyLogsColumns) {
+	"request_id":         "request_id String",
-		return fmt.Errorf("got %d columns, want %d", len(columns), len(proxyLogsColumns))
+	"method":             "method String",
+	"path":               "path String",
+	"query":              "query String",
+	"client_ip":          "client_ip String",
+	"request_headers":    "request_headers String",
+	"request_body":       "request_body String",
+	"request_truncated":  "request_truncated Bool DEFAULT false",
+	"status_code":        "status_code Int32",
+	"response_headers":   "response_headers String",
+	"response_body":      "response_body String",
+	"response_truncated": "response_truncated Bool DEFAULT false",
+	"is_stream":          "is_stream Bool DEFAULT false",
+	"latency_ms":         "latency_ms Int64",
+	"started_at":         "started_at DateTime64(3)",
+	"finished_at":        "finished_at DateTime64(3)",
+	"error":              "error String",
+}
+
+func missingProxyLogColumns(columns []schemaColumn) ([]string, error) {
+	existing := make(map[string]string, len(columns))
+	for _, column := range columns {
+		existing[column.name] = column.typ
 	}
-	for i, want := range proxyLogsColumns {
+	var statements []string
-		if columns[i] != want {
+	for _, required := range proxyLogsColumns {
-			return fmt.Errorf("column %d is %s %s, want %s %s", i+1, columns[i].name, columns[i].typ, want.name, want.typ)
+		if typ, ok := existing[required.name]; ok {
+			if typ != required.typ {
+				return nil, fmt.Errorf("column %s has type %s, want %s", required.name, typ, required.typ)
+			}
+			continue
+		}
+		statements = append(statements, "ALTER TABLE proxy_logs ADD COLUMN IF NOT EXISTS "+proxyLogColumnDDL[required.name])
+	}
+	return statements, nil
+}
+
+func validateProxyLogsSchema(columns []schemaColumn, table schemaTable) error {
+	existing := make(map[string]string, len(columns))
+	for _, column := range columns {
+		existing[column.name] = column.typ
+	}
+	for _, want := range proxyLogsColumns {
+		if typ, ok := existing[want.name]; !ok || typ != want.typ {
+			return fmt.Errorf("column %s is %s, want %s", want.name, typ, want.typ)
 		}
 	}
+	return validateProxyLogsTable(table)
+}
+
+func validateProxyLogsTable(table schemaTable) error {
 	if table.engine != "MergeTree" {
 		return fmt.Errorf("engine is %q, want MergeTree", table.engine)
 	}

docs/compose/plans/2026-07-12-db-log-reliability.md

@@ -0,0 +1,83 @@
+# Database And Log Reliability Implementation Plan
+
+> [!NOTE]
+> This document may not reflect the current implementation.
+> See the final report for up-to-date state:
+> [Final Report](../reports/db-log-reliability.md)
+
+> **For agentic workers:** REQUIRED SUB-SKILL: Use compose:subagent (recommended) or compose:execute to implement this plan task-by-task. Steps use checkbox (`- [ ]`) syntax for tracking.
+
+**Goal:** Preserve queued logs during temporary database outages, prevent active ClickHouse connections from being closed during replacement, and migrate missing log columns safely.
+
+**Architecture:** Add reference-counted connection leases to `db.Pool`, consume those leases from the logger adapter, and retain an in-memory worker batch while the backend is unhealthy. Extend migration with fixed, additive DDL for missing known columns while retaining strict validation of existing columns and table keys.
+
+**Tech Stack:** Go, ClickHouse Go driver, standard library concurrency primitives, Go tests.
+
+## Global Constraints
+
+- Change only findings 2, 3, and 4 from the review.
+- Keep `Submit` non-blocking and retain existing queue entry and byte budgets.
+- Do not add disk persistence or retry ambiguous `Send` failures.
+- Do not destructively modify existing ClickHouse columns, engine, partition key, or sorting key.
+
+---
+
+### Task 1: Additive Schema Migration
+
+**Files:**
+- Modify: `db/migrate.go`
+- Modify: `db/clickhouse_test.go`
+
+**Interfaces:**
+- Produces: migration that adds missing entries from `proxyLogsColumns` using fixed `ALTER TABLE proxy_logs ADD COLUMN IF NOT EXISTS` statements.
+
+- [ ] Add tests proving a missing known column executes its fixed additive DDL and succeeds after refreshed metadata; extra columns are accepted; wrong existing types and table keys remain rejected.
+- [ ] Run `go test ./db -run 'Test.*Migration|TestValidateProxyLogsSchema' -count=1` and confirm the new tests fail.
+- [ ] Implement a fixed column-definition map, detect missing columns by name, execute additive DDL, re-query metadata, and validate required columns by name and type without rejecting extras.
+- [ ] Run `go test ./db -count=1` and confirm it passes.
+
+### Task 2: Connection Leases
+
+**Files:**
+- Modify: `db/clickhouse.go`
+- Modify: `db/clickhouse_test.go`
+- Modify: `logger/queue.go`
+
+**Interfaces:**
+- Produces: `Pool.Acquire() (clickhouse.Conn, uint64, func())`; release is idempotent and closes a retired connection only after its final lease ends.
+- Consumes: logger `poolBackend` acquires one lease per flush snapshot and releases it after the flush attempt.
+
+- [ ] Add tests proving replacement publishes the new connection without closing a leased old connection, then closes the old connection after release; `Close` rejects new leases and waits within its context for active leases.
+- [ ] Run `go test ./db -run 'Test.*Lease|TestClose' -count=1` and confirm failure.
+- [ ] Add per-generation connection state with active count and retired flag; implement `Acquire`; retire instead of immediately closing on replacement; update `Close` to wait for lease drain under context.
+- [ ] Update `poolBackend` snapshot/release handling so every acquired connection is released after a flush attempt.
+- [ ] Run `go test ./db ./tests/logger -count=1` and confirm passing output.
+
+### Task 3: Retain Batches While Database Is Unhealthy
+
+**Files:**
+- Modify: `logger/queue.go`
+- Modify: `tests/logger/queue_test.go`
+
+**Interfaces:**
+- Consumes: existing `Backend.Healthy()` and leased backend snapshots.
+- Produces: worker batches remain reserved and retry after health recovery; shutdown cancellation records and releases unsubmitted entries.
+
+- [ ] Add a test that queues one entry while unhealthy, verifies no prepare call and no failed count, restores health, and verifies exactly one successful prepare/send.
+- [ ] Add a test that shutdown deadline releases a retained unhealthy batch and records it as failed.
+- [ ] Run the two focused tests and confirm they fail.
+- [ ] Change worker flushing so an unhealthy backend returns a retained outcome; wait on a bounded timer or cancellation without consuming additional entries; release only after success/final failure or shutdown.
+- [ ] Run `go test ./tests/logger -count=1` and confirm passing output.
+
+### Task 4: Verification And Security Review
+
+**Files:**
+- Review only: all changed files and their tests.
+
+**Interfaces:**
+- Produces: fresh verification evidence and a list of any new correctness or security issues introduced by the changes.
+
+- [ ] Run `gofmt` on changed Go files.
+- [ ] Run `go test -count=1 ./...`; expect only the pre-existing deployment test concerning ClickHouse host ports to fail.
+- [ ] Run `go test -count=1 ./db ./tests/logger`, `go vet ./...`, and `go build ./...`; require success.
+- [ ] Review lease acquisition/release paths, cancellation, lock ordering, migration identifier construction, and queue budget accounting for new vulnerabilities.

docs/compose/reports/2026-07-10-clickhouse-migration-security-fixes.md

@@ -4,7 +4,7 @@
 
 本次修订覆盖队列并发关闭和单一总 shutdown deadline、响应完整性、请求体 fail-closed、ClickHouse 部分/模糊提交、64 MiB 默认队列字节预算、严格正值配置、响应体总/idle timeout、SSE 协议边界、WebSocket 101 与 shutdown、DSN/TLS、schema 校验、Compose 安全默认、固定 502、可信代理、batch 清理和 server 启动统一清理。
 
-Compose 使用固定 ClickHouse 镜像 `clickhouse/clickhouse-server:25.3.3.42-alpine`，默认不发布 ClickHouse 端口，并在配置展开阶段拒绝空的 `UPSTREAM_URL`、`CLICKHOUSE_URL` 和 `CLICKHOUSE_PASSWORD`。`.env.example` 不提供密码或 DSN 默认值。
+Compose 使用固定 ClickHouse 镜像 `clickhouse/clickhouse-server:26.3`，默认不发布 ClickHouse 端口，并在配置展开阶段拒绝空的 `UPSTREAM_URL`、`CLICKHOUSE_URL` 和 `CLICKHOUSE_PASSWORD`。`.env.example` 不提供密码或 DSN 默认值。
 
 ## 关键取舍
 

docs/compose/reports/db-log-reliability.md

@@ -0,0 +1,54 @@
+---
+feature: db-log-reliability
+status: delivered
+specs: []
+plans:
+  - docs/compose/plans/2026-07-12-db-log-reliability.md
+branch: main
+commits: uncommitted
+---
+
+# Database And Log Reliability - Final Report
+
+## What Was Built
+
+Temporary ClickHouse outages no longer cause worker-held log batches to be immediately discarded. Each worker retains at most one configured-size batch while the backend is unhealthy, preserving the existing global entry and byte budgets and the non-blocking submission policy.
+
+ClickHouse connections now use generation-bound leases for batch writes and health checks. Replaced or closed connections remain alive until active users release them. Existing `proxy_logs` tables can gain missing known columns through fixed additive DDL after table-level compatibility checks.
+
+## Architecture
+
+`db.Pool.Acquire` returns the current connection, its generation, and an idempotent release function. Retired connections are tracked until all leases are released and the underlying close operation completes. `Pool.Close` prevents new leases and waits within its context for active and in-progress closes.
+
+`logger.Queue` retains a full local batch when `Backend.Healthy` is false and stops consuming further channel entries until recovery or shutdown. The pool adapter acquires one connection lease per flush attempt and releases it after the attempt.
+
+`db/migrate.go` validates the table engine, partition key, sorting key, and types of existing required columns before executing static `ADD COLUMN IF NOT EXISTS` statements. It permits unrelated extra columns and revalidates after migration.
+
+### Design Decisions
+
+We kept outage buffering in memory because the existing bounded queue already defines memory ownership and overload behavior; adding a durable WAL would substantially expand scope. Ambiguous `Send` failures remain non-retryable to avoid duplicate records.
+
+We use static column definitions rather than metadata-derived SQL so migration input cannot introduce identifiers or DDL fragments.
+
+## Usage
+
+No configuration or API changes are required. Existing queue size, byte budget, batch size, and ClickHouse settings continue to control operation.
+
+## Verification
+
+`go test -count=1 ./db ./tests/logger`, `go vet ./...`, `go build ./...`, and `git diff --check` pass. `go test -count=1 ./...` has one pre-existing failure in `tests/deployment`: the unchanged Compose file publishes ClickHouse host ports. Race detection remains unavailable because this Windows environment has CGO disabled.
+
+Independent final review found no new or unresolved high/medium-risk issues in the changed reliability paths.
+
+## Journey Log
+
+- [lesson] Connection leases must cover health checks as well as database writes.
+- [pivot] Pool shutdown now tracks in-progress connection closes so repeated close calls cannot report completion early.
+- [lesson] Retaining an unhealthy batch must also stop channel consumption at `batchSize` to prevent recovery spikes.
+- [pivot] Migration validates table-level invariants before any additive DDL to avoid modifying incompatible tables.
+
+## Source Materials
+
+| File | Role | Notes |
+|------|------|-------|
+| `docs/compose/plans/2026-07-12-db-log-reliability.md` | Implementation plan | Complete |

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
 
-迭代 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
 

logger/queue.go

@@ -8,6 +8,7 @@ import (
 	"sync"
 	"sync/atomic"
 	"time"
+	"unsafe"
 
 	"github.com/ClickHouse/clickhouse-go/v2/lib/driver"
 
@@ -46,19 +47,28 @@ func (p clickHouseBatchPreparer) PrepareBatch(ctx context.Context, query string)
 }
 
 type poolBackend struct {
-	pool *db.Pool
+	pool       *db.Pool
+	conn       driver.Conn
+	generation uint64
+	release    func()
 }
 
 func (b poolBackend) PrepareBatch(ctx context.Context, query string) (Batch, error) {
-	conn := b.pool.Get()
+	if b.conn == nil {
-	if conn == nil {
 		return nil, errors.New("pool nil")
 	}
-	return clickHouseBatchPreparer{conn: conn}.PrepareBatch(ctx, query)
+	return clickHouseBatchPreparer{conn: b.conn}.PrepareBatch(ctx, query)
 }
 
 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, release := b.pool.Acquire()
+	return poolBackend{pool: b.pool, conn: conn, generation: generation, release: release}, generation
+}
+func (b poolBackend) MarkUnhealthyGeneration(generation uint64) {
+	b.pool.MarkUnhealthyGeneration(generation)
+}
 
 type Stats struct {
 	Enqueued  uint64
@@ -153,8 +163,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) +
+	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))
+		len(e.RequestHeaders)+len(e.RequestBody)+len(e.ResponseHeaders)+len(e.ResponseBody))
 }
 
 func (q *Queue) Stats() Stats {
@@ -195,9 +205,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 +233,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 +252,7 @@ func (q *Queue) Submit(e *LogEntry) {
 		q.dropped.Add(1)
 		return
 	}
-	entry := cloneLogEntry(e)
+	size := EstimatedBytes(e)
-	size := EstimatedBytes(entry)
 	if !q.mu.TryRLock() {
 		q.dropped.Add(1)
 		return
@@ -246,6 +262,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)
@@ -310,43 +327,77 @@ func (q *Queue) run(workCtx context.Context) {
 	ticker := time.NewTicker(q.batchInterval)
 	defer ticker.Stop()
 
-	flush := func() {
+	flush := func(final bool) bool {
 		if len(batch) == 0 {
-			return
+			return true
 		}
-		entries := make([]*LogEntry, len(batch))
+		if q.backend == nil || !q.backend.Healthy() {
-		for i, item := range batch {
+			if !final {
-			entries[i] = item.entry
+				return false
+			}
+			q.failed.Add(uint64(len(batch)))
+		} else {
+			q.flush(q.flushContext(workCtx), entriesFromBatch(batch))
 		}
-		q.flush(q.flushContext(workCtx), entries)
 		for _, item := range batch {
 			q.release(item.size)
 		}
 		clear(batch)
 		batch = batch[:0]
+		return true
 	}
 
 	for {
+		if len(batch) >= q.batchSize && (q.backend == nil || !q.backend.Healthy()) {
+			if q.isStopped() {
+				flush(true)
+				q.discardQueued()
+				return
+			}
+			select {
+			case <-ticker.C:
+				flush(false)
+			case <-workCtx.Done():
+				flush(true)
+				q.discardQueued()
+				return
+			}
+			continue
+		}
 		select {
 		case item, ok := <-q.ch:
 			if !ok {
-				flush()
+				flush(true)
 				return
 			}
 			batch = append(batch, item)
 			if len(batch) >= q.batchSize {
-				flush()
+				flush(false)
 			}
 		case <-ticker.C:
-			flush()
+			flush(false)
 		case <-workCtx.Done():
-			flush()
+			flush(true)
 			q.discardQueued()
 			return
 		}
 	}
 }
 
+func (q *Queue) isStopped() bool {
+	q.mu.RLock()
+	defer q.mu.RUnlock()
+	return q.stopped
+}
+
+func entriesFromBatch(batch []queuedEntry) []*LogEntry {
+	entries := make([]*LogEntry, len(batch))
+	for i, item := range batch {
+		entries[i] = item.entry
+	}
+	return entries
+}
+
 func (q *Queue) flushContext(workCtx context.Context) context.Context {
 	q.mu.RLock()
 	defer q.mu.RUnlock()
@@ -357,23 +408,24 @@ func (q *Queue) flushContext(workCtx context.Context) context.Context {
 }
 
 func (q *Queue) flush(ctx context.Context, entries []*LogEntry) {
-	if q.backend == nil || !q.backend.Healthy() {
-		q.failed.Add(uint64(len(entries)))
-		return
-	}
-
 	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, release := backendSnapshot(q.backend)
+		failedGeneration = generation
+		result := Flush(ctx, attemptBackend, retry)
+		if release != nil {
+			release()
+		}
 		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 +440,41 @@ 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, func()) {
+	if versioned, ok := backend.(generationBackend); ok {
+		snapshot, generation := versioned.Snapshot()
+		if leased, ok := snapshot.(interface{ Release() }); ok {
+			return snapshot, generation, leased.Release
+		}
+		return snapshot, generation, nil
+	}
+	return backend, 0, nil
+}
+
+func (b poolBackend) Release() {
+	if b.release != nil {
+		b.release()
+	}
+}
+
+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 +510,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
 }

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,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)
+	runErr := waitForShutdown(rootCtx.Done(), serverErr)
-	signal.Notify(stop, syscall.SIGINT, syscall.SIGTERM)
-	runErr := waitForShutdown(stop, serverErr)
-	signal.Stop(stop)
 
 	// HTTP、升级连接、日志排空和数据库共享同一个关闭总预算。
 	shutdownCtx, shutdownCancel := context.WithTimeout(context.Background(), 30*time.Second)

main_test.go

@@ -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 := make(chan struct{})
-		stop <- os.Interrupt
+		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)
 		}
 	})

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.ContentLength == 0 {
+	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"
@@ -25,14 +26,17 @@ type LogSubmitter interface {
 
 // Handler 构造反代 HTTP handler。
 type Handler struct {
-	rp           *httputil.ReverseProxy
+	rp              *httputil.ReverseProxy
-	filter       *config.Filter
+	filter          *config.Filter
-	queue        LogSubmitter
+	queue           LogSubmitter
-	maxBodyBytes int64
+	maxBodyBytes    int64
-	trusted      []netip.Prefix
+	maxRequestBytes int64
-	connMu       sync.Mutex
+	trusted         []netip.Prefix
-	conns        map[net.Conn]struct{}
+	connMu          sync.Mutex
-	connChanged  chan struct{}
+	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,
+		rp:              rp,
-		filter:       filter,
+		filter:          filter,
-		queue:        queue,
+		queue:           queue,
-		maxBodyBytes: maxBody,
+		maxBodyBytes:    maxBody,
-		trusted:      append([]netip.Prefix(nil), opts.TrustedProxies...),
+		maxRequestBytes: opts.MaxRequestBytes,
-		conns:        make(map[net.Conn]struct{}),
+		trusted:         append([]netip.Prefix(nil), opts.TrustedProxies...),
-		connChanged:  make(chan struct{}),
+		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{})

proxy/sse.go

@@ -552,11 +552,16 @@ func (g *geminiChunk) UnmarshalJSON(data []byte) error {
 }
 
 type geminiCandidate struct {
-	raw     map[string]json.RawMessage
+	raw   map[string]json.RawMessage
-	index   int
+	index int
-	role    string
+	role  string
-	text    strings.Builder
+	parts []*geminiPart
-	partRaw map[string]json.RawMessage
+}
+
+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 {
+				var contentRaw struct {
-					assembled.text.WriteString(part.Text)
+					Parts []map[string]json.RawMessage `json:"parts"`
 				}
-				if len(assembled.partRaw) == 0 {
+				if rawContent, ok := candidate.Raw["content"]; ok {
-					var contentRaw struct {
+					_ = json.Unmarshal(rawContent, &contentRaw)
-						Parts []map[string]json.RawMessage `json:"parts"`
+				}
+				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()}}}
+		parts := make([]any, 0, len(candidate.parts))
-		if len(candidate.partRaw) > 0 {
+		for _, part := range candidate.parts {
-			partRaw := cloneRawMap(candidate.partRaw)
+			partRaw := cloneRawMap(part.raw)
-			partRaw["text"] = mustJSON(candidate.text.String())
+			if part.text.Len() > 0 || len(partRaw) == 0 {
-			contentRaw["parts"] = []any{rawMapToMap(partRaw)}
+				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()
@@ -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
 	}

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/deployment/compose_test.go

@@ -23,8 +23,8 @@ func TestComposeUsesSafeDeploymentDefaults(t *testing.T) {
 		}
 	}
 
-	if !regexp.MustCompile(`(?m)^\s+image: clickhouse/clickhouse-server:\d+\.\d+\.\d+\.\d+-alpine\s*$`).MatchString(compose) {
+	if !regexp.MustCompile(`(?m)^\s+image: clickhouse/clickhouse-server:26\.3\s*$`).MatchString(compose) {
-		t.Error("ClickHouse image must use an exact version tag")
+		t.Error("ClickHouse image must use version 26.3")
 	}
 
 	serviceStart := strings.LastIndex(compose, "  thief_clickhouse:")
@@ -32,8 +32,16 @@ func TestComposeUsesSafeDeploymentDefaults(t *testing.T) {
 		t.Fatal("compose.yml is missing the thief_clickhouse service")
 	}
 	clickhouseService := compose[serviceStart:]
-	if strings.Contains(clickhouseService, "\n    ports:") {
+	for _, port := range []string{
-		t.Error("ClickHouse must not publish host ports by default")
+		`${CLICKHOUSE_LISTEN_IP:-127.0.0.1}:${CLICKHOUSE_HTTP_PORT:-8123}:8123`,
+		`${CLICKHOUSE_LISTEN_IP:-127.0.0.1}:${CLICKHOUSE_NATIVE_PORT:-9000}:9000`,
+	} {
+		if !strings.Contains(clickhouseService, port) {
+			t.Errorf("ClickHouse port must bind to loopback by default with %q", port)
+		}
+	}
+	if !strings.Contains(envExample, "CLICKHOUSE_LISTEN_IP=127.0.0.1\n") {
+		t.Error(".env.example must bind ClickHouse to loopback by default")
 	}
 
 	for _, emptySecret := range []string{"CLICKHOUSE_URL=\n", "CLICKHOUSE_PASSWORD=\n"} {

tests/logger/queue_test.go

@@ -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
+	calls     int
-	queries []string
+	queries   []string
-	batches []*fakeBatch
+	batches   []*fakeBatch
-	scripts []batchScript
+	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 {
+	if got := logger.EstimatedBytes(entry); got <= 55 {
-		t.Fatalf("EstimatedBytes=%d want %d", got, want)
+		t.Fatalf("EstimatedBytes=%d must include fixed entry overhead", got)
 	}
 }
 
@@ -358,6 +381,41 @@ func TestQueueSubmitOwnsEntrySnapshot(t *testing.T) {
 	}
 }
 
+func TestQueueRetainsBatchUntilBackendRecovers(t *testing.T) {
+	backend := newQueueBackend(func(context.Context, string) (logger.Batch, error) {
+		return &fakeBatch{}, nil
+	})
+	backend.healthy.Store(false)
+	q := logger.NewQueueWithBackend(backend, 1, 1, 1, 10*time.Millisecond)
+	q.Start(context.Background())
+	q.Submit(&logger.LogEntry{RequestID: "retained"})
+
+	time.Sleep(50 * time.Millisecond)
+	if calls := backend.calls.Load(); calls != 0 {
+		t.Fatalf("PrepareBatch calls=%d while unhealthy", calls)
+	}
+	if stats := q.Stats(); stats.Failed != 0 || stats.Bytes == 0 {
+		t.Fatalf("unhealthy stats=%+v want retained bytes and no failure", stats)
+	}
+
+	backend.healthy.Store(true)
+	deadline := time.Now().Add(time.Second)
+	for backend.calls.Load() == 0 && time.Now().Before(deadline) {
+		time.Sleep(10 * time.Millisecond)
+	}
+	ctx, cancel := context.WithTimeout(context.Background(), time.Second)
+	defer cancel()
+	if err := q.Stop(ctx); err != nil {
+		t.Fatalf("Stop: %v", err)
+	}
+	if calls := backend.calls.Load(); calls != 1 {
+		t.Fatalf("PrepareBatch calls=%d want 1", calls)
+	}
+	if stats := q.Stats(); stats.Failed != 0 || stats.Bytes != 0 {
+		t.Fatalf("recovered stats=%+v", stats)
+	}
+}
+
 func TestQueueByteBudgetTracksOwnedEntrySnapshot(t *testing.T) {
 	batch := &fakeBatch{}
 	backend := newQueueBackend(func(context.Context, string) (logger.Batch, error) {
@@ -393,7 +451,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 +463,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 +497,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 +595,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)
 	}

tests/proxy/robustness_test.go

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

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

tests/proxy/websocket_test.go

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