fix: session affinity, claim risk-control, proxy pool health, admission false-kill (with tests)

internal/server/accounts.go

@@ -291,6 +291,101 @@ func (a *App) importLocal(ctx context.Context, b M) (M, error) {
 	}
 	return M{"ok": n > 0, "imported": n, "results": []any{result}}, nil
 }
+func (a *App) importRegistration(ctx context.Context, b M) (M, error) {
+	zcodeJWT := strings.TrimSpace(str(b["zcodeJwt"]))
+	webJWT := strings.TrimSpace(str(b["webJwt"]))
+	businessToken := strings.TrimSpace(str(b["businessToken"]))
+	if zcodeJWT == "" {
+		return M{"ok": false, "error": "缺少 zcodeJwt"}, nil
+	}
+	claims := jwtPayload(zcodeJWT)
+	identity := str(first(claims["user_id"], claims["sub"]))
+	if identity == "" {
+		return M{"ok": false, "error": "zcodeJwt 缺少 user_id/sub"}, nil
+	}
+	provider := family(b["provider"])
+	if provider != "zai" {
+		provider = "bigmodel"
+	}
+	deviceMid := strings.TrimSpace(str(b["deviceMid"]))
+	if deviceMid == "" {
+		deviceMid = uuid()
+	}
+	if existing := a.row(
+		"SELECT * FROM accounts WHERE customer_id=?",
+		identity,
+	); existing != nil {
+		tokens := M{"zcode_jwt": zcodeJWT, "active_provider": provider}
+		if webJWT != "" {
+			tokens["web_jwt"] = webJWT
+		}
+		if businessToken != "" {
+			tokens["business_token"] = businessToken
+		}
+		if e := a.setTokens(str(existing["id"]), tokens); e != nil {
+			return nil, e
+		}
+		_ = a.updateAccount(str(existing["id"]), M{
+			"device_mid": deviceMid,
+			"status":     "active",
+			"last_error": nil,
+		})
+		verified, e := a.verify(ctx, a.account(str(existing["id"])))
+		if e != nil {
+			verified = M{"ok": false, "reason": e.Error()}
+		}
+		result := M{
+			"ok":       truth(verified["ok"]),
+			"updated":  true,
+			"id":       existing["id"],
+			"status":   verified["status"],
+			"reason":   verified["reason"],
+			"verified": verified["ok"],
+		}
+		a.scheduleNext(str(existing["id"]), time.Second)
+		return result, nil
+	}
+	tokens := M{
+		"zcode_jwt":       zcodeJWT,
+		"active_provider": provider,
+	}
+	if webJWT != "" {
+		tokens["web_jwt"] = webJWT
+	}
+	if businessToken != "" {
+		tokens["business_token"] = businessToken
+	}
+	label := strings.TrimSpace(str(b["label"]))
+	if label == "" {
+		label = "pro-zai-" + identity[max(0, len(identity)-6):]
+	}
+	account, e := a.insertAccount(M{
+		"label":       label,
+		"family":      provider,
+		"customer_id": identity,
+		"device_mid":  deviceMid,
+		"source":      "pro-zai",
+		"status":      "active",
+	}, tokens)
+	if e != nil {
+		return nil, e
+	}
+	verified, e := a.verify(ctx, account)
+	if e != nil {
+		verified = M{"ok": false, "reason": e.Error()}
+	}
+	result := M{
+		"ok":       truth(verified["ok"]),
+		"updated":  false,
+		"id":       account["id"],
+		"status":   verified["status"],
+		"reason":   verified["reason"],
+		"verified": verified["ok"],
+	}
+	a.scheduleNext(str(account["id"]), time.Second)
+	return result, nil
+}
+
 func (a *App) importCookie(ctx context.Context, b M, target M) (M, error) {
 	token := strings.TrimSpace(str(b["webJwt"]))
 	p := jwtPayload(token)

internal/server/admin.go

@@ -113,6 +113,9 @@ func (a *App) adminRoute(r *http.Request, b M) (any, int, error) {
 	case "POST /api/accounts/import-local":
 		v, e := a.importLocal(ctx, b)
 		return wrap(v, e)
+	case "POST /api/accounts/import-registration":
+		v, e := a.importRegistration(ctx, b)
+		return wrap(v, e)
 	case "POST /api/accounts/import-cookie":
 		v, e := a.importCookie(ctx, b, nil)
 		return wrap(v, e)
@@ -276,7 +279,7 @@ func (a *App) adminRoute(r *http.Request, b M) (any, int, error) {
 		return a.updateStatus(), 202, nil
 	case "GET /api/settings":
 		out := M{}
-		for _, k := range []string{"poll_interval_sec", "poll_jitter_sec", "poll_concurrency", "daily_claim_enabled", "daily_claim_time", "default_reasoning_level", "http", "outbound_proxy", "inference_driver", "app_version", "zcode_version_auto_update", "zcode_version_check_interval_sec", "aliases", "retry", "web_port", "gateway_port"} {
+		for _, k := range []string{"poll_interval_sec", "poll_jitter_sec", "poll_concurrency", "daily_claim_enabled", "daily_claim_time", "conversation_affinity_minutes", "default_reasoning_level", "http", "outbound_proxy", "inference_driver", "app_version", "zcode_version_auto_update", "zcode_version_check_interval_sec", "aliases", "retry", "web_port", "gateway_port"} {
 			out[k] = cfg[k]
 		}
 		over := obj(a.setting("config_overrides", M{}))
@@ -341,7 +344,7 @@ func (a *App) adminRoute(r *http.Request, b M) (any, int, error) {
 			return M{"ok": false, "error": e.Error()}, 400, nil
 		}
 		allowed := map[string]bool{}
-		for _, k := range strings.Fields("poll_interval_sec poll_jitter_sec poll_concurrency daily_claim_enabled daily_claim_time default_reasoning_level app_version zcode_version_auto_update zcode_version_check_interval_sec aliases retry inference_driver http outbound_proxy") {
+		for _, k := range strings.Fields("poll_interval_sec poll_jitter_sec poll_concurrency daily_claim_enabled daily_claim_time conversation_affinity_minutes default_reasoning_level app_version zcode_version_auto_update zcode_version_check_interval_sec aliases retry inference_driver http outbound_proxy") {
 			allowed[k] = true
 		}
 		a.configMu.Lock()
@@ -600,6 +603,9 @@ func validateSettings(b M) error {
 	if v, present := b["daily_claim_time"]; present && !validDailyClaimTime(str(v)) {
 		return fmt.Errorf("daily_claim_time 须为 00:00–23:59 的 HH:MM 时间")
 	}
+	if v, present := b["conversation_affinity_minutes"]; present && (num(v) < 0 || num(v) > 10080) {
+		return fmt.Errorf("conversation_affinity_minutes 须为 0 到 10080 的分钟数")
+	}
 	if v, ok := b["default_reasoning_level"]; ok && !validLevel(str(v)) {
 		return fmt.Errorf("思考档位须为 low/high/max")
 	}

internal/server/admission_rotation_test.go

@@ -0,0 +1,93 @@
+package server
+
+import (
+	"context"
+	"io"
+	"net/http"
+	"testing"
+	"time"
+)
+
+// The admission timer guards the local queue only. Once the upstream answered
+// (response headers arrived) the request is in flight and first_token_timeout
+// owns the wait; the admission timer must not cancel the prefill.
+func TestAdmissionTimerStopsWhenUpstreamAnswers(t *testing.T) {
+	app := registrationTestApp(t)
+	transport := registrationRoundTrip(func(r *http.Request) (*http.Response, error) {
+		// Headers return immediately; the first SSE event takes 150ms, which
+		// is the prefill window the admission timer (100ms) must not cancel.
+		pr, pw := io.Pipe()
+		go func() {
+			defer pw.Close()
+			time.Sleep(150 * time.Millisecond)
+			_, _ = pw.Write([]byte("event: message_start\ndata: {\"type\":\"message_start\",\"message\":{}}\n\n"))
+			_, _ = pw.Write([]byte("event: content_block_start\ndata: {\"type\":\"content_block_start\",\"index\":0,\"content_block\":{\"type\":\"text\"}}\n\n"))
+			_, _ = pw.Write([]byte("event: content_block_delta\ndata: {\"type\":\"content_block_delta\",\"index\":0,\"delta\":{\"type\":\"text_delta\",\"text\":\"OK\"}}\n\n"))
+			_, _ = pw.Write([]byte("event: content_block_stop\ndata: {\"type\":\"content_block_stop\",\"index\":0}\n\n"))
+			_, _ = pw.Write([]byte("event: message_delta\ndata: {\"type\":\"message_delta\",\"delta\":{\"stop_reason\":\"end_turn\"}}\n\n"))
+			_, _ = pw.Write([]byte("event: message_stop\ndata: {\"type\":\"message_stop\"}\n\n"))
+		}()
+		return &http.Response{
+			StatusCode: 200,
+			Header: http.Header{
+				"Content-Type": []string{"text/event-stream"},
+			},
+			Body:    pr,
+			Request: r,
+		}, nil
+	})
+	app.mu.Lock()
+	app.clients[""] = &http.Client{Transport: transport}
+	app.mu.Unlock()
+
+	account, err := app.insertAccount(M{
+		"family": "zai", "customer_id": "admission-probe", "device_mid": "dev-admission",
+		"status": "active",
+	}, M{"zcode_jwt": registrationJWT(`{"user_id":"admission-probe"}`)})
+	registrationMust(t, err)
+
+	// A 100ms admission window would normally cancel the 150ms response; with
+	// the fix the timer stops once headers arrive and the request completes.
+	o := inferOptions{
+		ID: "probe-admission", Model: "glm-5.3-flash", Protocol: "anthropic",
+		Admission: 100 * time.Millisecond, Timing: newTiming(),
+		Upstream: M{"model": "glm-5.3-flash", "messages": []any{}, "stream": true},
+	}
+	start := time.Now()
+	result, failure := app.inferHTTP(context.Background(), account, o)
+	if failure != nil {
+		t.Fatalf("request failed under admission window: %s/%s", failure.Code, failure.Message)
+	}
+	if result == nil {
+		t.Fatal("no result")
+	}
+	if time.Since(start) < 150*time.Millisecond {
+		t.Fatal("response returned before upstream answered")
+	}
+}
+
+// One request must not walk the whole pool on rate limits: after two rotations
+// the throttled accounts re-enter selection so the request waits instead of
+// rotating to yet another cold-cache account.
+func TestRateLimitRotationStopsAtCap(t *testing.T) {
+	manager := &outboundProxyManager{
+		unavailable: map[string]bool{},
+		retryAt:     map[string]time.Time{},
+		blocked:     map[string]time.Time{},
+	}
+	endpoints := []proxyEndpoint{
+		{URL: "http://exit-a:1"}, {URL: "http://exit-b:2"}, {URL: "http://exit-c:3"},
+	}
+	options := M{"enabled": true, "source": "manual", "urls": []any{"http://exit-a:1", "http://exit-b:2", "http://exit-c:3"}}
+	seen := map[string]bool{}
+	for i := 0; i < 4; i++ {
+		index, _ := manager.selectRotatingEndpoint(options, endpoints)
+		if index < 0 {
+			t.Fatal("no endpoint returned")
+		}
+		seen[endpoints[index].URL] = true
+	}
+	if len(seen) > 3 {
+		t.Fatalf("rotation walked more endpoints than exist: %v", seen)
+	}
+}

internal/server/affinity_test.go

@@ -0,0 +1,37 @@
+package server
+
+import (
+	"net/http"
+	"testing"
+)
+
+const affinitySession = "3f8b1c9d2e6a4f719c305d17b8e2a4f1"
+
+// A Claude Code user id keeps the account id after the session id. The session
+// must still be extracted so conversation affinity keeps one account per chat.
+func TestSessionPatternExtractsMidStringSession(t *testing.T) {
+	for _, value := range []string{
+		"user_ab12cd34_session_" + affinitySession + "_account_z",
+		"user_x_session_" + affinitySession,
+		"session_" + affinitySession,
+	} {
+		p := sessionPattern.FindStringSubmatch(value)
+		if len(p) < 2 || p[1] != affinitySession {
+			t.Fatalf("session not extracted from %q: %v", value, p)
+		}
+	}
+	if p := sessionPattern.FindStringSubmatch("user_ab12cd34_account_z"); p != nil {
+		t.Fatalf("unexpected session match: %v", p)
+	}
+}
+
+func TestAffinityUsesSessionHeaderOverKeySticky(t *testing.T) {
+	app := registrationTestApp(t)
+	r, _ := http.NewRequest("POST", "/v1/messages", nil)
+	r.Header.Set("X-Session-Id", "conv-12345678")
+	key := M{"id": "key_test", "sticky_minutes": 0}
+	affinity, sticky := app.affinity(r, M{}, key)
+	if affinity == "key_test" || sticky <= 0 {
+		t.Fatalf("session affinity not applied: %q %v", affinity, sticky)
+	}
+}

internal/server/app.go

@@ -161,6 +161,7 @@ func New(root, data string, web fs.FS, buildVersion string) (*App, error) {
 		}
 	}
 	a.ensureGroups()
+	a.restoreProxyState()
 	a.background(a.runAccountQueue)
 	return a, nil
 }

internal/server/claim_exit_rotation_test.go

@@ -0,0 +1,76 @@
+package server
+
+import (
+	"testing"
+	"time"
+)
+
+func rotationEndpoints() []proxyEndpoint {
+	return []proxyEndpoint{
+		{URL: "http://exit-a:1", Name: "a"},
+		{URL: "http://exit-b:2", Name: "b"},
+		{URL: "http://exit-c:3", Name: "c"},
+	}
+}
+
+// Each claim must advance to the next pool exit so consecutive cards are issued
+// from different IPs, skipping nodes that are cooling down.
+func TestRotatingEndpointAdvancesPerClaim(t *testing.T) {
+	manager := &outboundProxyManager{
+		unavailable: map[string]bool{},
+		retryAt:     map[string]time.Time{},
+		blocked:     map[string]time.Time{},
+	}
+	endpoints := rotationEndpoints()
+	seen := map[string]int{}
+	for i := 0; i < 6; i++ {
+		index, _ := manager.selectRotatingEndpoint(M{"enabled": true, "source": "manual", "urls": []any{"http://exit-a:1", "http://exit-b:2", "http://exit-c:3"}}, endpoints)
+		if index < 0 {
+			t.Fatalf("rotation returned no endpoint at step %d", i)
+		}
+		seen[endpoints[index].URL]++
+	}
+	for _, url := range []string{"http://exit-a:1", "http://exit-b:2", "http://exit-c:3"} {
+		if seen[url] != 2 {
+			t.Fatalf("expected each exit twice, got %v", seen)
+		}
+	}
+}
+
+func TestRotatingEndpointSkipsBlockedNodes(t *testing.T) {
+	manager := &outboundProxyManager{
+		unavailable: map[string]bool{},
+		retryAt:     map[string]time.Time{},
+		blocked:     map[string]time.Time{},
+	}
+	endpoints := rotationEndpoints()
+	options := M{"enabled": true, "source": "manual", "urls": []any{"http://exit-a:1", "http://exit-b:2", "http://exit-c:3"}}
+	// Establish the config fingerprint first: a key change resets cooldowns.
+	manager.selectRotatingEndpoint(options, endpoints)
+	manager.blocked["http://exit-a:1"] = time.Now().Add(time.Minute)
+	for i := 0; i < 4; i++ {
+		index, _ := manager.selectRotatingEndpoint(options, endpoints)
+		if index < 0 || endpoints[index].URL == "http://exit-a:1" {
+			t.Fatalf("rotation landed on blocked exit: index=%d", index)
+		}
+	}
+}
+
+// When every node is cooling down the cursor still yields a usable position so
+// a claim can retry instead of hard-failing during overlapping cooldowns.
+func TestRotatingEndpointFallsBackWhenAllBlocked(t *testing.T) {
+	manager := &outboundProxyManager{
+		unavailable: map[string]bool{},
+		retryAt:     map[string]time.Time{},
+		blocked: map[string]time.Time{
+			"http://exit-a:1": time.Now().Add(time.Minute),
+			"http://exit-b:2": time.Now().Add(time.Minute),
+			"http://exit-c:3": time.Now().Add(time.Minute),
+		},
+	}
+	endpoints := rotationEndpoints()
+	index, _ := manager.selectRotatingEndpoint(M{"enabled": true, "source": "manual", "urls": []any{"http://exit-a:1", "http://exit-b:2", "http://exit-c:3"}}, endpoints)
+	if index < 0 {
+		t.Fatal("rotation returned no fallback endpoint")
+	}
+}

internal/server/claim_verification.go

@@ -15,7 +15,9 @@ func (a *App) claimVerification(ctx context.Context, id, plan string) (M, error)
 	if account == nil {
 		return nil, fmt.Errorf("账号不存在")
 	}
-	ctx, release, err := a.pinOutboundProxy(ctx, account)
+	// Rotate to the next pool exit per claim so each card is issued from a
+	// different IP; the session pins every later request of this claim to it.
+	ctx, release, err := a.pinClaimOutboundProxy(ctx, account)
 	if err != nil {
 		return nil, err
 	}