fix: recover from unknown opaque state after affinity loss

Author: Futureppo

README.md

@@ -316,6 +316,8 @@ X-Grok-Session-ID: conversation-123
 
 `previous_response_id`、显式 session ID 和状态签名属于 hard affinity；`GROK_AUTHS_DIR` 下的 `.grokcli2api-affinity.json` 只保存加入 tenant namespace 后的哈希 binding、账号、模型、backend、上游 session、下一 turn 与过期时间。`prompt_cache_key`、用户字段等 soft affinity 与 `store:false` 工具回放仍只保存在内存中。所有映射都受 TTL 与容量上限控制；本地 API Key 不直接选择账号，但会隔离不同 tenant 的 affinity、Response ownership、状态签名与工具回放。客户端 IP 不参与亲和。
 
+容器重建或 affinity 状态丢失后，请求中的 Responses `reasoning.encrypted_content` 或 Anthropic thinking signature 可能在当前 tenant 中没有 binding。服务会删除这部分未知 opaque state，并以新上游 session 继续其余对话，而不会把密文转发给随机账号；已经存在但互相指向不同账号/session 的有效 binding 仍会被拒绝。未知 `previous_response_id` 继续返回 404。
+
 ## 配置说明
 
 程序会从当前工作目录的 `.env` 文件中加载尚未设置的环境变量。完整模板与高级客户端标识选项见 [`.env.example`](.env.example)。

README_EN.md

@@ -316,6 +316,8 @@ The service recognizes affinity inputs in this priority order:
 
 `previous_response_id`, an explicit session ID, and state signatures create hard affinity. The `.grokcli2api-affinity.json` file under `GROK_AUTHS_DIR` contains only the tenant-namespaced hashed binding, account, model, backend, upstream session, next turn, and expiry. Soft affinity from cache/user fields and `store:false` tool replay remain in memory only. Every mapping is bounded by the configured TTL and capacity. A local API key does not directly choose an account, but it isolates affinity, Response ownership, state signatures, and replay data between tenants; client IP addresses are not used.
 
+After a container rebuild or affinity-state loss, Responses `reasoning.encrypted_content` or an Anthropic thinking signature may have no binding in the current tenant. The service removes only that unknown opaque state and continues the remaining conversation with a fresh upstream session instead of forwarding ciphertext through a random account. Valid bindings that point to different accounts or sessions are still rejected. An unknown `previous_response_id` continues to return 404.
+
 ## Configuration
 
 The service loads environment variables that are not already set from a `.env` file in the current working directory. See [`.env.example`](.env.example) for the complete template and advanced client-identity options.

internal/inference/plan.go

@@ -151,6 +151,79 @@ func (p *RequestPlan) StateHandles() []StateHandle {
 	return append([]StateHandle(nil), p.stateHandles...)
 }
 
