Add tests for proxy and router digest functionality

- Implement comprehensive tests for the ingress proxy, validating request handling and response generation.
- Introduce usage observer tests to ensure accurate tracking of token usage across various scenarios.
- Create a new digest module in the auto router to manage conversation summaries, enforcing size limits and maintaining relevant context.
- Add tests for the digest functionality to verify that it correctly summarizes user interactions and adheres to defined byte limits.

Author: nextroad-dev

README.md

@@ -17,6 +17,14 @@ docker run -d \
 
 打开 <http://127.0.0.1:8080/admin/>，按页面提示设置管理员密码。之后在管理页面添加提供商和模型、配置模型组，并创建推理密钥。上游 API 密钥与管理员密码在管理页面配置。
 
+首次设置需要证明你是部署者：尚未设置密码时，服务启动会在日志中打印一次性初始化令牌 `bootstrap_token`。从非本机连接（包括经 Docker 端口映射或反向代理访问）设置密码时，页面会要求填写该令牌；设置成功后令牌立即失效。获取方式：
+
+```sh
+docker logs auto-router 2>&1 | grep bootstrap_token
+```
+
+直接在本机运行二进制并通过 loopback 访问时无需令牌。
+
 数据保存在 Docker 命名卷 `auto-router-data` 中，容器重建或升级不会删除该卷。升级到最新镜像：
 
 ```sh
@@ -28,6 +36,22 @@ docker rm auto-router
 
 不要删除数据卷；删除 `auto-router-data` 会同时删除数据库及管理配置。默认端口只绑定本机。若需对外提供服务，请先配置 TLS 反向代理，不要直接将管理页面暴露到公网。
 
+### 反向代理与启动参数
+
+以下参数在启动时通过命令行设置（Docker 中直接追加在镜像名之后，镜像入口已包含 `-listen 0.0.0.0:8080`）：
+
+| 参数 | 默认 | 说明 |
+| --- | --- | --- |
+| `-trusted-proxies` | 空 | 逗号分隔的反向代理 CIDR 或 IP。只有来自这些地址的 `X-Forwarded-For` / `X-Forwarded-Proto` 才会被采信；为空时一律以 TCP 对端地址为客户端地址。多级代理时从右向左跳过可信跳，取第一个不可信地址。`Forwarded`、`X-Forwarded-*`、`X-Real-IP` 永远不会转发给上游提供商。 |
+| `-secure-cookies` | `auto` | 管理会话 Cookie 的 `Secure` 属性：`auto` 在 TLS 直连或可信代理声明 `https` 时设置；`always` / `never` 强制。 |
+| `-upstream-body-timeout` | `10m` | 非流式上游响应在收到响应头后，两次读取之间允许的最长停顿；`0` 关闭。 |
+| `-stream-idle-timeout` | `0` | 流式上游响应两次读取之间允许的最长停顿；默认关闭，避免截断长时间推理。 |
+| `-max-buffered-response-bytes` | 32 MiB | 需要整体缓冲以转换协议（Anthropic/Gemini 非流式）的上游响应上限；超出返回 502 `upstream_response_too_large`。 |
+
+例如 Nginx 与服务同机、对外提供 HTTPS：`-trusted-proxies 127.0.0.1`。Docker 网络内的代理请填写其所在网段，例如 `-trusted-proxies 172.16.0.0/12`。
+
+管理登录限流按客户端地址分别计数（每地址每分钟 10 次失败），另有全局每分钟 100 次失败的兜底上限。
+
 ## 自动路由
 
 将请求中的 `model` 设为 `auto`，即可请求自动路由。启用 Jev 时，它只在当前有合格成员的 `simple`、`medium`、`complex` 组中选择任务组；本地策略先按请求特征、模型能力、上下文窗口及硬限制过滤成员，再严格按组内配置顺序选择主模型。不会跨组重试，组内也不做评分排序。

cmd/auto-router/main.go

@@ -22,6 +22,7 @@ import (
 	"github.com/nextroad-dev/Auto-Router/internal/logging"
 	"github.com/nextroad-dev/Auto-Router/internal/models"
 	"github.com/nextroad-dev/Auto-Router/internal/modelsdev"
+	"github.com/nextroad-dev/Auto-Router/internal/netx"
 	"github.com/nextroad-dev/Auto-Router/internal/providers"
 	"github.com/nextroad-dev/Auto-Router/internal/providers/direct"
 	"github.com/nextroad-dev/Auto-Router/internal/proxy"
@@ -87,6 +88,11 @@ func run(ctx context.Context, args []string, output io.Writer) error {
 	flags := flag.NewFlagSet("auto-router", flag.ContinueOnError)
 	flags.SetOutput(output)
 	listenAddress := flags.String("listen", "", "override the default HTTP listen address")
+	trustedProxies := flags.String("trusted-proxies", "", "comma-separated reverse proxy CIDRs or IPs whose X-Forwarded-For/X-Forwarded-Proto are trusted (default: none)")
+	secureCookies := flags.String("secure-cookies", "", "session cookie Secure attribute: auto (HTTPS or trusted proxy https), always or never (default auto)")
+	upstreamBodyTimeout := flags.Duration("upstream-body-timeout", -1, "maximum stall between reads of a non-stream upstream body; 0 disables (default 10m)")
+	streamIdleTimeout := flags.Duration("stream-idle-timeout", -1, "maximum stall between reads of a streaming upstream body; 0 disables (default 0)")
+	maxBufferedResponseBytes := flags.Int64("max-buffered-response-bytes", 0, "maximum converted (buffered) upstream response size in bytes (default 32 MiB)")
 	configPath := flags.String("config", "", "deprecated; ignored")
 	syncModels := flags.Bool("sync-models", false, "deprecated; use the WebUI to synchronize models.dev")
 	recoverAdmin := flags.Bool("recover-admin", false, "offline recovery: rotate administrator credentials while holding the fixed listener port")
@@ -147,6 +153,24 @@ func run(ctx context.Context, args []string, output io.Writer) error {
 	if *listenAddress != "" {
 		cfg.HTTP.Address = *listenAddress
 	}
+	if *trustedProxies != "" {
+		cfg.HTTP.TrustedProxies = strings.Split(*trustedProxies, ",")
+	}
+	if *secureCookies != "" {
+		cfg.Admin.SecureCookies = *secureCookies
+	}
+	if *upstreamBodyTimeout >= 0 {
+		cfg.HTTP.UpstreamBodyTimeout = config.Duration(*upstreamBodyTimeout)
+	}
+	if *streamIdleTimeout >= 0 {
+		cfg.HTTP.StreamIdleTimeout = config.Duration(*streamIdleTimeout)
+	}
+	if *maxBufferedResponseBytes != 0 {
+		cfg.HTTP.MaxBufferedResponseBytes = *maxBufferedResponseBytes
+	}
+	if err := cfg.Validate(); err != nil {
+		return err
+	}
 	sources := config.Sources{}
 	if *recoverAdmin {
 		return runAdminRecovery(ctx, cfg, output)
@@ -236,6 +260,11 @@ func run(ctx context.Context, args []string, output io.Writer) error {
 	// Build the direct provider executor set from the catalog already loaded from
 	// SQLite. Providers without a base_url remain routable in the catalog but
 	// return provider_not_configured when selected.
+	globalExecutorRegistry.Configure(direct.Options{
+		StreamIdleTimeout:        time.Duration(cfg.HTTP.StreamIdleTimeout),
+		BodyTimeout:              time.Duration(cfg.HTTP.UpstreamBodyTimeout),
+		MaxBufferedResponseBytes: cfg.HTTP.MaxBufferedResponseBytes,
+	})
 	globalExecutorRegistry.Build(store.Load())
 	executor := &providerRouter{registry: globalExecutorRegistry}
 	defer func() {
@@ -277,11 +306,15 @@ func run(ctx context.Context, args []string, output io.Writer) error {
 	if err != nil {
 		return err
 	}
+	bootstrap, err := newBootstrapToken(ctx, db, logger)
+	if err != nil {
+		return err
+	}
 	syncGate := newRegistrySyncGate()
 	server := newHTTPServer(serverAssembly{
 		Config: cfg, Store: store, DB: db, Checker: db, Executor: executor, Jev: jevClient,
 		Recorder: recorder, Settings: settingsStore, Authenticator: authenticator, Sessions: sessions,
-		Groups: &groupStore, ResolveModelMetadata: modelMetadataResolver.Resolve,
+		Groups: &groupStore, ResolveModelMetadata: modelMetadataResolver.Resolve, Bootstrap: bootstrap,
 		SyncRegistry: func(syncCtx context.Context) (any, error) {
 			return syncGate.Do(syncCtx, func(syncCtx context.Context) (any, error) {
 				current := settingsStore.Config()
@@ -309,6 +342,27 @@ func run(ctx context.Context, args []string, output io.Writer) error {
 	return serve(ctx, server, listener, time.Duration(cfg.HTTP.ShutdownTimeout), logger)
 }
 
+// newBootstrapToken issues the one-time first-run setup token when no owner
+// password exists yet, and prints it to the process log: reading the log is the
+// proof a remote operator gives that they control this deployment. A local
+// browser does not need it.
+func newBootstrapToken(ctx context.Context, db *sql.DB, logger *slog.Logger) (*api.BootstrapToken, error) {
+	set, err := storage.PasswordSet(ctx, db)
+	if err != nil {
+		return nil, fmt.Errorf("read password setup state: %w", err)
+	}
+	if set {
+		return nil, nil
+	}
+	token, value, err := api.NewBootstrapToken()
+	if err != nil {
+		return nil, fmt.Errorf("generate bootstrap token: %w", err)
+	}
+	logger.Warn("the management password is not set; remote first-run setup requires this one-time bootstrap token",
+		"bootstrap_token", value, "setup", "/admin/")
+	return token, nil
+}
+
 // providerRouter selects the direct executor for the provider chosen by the
 // catalog/router. It does not execute requests itself; the registry lookup and
 // the executor call are deliberately one small boundary.
@@ -325,7 +379,9 @@ func (r *providerRouter) Do(ctx context.Context, request *providers.Request) (*p
 	if r == nil || r.registry == nil {
 		return nil, fmt.Errorf("%w: no direct provider registry", providers.ErrProviderNotConfigured)
 	}
-	executor, ok := r.registry.Get(request.ProviderKey)
+	// The executor comes from the catalog generation the request was routed
+	// against, never from whatever is newest when the call is made.
+	executor, ok := r.registry.GetGeneration(request.CatalogGeneration, request.ProviderKey)
 	if !ok {
 		return nil, fmt.Errorf("%w: provider %q has no configured executor", providers.ErrProviderNotConfigured, request.ProviderKey)
 	}
@@ -403,8 +459,9 @@ func reloadCatalog(ctx context.Context, db *sql.DB, logger *slog.Logger, store *
 		// to publish an empty one.
 		return err
 	}
-	published := store.Swap(catalog)
-	globalExecutorRegistry.Build(published)
+	// The executor set is built inside the publication, before the catalog is
+	// visible, so catalog and executors form one request-consistent generation.
+	published := store.SwapWith(catalog, globalExecutorRegistry.Build)
 	logCatalogSnapshot(logger, published, true)
 	return nil
 }
@@ -551,6 +608,9 @@ type serverAssembly struct {
 	SyncRegistry         func(context.Context) (any, error)
 	ResolveModelMetadata func(context.Context, string, string) (modelsdev.MetadataMatch, bool, error)
 	Groups               *atomic.Pointer[storage.ModelGroups]
+	// Bootstrap authorizes first-run setup from a non-local client. Nil when the
+	// owner password already exists.
+	Bootstrap *api.BootstrapToken
 }
 
 // newHTTPServer builds the HTTP handler. The debug endpoints are mounted only when
@@ -587,6 +647,12 @@ func newHTTPServer(assembly serverAssembly, logger *slog.Logger) *http.Server {
 	routeDebugHandler := gateDebugHandler(api.NewDebugRoute(api.RouteDebugOptions{
 		Router: autoRouter, MaxRequestBytes: cfg.HTTP.MaxRequestBytes,
 	}), settingsStore, false)
+	// The configuration was validated at startup, so the list parses.
+	trustedProxies, err := netx.ParseTrustedProxies(cfg.HTTP.TrustedProxies)
+	if err != nil {
+		logger.Error("http.trusted_proxies could not be parsed; no proxy is trusted", "error", err)
+		trustedProxies = nil
+	}
 	authenticator := assembly.Authenticator
 	if authenticator == nil {
 		authenticator, _ = api.NewDatabaseAuthenticator(nil, logger)
@@ -604,6 +670,7 @@ func newHTTPServer(assembly serverAssembly, logger *slog.Logger) *http.Server {
 		Recorder:           recorder,
 		Live:               liveProxyValues{store: settingsStore, groups: assembly.Groups},
 		OverrideAuthorizer: api.OverrideAuthorizer(authenticator),
+		TrustedProxies:     trustedProxies,
 	})
 	var adminHandler http.Handler
 	// sessionStore is created only when the management surface is mounted, and it is
@@ -648,6 +715,9 @@ func newHTTPServer(assembly serverAssembly, logger *slog.Logger) *http.Server {
 			},
 			SyncRegistry:         assembly.SyncRegistry,
 			ResolveModelMetadata: assembly.ResolveModelMetadata,
+			TrustedProxies:       trustedProxies,
+			SecureCookies:        cfg.Admin.SecureCookies,
+			Bootstrap:            assembly.Bootstrap,
 		})
 		sessionStore = sessions
 		logger.Info("the Admin API is mounted",

docs/openapi.yaml

@@ -497,9 +497,12 @@ paths:
             application/json:
               schema:
                 type: object
-                required: [password_set]
+                required: [password_set, bootstrap_token_required]
                 properties:
                   password_set: {type: boolean}
+                  bootstrap_token_required:
+                    type: boolean
+                    description: True when the password is not set and this client is not local, so setup must present the one-time bootstrap token printed in the service log.
         '503':
           $ref: '#/components/responses/AdminError'
 
@@ -508,7 +511,7 @@ paths:
       tags: [Setup]
       summary: Set the initial owner password
       operationId: setInitialPassword
-      description: Available only until the initial password is set. Browser requests must be same-origin; requests identified as cross-site are rejected. Password length is 12–1024 bytes.
+      description: Available only until the initial password is set. Browser requests must be same-origin; requests identified as cross-site are rejected. A client that is not on the local machine (a loopback peer without forwarding headers, or a loopback client reported by a trusted proxy) must also present the one-time bootstrap token printed in the service log at startup; it is invalidated by a successful setup. Password length is 12–1024 bytes.
       security: []
       requestBody:
         required: true
@@ -519,6 +522,7 @@ paths:
               required: [password]
               properties:
                 password: {type: string, minLength: 12, maxLength: 1024, writeOnly: true}
+                bootstrap_token: {type: string, writeOnly: true, description: Required from a non-local client.}
       responses:
         '201':
           description: Initial password set.
@@ -2423,6 +2427,7 @@ components:
         - truncated_evidence
         - upstream_unavailable
         - upstream_timeout
+        - upstream_response_too_large
         - client_closed_request
         - internal_error
         - invalid_api_key
@@ -2447,6 +2452,7 @@ components:
         - invalid_api_key
         - insufficient_scope
         - cross_site_request
+        - bootstrap_token_required
         - too_many_attempts
         - request_too_large
         - unknown_model

internal/api/admin.go

@@ -19,6 +19,7 @@ import (
 	"github.com/nextroad-dev/Auto-Router/internal/logging"
 	"github.com/nextroad-dev/Auto-Router/internal/models"
 	"github.com/nextroad-dev/Auto-Router/internal/modelsdev"
+	"github.com/nextroad-dev/Auto-Router/internal/netx"
 	"github.com/nextroad-dev/Auto-Router/internal/proxy"
 	"github.com/nextroad-dev/Auto-Router/internal/settings"
 	"github.com/nextroad-dev/Auto-Router/internal/storage"
@@ -103,6 +104,15 @@ type AdminOptions struct {
 	StartedAt time.Time
 	// Logger receives one INFO line per accepted write. Nil disables it.
 	Logger *slog.Logger
+	// TrustedProxies decides whose forwarding headers are believed for the login
+	// throttle's source address, the local-client check of first-run setup and the
+	// browser-facing scheme. Nil trusts nobody.
+	TrustedProxies *netx.TrustedProxies
+	// SecureCookies is admin.secure_cookies: auto, always or never. Empty is auto.
+	SecureCookies string
+	// Bootstrap authorizes first-run setup from a non-local client. Nil means no
+	// token was issued, so only a local client can set the initial password.
+	Bootstrap *BootstrapToken
 }
 
 // adminHandler serves /admin/v1.
@@ -121,6 +131,9 @@ type adminHandler struct {
 	maxPageSize          int
 	startedAt            time.Time
 	logger               *slog.Logger
+	trustedProxies       *netx.TrustedProxies
+	secureCookies        string
+	bootstrap            *BootstrapToken
 	// authenticator validates a login submission through the same constant-time
 	// loop the header path uses. It is nil when the surface was built without one.
 	authenticator *Authenticator
@@ -150,6 +163,9 @@ func NewAdmin(options AdminOptions) http.Handler {
 		logger:               options.Logger,
 		authenticator:        options.Authenticator,
 		sessions:             options.Sessions,
+		trustedProxies:       options.TrustedProxies,
+		secureCookies:        options.SecureCookies,
+		bootstrap:            options.Bootstrap,
 	}
 	if handler.authenticator != nil && handler.authenticator.Enabled() && handler.sessions == nil {
 		// The management surface accepts a session only when it can validate one. A

internal/api/admin_password.go

@@ -18,13 +18,17 @@ func (h *adminHandler) handleSetupStatus(w http.ResponseWriter, r *http.Request)
 		storageFailure(w, h, "read password setup state", err)
 		return
 	}
+	// bootstrap_token_required tells the setup page whether it must ask for the
+	// token printed in the process log. It is only meaningful before setup.
 	writeAdminJSON(w, http.StatusOK, struct {
-		PasswordSet bool `json:"password_set"`
-	}{set})
+		PasswordSet            bool `json:"password_set"`
+		BootstrapTokenRequired bool `json:"bootstrap_token_required"`
+	}{set, !set && !h.trustedProxies.IsLocalRequest(r)})
 }
 
 type setPasswordRequest struct {
-	Password *string `json:"password"`
+	Password       *string `json:"password"`
+	BootstrapToken *string `json:"bootstrap_token"`
 }
 
 // handleSetupPassword sets the owner's password exactly once. Registration is
@@ -38,18 +42,34 @@ func (h *adminHandler) handleSetupPassword(w http.ResponseWriter, r *http.Reques
 		writeAdminError(w, http.StatusForbidden, "cross_site_request", "the request did not originate from this management surface", "")
 		return
 	}
-	if h.sessions != nil && !h.sessions.AllowLogin() {
+	source := h.trustedProxies.ClientIP(r)
+	if h.sessions != nil && !h.sessions.AllowLogin(source) {
 		writeAdminError(w, http.StatusTooManyRequests, "too_many_attempts", "too many password setup attempts; try again later", "")
 		return
 	}
 	var body setPasswordRequest
 	if _, err := readAdminBody(r, &body); err != nil {
-		if h.sessions != nil {
-			h.sessions.RecordFailure()
-		}
+		h.sessions.RecordFailure(source)
 		writeBodyError(w, err)
 		return
 	}
+	// Before a password exists there is no credential to check, so a caller that
+	// is not on this machine must prove it can read the process log by presenting
+	// the one-time bootstrap token. A local caller is trusted as the operator.
+	if !h.trustedProxies.IsLocalRequest(r) {
+		presented := ""
+		if body.BootstrapToken != nil {
+			presented = *body.BootstrapToken
+		}
+		if !h.bootstrap.Verify(presented) {
+			h.sessions.RecordFailure(source)
+			if h.logger != nil {
+				h.logger.Warn("initial password setup refused", "path", r.URL.Path, "reason", "bootstrap_token_required")
+			}
+			writeAdminError(w, http.StatusForbidden, "bootstrap_token_required", "initial setup from a remote client requires the bootstrap token printed in the service log", "bootstrap_token")
+			return
+		}
+	}
 	if body.Password == nil {
 		writeAdminError(w, http.StatusBadRequest, "invalid_request", "the request body must carry a string password field", "password")
 		return
@@ -65,9 +85,8 @@ func (h *adminHandler) handleSetupPassword(w http.ResponseWriter, r *http.Reques
 		}
 		return
 	}
-	if h.sessions != nil {
-		h.sessions.RecordSuccess()
-	}
+	h.bootstrap.Consume()
+	h.sessions.RecordSuccess(source)
 	writeAdminJSON(w, http.StatusCreated, struct {
 		PasswordSet bool `json:"password_set"`
 	}{true})
@@ -117,19 +136,8 @@ func (h *adminHandler) handlePasswordChange(w http.ResponseWriter, r *http.Reque
 	if h.sessions != nil {
 		h.sessions.Clear()
 	}
-	clearSessionCookie(w)
+	h.clearSessionCookie(w, r)
 	writeAdminJSON(w, http.StatusOK, struct {
 		Changed bool `json:"changed"`
 	}{true})
 }
-
-func clearSessionCookie(w http.ResponseWriter) {
-	http.SetCookie(w, &http.Cookie{
-		Name:     sessionCookieName,
-		Value:    "",
-		Path:     "/admin",
-		HttpOnly: true,
-		SameSite: http.SameSiteStrictMode,
-		MaxAge:   -1,
-	})
-}