refactor(codex): align continuity helpers with review feedback

Align websocket continuity resolution with the HTTP Codex path, make auth-affinity principal keys use a stable string representation, and extract small helpers that remove duplicated continuity and affinity logic without changing the validated cache-hit behavior.
This commit is contained in:
VooDisss
2026-03-27 18:11:57 +02:00
parent 511b8a992e
commit 62b17f40a1
5 changed files with 86 additions and 98 deletions

View File

@@ -216,13 +216,31 @@ func requestExecutionMetadata(ctx context.Context) map[string]any {
} else if ctx != nil {
if ginCtx, ok := ctx.Value("gin").(*gin.Context); ok && ginCtx != nil {
if apiKey, exists := ginCtx.Get("apiKey"); exists && apiKey != nil {
meta[authAffinityMetadataKey] = fmt.Sprintf("principal:%v", apiKey)
if principal := stablePrincipalMetadataKey(apiKey); principal != "" {
meta[authAffinityMetadataKey] = principal
}
}
}
}
return meta
}
func stablePrincipalMetadataKey(raw any) string {
var keyStr string
switch v := raw.(type) {
case string:
keyStr = v
case fmt.Stringer:
keyStr = v.String()
default:
keyStr = fmt.Sprintf("%v", raw)
}
if trimmed := strings.TrimSpace(keyStr); trimmed != "" {
return "principal:" + trimmed
}
return ""
}
func pinnedAuthIDFromContext(ctx context.Context) string {
if ctx == nil {
return ""

View File

@@ -2271,6 +2271,20 @@ func (m *Manager) AuthAffinity(key string) string {
return strings.TrimSpace(m.affinity[key])
}
func (m *Manager) applyAuthAffinity(opts *cliproxyexecutor.Options) {
if m == nil || opts == nil || pinnedAuthIDFromMetadata(opts.Metadata) != "" {
return
}
if affinityKey := authAffinityKeyFromMetadata(opts.Metadata); affinityKey != "" {
if affinityAuthID := m.AuthAffinity(affinityKey); affinityAuthID != "" {
if opts.Metadata == nil {
opts.Metadata = make(map[string]any)
}
opts.Metadata[cliproxyexecutor.PinnedAuthMetadataKey] = affinityAuthID
}
}
}
func (m *Manager) SetAuthAffinity(key, authID string) {
key = strings.TrimSpace(key)
authID = strings.TrimSpace(authID)
@@ -2378,18 +2392,7 @@ func (m *Manager) pickNextLegacy(ctx context.Context, provider, model string, op
}
func (m *Manager) pickNext(ctx context.Context, provider, model string, opts cliproxyexecutor.Options, tried map[string]struct{}) (*Auth, ProviderExecutor, error) {
if pinnedAuthID := pinnedAuthIDFromMetadata(opts.Metadata); pinnedAuthID == "" {
if affinityKey := authAffinityKeyFromMetadata(opts.Metadata); affinityKey != "" {
if affinityAuthID := m.AuthAffinity(affinityKey); affinityAuthID != "" {
meta := opts.Metadata
if meta == nil {
meta = make(map[string]any)
opts.Metadata = meta
}
meta[cliproxyexecutor.PinnedAuthMetadataKey] = affinityAuthID
}
}
}
m.applyAuthAffinity(&opts)
if !m.useSchedulerFastPath() {
return m.pickNextLegacy(ctx, provider, model, opts, tried)
}
@@ -2504,18 +2507,7 @@ func (m *Manager) pickNextMixedLegacy(ctx context.Context, providers []string, m
}
func (m *Manager) pickNextMixed(ctx context.Context, providers []string, model string, opts cliproxyexecutor.Options, tried map[string]struct{}) (*Auth, ProviderExecutor, string, error) {
if pinnedAuthID := pinnedAuthIDFromMetadata(opts.Metadata); pinnedAuthID == "" {
if affinityKey := authAffinityKeyFromMetadata(opts.Metadata); affinityKey != "" {
if affinityAuthID := m.AuthAffinity(affinityKey); affinityAuthID != "" {
meta := opts.Metadata
if meta == nil {
meta = make(map[string]any)
opts.Metadata = meta
}
meta[cliproxyexecutor.PinnedAuthMetadataKey] = affinityAuthID
}
}
}
m.applyAuthAffinity(&opts)
if !m.useSchedulerFastPath() {
return m.pickNextMixedLegacy(ctx, providers, model, opts, tried)
}