+// WithoutOpaqueState returns an immutable copy of the plan with the selected
+// opaque continuation handles removed. A missing local ownership binding can
+// happen after affinity state is lost during an upgrade or container rebuild;
+// the ciphertext must not be forwarded through a newly selected account, but
+// the remaining public conversation can safely start a fresh upstream session.
+func (p *RequestPlan) WithoutOpaqueState(handles []StateHandle) *RequestPlan {
+	if p == nil || len(handles) == 0 {
+		return p
+	}
+	dropped := make(map[string]struct{}, len(handles))
+	for _, handle := range handles {
+		if handle.Kind != StateOpaqueToken {
+			continue
+		}
+		if value := strings.TrimSpace(handle.Value); value != "" {
+			dropped[value] = struct{}{}
+		}
+	}
+	if len(dropped) == 0 {
+		return p
+	}
+
+	body := cloneMap(p.body)
+	switch p.protocol {
+	case ProtocolResponses:
+		input, _ := body["input"].([]any)
+		clean := make([]any, 0, len(input))
+		for _, rawItem := range input {
+			item, _ := rawItem.(map[string]any)
+			if strings.EqualFold(trimmedString(item["type"]), "reasoning") {
+				if encrypted, ok := item["encrypted_content"].(string); ok {
+					if _, remove := dropped[strings.TrimSpace(encrypted)]; remove {
+						continue
+					}
+				}
+			}
+			clean = append(clean, rawItem)
+		}
+		if input != nil {
+			body["input"] = clean
+		}
+	case ProtocolMessages:
+		messages, _ := body["messages"].([]any)
+		for _, rawMessage := range messages {
+			message, _ := rawMessage.(map[string]any)
+			parts, _ := message["content"].([]any)
+			if parts == nil {
+				continue
+			}
+			clean := make([]any, 0, len(parts))
+			for _, rawPart := range parts {
+				part, _ := rawPart.(map[string]any)
+				kind := trimmedString(part["type"])
+				if strings.EqualFold(kind, "thinking") || strings.EqualFold(kind, "redacted_thinking") {
+					if signature, ok := part["signature"].(string); ok {
+						if _, remove := dropped[strings.TrimSpace(signature)]; remove {
+							continue
+						}
+					}
+				}
+				clean = append(clean, rawPart)
+			}
+			message["content"] = clean
+		}
+	}
+
+	copy := *p
+	copy.body = body
+	copy.stateHandles = requestStateHandles(copy.protocol, body)
+	copy.hasState = len(copy.stateHandles) > 0
+	return &copy
+}
+
 // NewRequestPlan validates only public required structure and fields that are
 // intentionally retained. Unsupported fields and content are left for the
 // account-specific renderer to discard silently.

internal/inference/plan_test.go

@@ -348,6 +348,58 @@ func TestEncryptedResponsesReasoningIsStatefulOnlyOnNativeBackend(t *testing.T)
 	}
 }
 
+func TestWithoutOpaqueStateReturnsSanitizedImmutablePlan(t *testing.T) {
+	t.Run("responses", func(t *testing.T) {
+		plan := mustPlan(t, ProtocolResponses, map[string]any{
+			"model": "m", "input": []any{
+				map[string]any{"type": "reasoning", "encrypted_content": "lost-state"},
+				map[string]any{"type": "reasoning", "encrypted_content": "owned-state"},
+				map[string]any{"type": "message", "role": "user", "content": "continue"},
+			},
+		})
+		clean := plan.WithoutOpaqueState([]StateHandle{{Kind: StateOpaqueToken, Value: "lost-state"}})
+		if len(plan.StateHandles()) != 2 || len(clean.StateHandles()) != 1 {
+			t.Fatalf("state handles original=%#v clean=%#v", plan.StateHandles(), clean.StateHandles())
+		}
+		original, err := plan.Render(modelcatalog.ModelDescriptor{Backend: modelcatalog.BackendResponses})
+		if err != nil || !contains(outputString(original.Body, false), "lost-state") {
+			t.Fatalf("original plan was mutated: body=%#v err=%v", original, err)
+		}
+		attempt, err := clean.Render(modelcatalog.ModelDescriptor{Backend: modelcatalog.BackendResponses})
+		if err != nil {
+			t.Fatal(err)
+		}
+		encoded := outputString(attempt.Body, false)
+		if contains(encoded, "lost-state") || !contains(encoded, "owned-state") || !attempt.PreservesState {
+			t.Fatalf("sanitized attempt=%#v encoded=%s", attempt, encoded)
+		}
+	})
+
+	t.Run("messages", func(t *testing.T) {
+		plan := mustPlan(t, ProtocolMessages, map[string]any{
+			"model": "m", "max_tokens": float64(32), "messages": []any{
+				map[string]any{"role": "assistant", "content": []any{
+					map[string]any{"type": "redacted_thinking", "data": "opaque", "signature": "lost-signature"},
+					map[string]any{"type": "thinking", "thinking": "keep", "signature": "owned-signature"},
+				}},
+				map[string]any{"role": "user", "content": "continue"},
+			},
+		})
+		clean := plan.WithoutOpaqueState([]StateHandle{{Kind: StateOpaqueToken, Value: "lost-signature"}})
+		if len(plan.StateHandles()) != 2 || len(clean.StateHandles()) != 1 {
+			t.Fatalf("state handles original=%#v clean=%#v", plan.StateHandles(), clean.StateHandles())
+		}
+		attempt, err := clean.Render(modelcatalog.ModelDescriptor{Backend: modelcatalog.BackendMessages})
+		if err != nil {
+			t.Fatal(err)
+		}
+		encoded := outputString(attempt.Body, false)
+		if contains(encoded, "lost-signature") || !contains(encoded, "owned-signature") || !attempt.PreservesState {
+			t.Fatalf("sanitized attempt=%#v encoded=%s", attempt, encoded)
+		}
+	})
+}
+
 func TestCleaningErrorsOnlyWhenNoMinimumInputRemains(t *testing.T) {
 	plan := mustPlan(t, ProtocolResponses, map[string]any{
 		"model": "m", "input": []any{

internal/server/continuity_e2e_test.go

@@ -16,15 +16,28 @@ import (
 	"github.com/Futureppo/grokcli2api-go/internal/modelcatalog"
 )
 
-func TestOpaqueStateOwnershipFailsClosedBeforeUpstream(t *testing.T) {
+func TestOpaqueStateOwnershipDropsUnknownBeforeUpstream(t *testing.T) {
 	var inferenceCalls atomic.Int32
+	var mu sync.Mutex
+	var inferenceBodies []map[string]any
 	upstream := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
 		if r.Method == http.MethodGet {
 			writeTestCatalog(w, modelcatalog.BackendResponses)
 			return
 		}
 		inferenceCalls.Add(1)
-		writeJSON(w, http.StatusInternalServerError, map[string]any{"error": "unexpected inference"})
+		var body map[string]any
+		if err := json.NewDecoder(r.Body).Decode(&body); err != nil {
+			t.Errorf("decode inference: %v", err)
+			return
+		}
+		mu.Lock()
+		inferenceBodies = append(inferenceBodies, body)
+		mu.Unlock()
+		writeJSON(w, http.StatusOK, map[string]any{
+			"id": "resp-ok", "object": "response", "status": "completed", "output": []any{},
+			"usage": map[string]any{"input_tokens": 1, "output_tokens": 1, "total_tokens": 2},
+		})
 	}))
 	defer upstream.Close()
 
@@ -36,54 +49,76 @@ func TestOpaqueStateOwnershipFailsClosedBeforeUpstream(t *testing.T) {
 	tenantA := s.pool.TenantID("tenant-a-key")
 	tenantB := s.pool.TenantID("tenant-b-key")
 
-	// A Messages thinking signature owned by tenant A must not be forwarded by
-	// tenant B when the selected account preserves it on a Responses backend.
+	// A Messages thinking signature owned by tenant A is unknown to tenant B.
+	// Drop it and start fresh instead of forwarding it through another account.
 	s.continuity.BindStateToken(tenantA, "foreign-signature", accounts[0], "grok-4", modelcatalog.BackendResponses, "session-a", 1)
-	assertStateRejected(t, s, "tenant-b-key", "/v1/messages", `{
+	assertStateStatus(t, s, "tenant-b-key", "/v1/messages", `{
 		"model":"grok-4","max_tokens":8,
 		"messages":[
 			{"role":"assistant","content":[{"type":"thinking","thinking":"summary","signature":"foreign-signature"}]},
 			{"role":"user","content":"continue"}
 		]
-	}`, "", http.StatusNotFound)
+	}`, "", http.StatusOK)
 
-	// Checking every retained item prevents an owned-first/foreign-second
-	// encrypted reasoning sequence from bypassing tenant isolation.
+	// Keep the locally owned state that selects the route while dropping an
+	// additional token which is unknown in this tenant.
 	s.continuity.BindStateToken(tenantB, "owned-ciphertext", accounts[0], "grok-4", modelcatalog.BackendResponses, "session-b", 2)
 	s.continuity.BindStateToken(tenantA, "foreign-ciphertext", accounts[0], "grok-4", modelcatalog.BackendResponses, "session-a", 2)
-	assertStateRejected(t, s, "tenant-b-key", "/v1/responses", `{
+	assertStateStatus(t, s, "tenant-b-key", "/v1/responses", `{
 		"model":"grok-4","input":[
 			{"type":"reasoning","encrypted_content":"owned-ciphertext"},
 			{"type":"reasoning","encrypted_content":"foreign-ciphertext"},
 			{"type":"message","role":"user","content":[{"type":"input_text","text":"continue"}]}
 		]
-	}`, "", http.StatusNotFound)
+	}`, "", http.StatusOK)
 
-	// An explicit client session cannot hide a stolen opaque token behind its
-	// higher affinity priority.
+	// An explicit client session remains authoritative while the unrelated
+	// unknown token is removed from its request.
 	sessionAffinity := auth.Affinity{Tenant: tenantB, Key: "session:client-session", Mode: auth.AffinityHard}
 	s.continuity.Bind(continuityKey(tenantB, sessionAffinity, "grok-4"), sessionAffinity, accounts[0], "grok-4", modelcatalog.BackendResponses, "session-b", 3)
-	assertStateRejected(t, s, "tenant-b-key", "/v1/responses", `{
+	assertStateStatus(t, s, "tenant-b-key", "/v1/responses", `{
 		"model":"grok-4","input":[
 			{"type":"reasoning","encrypted_content":"foreign-ciphertext"},
 			{"type":"message","role":"user","content":"continue"}
 		]
-	}`, "client-session", http.StatusNotFound)
+	}`, "client-session", http.StatusOK)
 
 	// Two valid handles for the same tenant still cannot be combined when they
 	// belong to different upstream accounts/sessions.
 	s.continuity.BindStateToken(tenantB, "route-one", accounts[0], "grok-4", modelcatalog.BackendResponses, "session-one", 4)
 	s.continuity.BindStateToken(tenantB, "route-two", accounts[1], "grok-4", modelcatalog.BackendResponses, "session-two", 4)
-	assertStateRejected(t, s, "tenant-b-key", "/v1/responses", `{
+	assertStateStatus(t, s, "tenant-b-key", "/v1/responses", `{
 		"model":"grok-4","input":[
 			{"type":"reasoning","encrypted_content":"route-one"},
 			{"type":"reasoning","encrypted_content":"route-two"},
 			{"type":"message","role":"user","content":"continue"}
 		]
 	}`, "", http.StatusServiceUnavailable)
 
-	if got := inferenceCalls.Load(); got != 0 {
-		t.Fatalf("unsafe state reached upstream %d times", got)
+	if got := inferenceCalls.Load(); got != 3 {
+		t.Fatalf("inference calls=%d, want 3 safe degraded requests", got)
+	}
+	mu.Lock()
+	defer mu.Unlock()
+	if len(inferenceBodies) != 3 {
+		t.Fatalf("inference bodies=%d, want 3", len(inferenceBodies))
+	}
+	encoded := make([]string, len(inferenceBodies))
+	for index, body := range inferenceBodies {
+		data, err := json.Marshal(body)
+		if err != nil {
+			t.Fatal(err)
+		}
+		encoded[index] = string(data)
+	}
+	if strings.Contains(encoded[0], "foreign-signature") {
+		t.Fatalf("foreign Messages state reached upstream: %s", encoded[0])
+	}
+	if !strings.Contains(encoded[1], "owned-ciphertext") || strings.Contains(encoded[1], "foreign-ciphertext") {
+		t.Fatalf("mixed state was not cleaned safely: %s", encoded[1])
+	}
+	if strings.Contains(encoded[2], "foreign-ciphertext") {
+		t.Fatalf("foreign session state reached upstream: %s", encoded[2])
 	}
 }
 
@@ -169,7 +204,7 @@ func TestDroppedStateStartsFreshWithoutFalseOwnershipError(t *testing.T) {
 	}
 }
 
-func assertStateRejected(t *testing.T, s *Server, apiKey, path, body, session string, wantStatus int) {
+func assertStateStatus(t *testing.T, s *Server, apiKey, path, body, session string, wantStatus int) {
 	t.Helper()
 	recorder := httptest.NewRecorder()
 	request := httptest.NewRequest(http.MethodPost, path, strings.NewReader(body))

internal/server/inference_execution.go

@@ -5,6 +5,7 @@ import (
 	"encoding/json"
 	"errors"
 	"io"
+	"log/slog"
 	"net/http"
 	"sort"
 	"strings"
@@ -27,6 +28,7 @@ func (e *continuityRequestError) Error() string { return e.Message }
 type inferenceExecution struct {
 	tenant          string
 	model           string
+	plan            *inference.RequestPlan
 	affinity        auth.Affinity
 	pinned          string
 	expectedBackend modelcatalog.Backend
@@ -50,7 +52,7 @@ type continuityLookup struct {
 
 func (s *Server) prepareInferenceExecution(r *http.Request, body map[string]any, plan *inference.RequestPlan) (inferenceExecution, error) {
 	tenant := tenantFromContext(r.Context())
-	execution := inferenceExecution{tenant: tenant, model: plan.Model(), affinity: requestSoftAffinity(r, body)}
+	execution := inferenceExecution{tenant: tenant, model: plan.Model(), plan: plan, affinity: requestSoftAffinity(r, body)}
 	execution.identity = grok.RequestIdentity{RequestID: grok.NewID(), SessionID: newLogicalSession()}
 	execution.identity.ConversationID = execution.identity.SessionID
 	turn := uint64(0)
@@ -63,11 +65,21 @@ func (s *Server) prepareInferenceExecution(r *http.Request, body map[string]any,
 	}
 
 	state := make([]continuityLookup, 0, len(plan.StateHandles()))
+	unknownOpaque := make([]inference.StateHandle, 0)
 	for _, handle := range plan.StateHandles() {
 		affinity := stateHandleAffinity(tenant, handle)
 		key, binding, found := s.continuity.Lookup(tenant, affinity, plan.Model())
+		if handle.Kind == inference.StateOpaqueToken && !found {
+			unknownOpaque = append(unknownOpaque, handle)
+			continue
+		}
 		state = append(state, continuityLookup{handle: handle, affinity: affinity, key: key, binding: binding, found: found})
 	}
+	if len(unknownOpaque) > 0 {
+		plan = plan.WithoutOpaqueState(unknownOpaque)
+		execution.plan = plan
+		slog.Warn("unknown opaque upstream state dropped", "model", plan.Model(), "protocol", plan.Protocol(), "count", len(unknownOpaque))
+	}
 
 	// An existing explicit session is authoritative. Only state tokens which
 	// that account's renderer will actually retain need ownership validation;
@@ -326,6 +338,7 @@ func (s *Server) executeNonStreaming(w http.ResponseWriter, r *http.Request, bod
 		s.writePlanError(w, plan.Protocol(), err)
 		return
 	}
+	plan = execution.plan
 	result, err := s.client.DoInference(r.Context(), plan, grok.InferenceOptions{
 		Affinity: execution.affinity, PinnedAccount: execution.pinned, ExpectedBackend: execution.expectedBackend, Identity: execution.identity,
 	})
@@ -365,6 +378,7 @@ func (s *Server) executeStreaming(w http.ResponseWriter, r *http.Request, body m
 		s.writePlanError(w, plan.Protocol(), err)
 		return
 	}
+	plan = execution.plan
 	stream, err := s.client.OpenInference(r.Context(), plan, grok.InferenceOptions{
 		Affinity: execution.affinity, PinnedAccount: execution.pinned, ExpectedBackend: execution.expectedBackend, Identity: execution.identity,
 	})