Skip to content

Commit 15ac7fb

Browse files
luispaterluispater
committed
refactor(auth): simplify home auth session management and remove ref counting
- Consolidated `homeRuntimeAuths` to store a map of session-scoped auth maps, replacing `homeRuntimeAuthSessions` and `homeRuntimeAuthRefs`. - Adjusted session cleanup logic to directly remove session-scoped auths without reference counting. - Added `GetExecutionSessionAuthByID` to retrieve auths scoped to a specific execution session. - Updated tests to reflect the new session-scoped caching behavior.
1 parent 8300ee8 commit 15ac7fb

3 files changed

Lines changed: 121 additions & 72 deletions

File tree

‎sdk/api/handlers/openai/openai_responses_websocket.go‎

Lines changed: 14 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -104,6 +104,15 @@ func (h *OpenAIResponsesAPIHandler) ResponsesWebsocket(c *gin.Context) {
104104
var lastRequest []byte
105105
lastResponseOutput := []byte("[]")
106106
pinnedAuthID := ""
107+
sessionAuthByID := func(authID string) (*coreauth.Auth, bool) {
108+
if h == nil || h.AuthManager == nil {
109+
return nil, false
110+
}
111+
if auth, ok := h.AuthManager.GetExecutionSessionAuthByID(passthroughSessionID, authID); ok {
112+
return auth, true
113+
}
114+
return h.AuthManager.GetByID(authID)
115+
}
107116
forceTranscriptReplayNextRequest := false
108117

109118
for {
@@ -130,8 +139,8 @@ func (h *OpenAIResponsesAPIHandler) ResponsesWebsocket(c *gin.Context) {
130139
appendWebsocketTimelineEvent(&wsTimelineLog, "request", payload, time.Now())
131140

132141
allowIncrementalInputWithPreviousResponseID := false
133-
if pinnedAuthID != "" && h != nil && h.AuthManager != nil {
134-
if pinnedAuth, ok := h.AuthManager.GetByID(pinnedAuthID); ok && pinnedAuth != nil {
142+
if pinnedAuthID != "" {
143+
if pinnedAuth, ok := sessionAuthByID(pinnedAuthID); ok && pinnedAuth != nil {
135144
allowIncrementalInputWithPreviousResponseID = websocketUpstreamSupportsIncrementalInput(pinnedAuth.Attributes, pinnedAuth.Metadata)
136145
}
137146
} else {
@@ -146,8 +155,8 @@ func (h *OpenAIResponsesAPIHandler) ResponsesWebsocket(c *gin.Context) {
146155
}
147156

148157
allowCompactionReplayBypass := false
149-
if pinnedAuthID != "" && h != nil && h.AuthManager != nil {
150-
if pinnedAuth, ok := h.AuthManager.GetByID(pinnedAuthID); ok && pinnedAuth != nil {
158+
if pinnedAuthID != "" {
159+
if pinnedAuth, ok := sessionAuthByID(pinnedAuthID); ok && pinnedAuth != nil {
151160
allowCompactionReplayBypass = responsesWebsocketAuthSupportsCompactionReplay(pinnedAuth)
152161
}
153162
} else {
@@ -228,7 +237,7 @@ func (h *OpenAIResponsesAPIHandler) ResponsesWebsocket(c *gin.Context) {
228237
if authID == "" || h == nil || h.AuthManager == nil {
229238
return
230239
}
231-
selectedAuth, ok := h.AuthManager.GetByID(authID)
240+
selectedAuth, ok := sessionAuthByID(authID)
232241
if !ok || selectedAuth == nil {
233242
return
234243
}

‎sdk/cliproxy/auth/conductor.go‎

Lines changed: 36 additions & 56 deletions
Original file line numberDiff line numberDiff line change
@@ -153,9 +153,7 @@ type Manager struct {
153153
scheduler *authScheduler
154154
// homeRuntimeAuths caches auths returned by Home so websocket sessions can
155155
// reuse an established upstream credential without dispatching every turn.
156-
homeRuntimeAuths map[string]*Auth
157-
homeRuntimeAuthSessions map[string]map[string]struct{}
158-
homeRuntimeAuthRefs map[string]int
156+
homeRuntimeAuths map[string]map[string]*Auth
159157
// providerOffsets tracks per-model provider rotation state for multi-provider routing.
160158
providerOffsets map[string]int
161159

@@ -195,16 +193,14 @@ func NewManager(store Store, selector Selector, hook Hook) *Manager {
195193
hook = NoopHook{}
196194
}
197195
manager := &Manager{
198-
store: store,
199-
executors: make(map[string]ProviderExecutor),
200-
selector: selector,
201-
hook: hook,
202-
auths: make(map[string]*Auth),
203-
homeRuntimeAuths: make(map[string]*Auth),
204-
homeRuntimeAuthSessions: make(map[string]map[string]struct{}),
205-
homeRuntimeAuthRefs: make(map[string]int),
206-
providerOffsets: make(map[string]int),
207-
modelPoolOffsets: make(map[string]int),
196+
store: store,
197+
executors: make(map[string]ProviderExecutor),
198+
selector: selector,
199+
hook: hook,
200+
auths: make(map[string]*Auth),
201+
homeRuntimeAuths: make(map[string]map[string]*Auth),
202+
providerOffsets: make(map[string]int),
203+
modelPoolOffsets: make(map[string]int),
208204
}
209205
// atomic.Value requires non-nil initial value.
210206
manager.runtimeConfig.Store(&internalconfig.Config{})
@@ -2724,10 +2720,24 @@ func (m *Manager) GetByID(id string) (*Auth, bool) {
27242720
defer m.mu.RUnlock()
27252721
auth, ok := m.auths[id]
27262722
if !ok {
2727-
auth, ok = m.homeRuntimeAuths[id]
2728-
if !ok {
2729-
return nil, false
2730-
}
2723+
return nil, false
2724+
}
2725+
return auth.Clone(), true
2726+
}
2727+
2728+
// GetExecutionSessionAuthByID retrieves a Home runtime auth scoped to an execution session.
2729+
func (m *Manager) GetExecutionSessionAuthByID(sessionID string, authID string) (*Auth, bool) {
2730+
sessionID = strings.TrimSpace(sessionID)
2731+
authID = strings.TrimSpace(authID)
2732+
if m == nil || sessionID == "" || authID == "" {
2733+
return nil, false
2734+
}
2735+
m.mu.RLock()
2736+
defer m.mu.RUnlock()
2737+
sessionAuths := m.homeRuntimeAuths[sessionID]
2738+
auth := sessionAuths[authID]
2739+
if auth == nil {
2740+
return nil, false
27312741
}
27322742
return auth.Clone(), true
27332743
}
@@ -3218,31 +3228,15 @@ func (m *Manager) clearHomeRuntimeAuthsLocked() {
32183228
if m == nil {
32193229
return
32203230
}
3221-
m.homeRuntimeAuths = make(map[string]*Auth)
3222-
m.homeRuntimeAuthSessions = make(map[string]map[string]struct{})
3223-
m.homeRuntimeAuthRefs = make(map[string]int)
3231+
m.homeRuntimeAuths = make(map[string]map[string]*Auth)
32243232
}
32253233

32263234
func (m *Manager) clearHomeRuntimeAuthsForSessionLocked(sessionID string) {
32273235
sessionID = strings.TrimSpace(sessionID)
32283236
if m == nil || sessionID == "" {
32293237
return
32303238
}
3231-
authIDs := m.homeRuntimeAuthSessions[sessionID]
3232-
if len(authIDs) == 0 {
3233-
delete(m.homeRuntimeAuthSessions, sessionID)
3234-
return
3235-
}
3236-
for authID := range authIDs {
3237-
refCount := m.homeRuntimeAuthRefs[authID]
3238-
if refCount <= 1 {
3239-
delete(m.homeRuntimeAuthRefs, authID)
3240-
delete(m.homeRuntimeAuths, authID)
3241-
continue
3242-
}
3243-
m.homeRuntimeAuthRefs[authID] = refCount - 1
3244-
}
3245-
delete(m.homeRuntimeAuthSessions, sessionID)
3239+
delete(m.homeRuntimeAuths, sessionID)
32463240
}
32473241

32483242
func (m *Manager) rememberHomeRuntimeAuth(sessionID string, auth *Auth) {
@@ -3256,24 +3250,14 @@ func (m *Manager) rememberHomeRuntimeAuth(sessionID string, auth *Auth) {
32563250
}
32573251
m.mu.Lock()
32583252
if m.homeRuntimeAuths == nil {
3259-
m.homeRuntimeAuths = make(map[string]*Auth)
3260-
}
3261-
if m.homeRuntimeAuthSessions == nil {
3262-
m.homeRuntimeAuthSessions = make(map[string]map[string]struct{})
3253+
m.homeRuntimeAuths = make(map[string]map[string]*Auth)
32633254
}
3264-
if m.homeRuntimeAuthRefs == nil {
3265-
m.homeRuntimeAuthRefs = make(map[string]int)
3266-
}
3267-
m.homeRuntimeAuths[authID] = auth.Clone()
3268-
sessionAuths := m.homeRuntimeAuthSessions[sessionID]
3255+
sessionAuths := m.homeRuntimeAuths[sessionID]
32693256
if sessionAuths == nil {
3270-
sessionAuths = make(map[string]struct{})
3271-
m.homeRuntimeAuthSessions[sessionID] = sessionAuths
3272-
}
3273-
if _, exists := sessionAuths[authID]; !exists {
3274-
sessionAuths[authID] = struct{}{}
3275-
m.homeRuntimeAuthRefs[authID]++
3257+
sessionAuths = make(map[string]*Auth)
3258+
m.homeRuntimeAuths[sessionID] = sessionAuths
32763259
}
3260+
sessionAuths[authID] = auth.Clone()
32773261
m.mu.Unlock()
32783262
}
32793263

@@ -3284,12 +3268,8 @@ func (m *Manager) homeRuntimeAuthByID(sessionID string, authID string) (*Auth, P
32843268
return nil, nil, "", false
32853269
}
32863270
m.mu.RLock()
3287-
sessionAuths := m.homeRuntimeAuthSessions[sessionID]
3288-
if _, ok := sessionAuths[authID]; !ok {
3289-
m.mu.RUnlock()
3290-
return nil, nil, "", false
3291-
}
3292-
auth := m.homeRuntimeAuths[authID]
3271+
sessionAuths := m.homeRuntimeAuths[sessionID]
3272+
auth := sessionAuths[authID]
32933273
m.mu.RUnlock()
32943274
if auth == nil || !authWebsocketsEnabled(auth) {
32953275
return nil, nil, "", false

‎sdk/cliproxy/auth/home_websocket_reuse_test.go‎

Lines changed: 71 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -27,9 +27,9 @@ func TestPickNextViaHomeReusesPinnedWebsocketAuthWithoutHomeDispatch(t *testing.
2727
}
2828
auth.EnsureIndex()
2929
manager.rememberHomeRuntimeAuth("session-1", auth)
30-
cachedAuth, ok := manager.GetByID("home-auth-1")
30+
cachedAuth, ok := manager.GetExecutionSessionAuthByID("session-1", "home-auth-1")
3131
if !ok || cachedAuth == nil || !authWebsocketsEnabled(cachedAuth) {
32-
t.Fatalf("GetByID() did not expose remembered websocket home auth: auth=%#v ok=%v", cachedAuth, ok)
32+
t.Fatalf("GetExecutionSessionAuthByID() did not expose remembered websocket home auth: auth=%#v ok=%v", cachedAuth, ok)
3333
}
3434

3535
ctx := cliproxyexecutor.WithDownstreamWebsocket(context.Background())
@@ -56,6 +56,61 @@ func TestPickNextViaHomeReusesPinnedWebsocketAuthWithoutHomeDispatch(t *testing.
5656
}
5757
}
5858

59+
func TestPickNextViaHomeKeepsSameAuthIDPayloadSessionScoped(t *testing.T) {
60+
manager := NewManager(nil, nil, nil)
61+
manager.SetConfig(&internalconfig.Config{Home: internalconfig.HomeConfig{Enabled: true}})
62+
manager.RegisterExecutor(schedulerTestExecutor{})
63+
64+
manager.rememberHomeRuntimeAuth("session-1", &Auth{
65+
ID: "home-auth-1",
66+
Provider: "test",
67+
Status: StatusActive,
68+
Attributes: map[string]string{
69+
"websockets": "true",
70+
homeUpstreamModelAttributeKey: "upstream-model-a",
71+
},
72+
})
73+
manager.rememberHomeRuntimeAuth("session-2", &Auth{
74+
ID: "home-auth-1",
75+
Provider: "test",
76+
Status: StatusActive,
77+
Attributes: map[string]string{
78+
"websockets": "true",
79+
homeUpstreamModelAttributeKey: "upstream-model-b",
80+
},
81+
})
82+
83+
ctx := cliproxyexecutor.WithDownstreamWebsocket(context.Background())
84+
optsSession1 := cliproxyexecutor.Options{
85+
Metadata: map[string]any{
86+
cliproxyexecutor.ExecutionSessionMetadataKey: "session-1",
87+
cliproxyexecutor.PinnedAuthMetadataKey: "home-auth-1",
88+
},
89+
}
90+
optsSession2 := cliproxyexecutor.Options{
91+
Metadata: map[string]any{
92+
cliproxyexecutor.ExecutionSessionMetadataKey: "session-2",
93+
cliproxyexecutor.PinnedAuthMetadataKey: "home-auth-1",
94+
},
95+
}
96+
97+
gotSession1, _, _, errSession1 := manager.pickNextViaHome(ctx, "gpt-5.4", optsSession1, nil)
98+
if errSession1 != nil {
99+
t.Fatalf("pickNextViaHome(session-1) error = %v", errSession1)
100+
}
101+
if got := gotSession1.Attributes[homeUpstreamModelAttributeKey]; got != "upstream-model-a" {
102+
t.Fatalf("pickNextViaHome(session-1) upstream model = %q, want upstream-model-a", got)
103+
}
104+
105+
gotSession2, _, _, errSession2 := manager.pickNextViaHome(ctx, "gpt-5.4", optsSession2, nil)
106+
if errSession2 != nil {
107+
t.Fatalf("pickNextViaHome(session-2) error = %v", errSession2)
108+
}
109+
if got := gotSession2.Attributes[homeUpstreamModelAttributeKey]; got != "upstream-model-b" {
110+
t.Fatalf("pickNextViaHome(session-2) upstream model = %q, want upstream-model-b", got)
111+
}
112+
}
113+
59114
func TestPickNextViaHomeDoesNotReuseTriedPinnedWebsocketAuth(t *testing.T) {
60115
manager := NewManager(nil, nil, nil)
61116
manager.SetConfig(&internalconfig.Config{Home: internalconfig.HomeConfig{Enabled: true}})
@@ -135,10 +190,12 @@ func TestPickNextViaHomeDoesNotReusePinnedNonWebsocketAuth(t *testing.T) {
135190
manager.RegisterExecutor(schedulerTestExecutor{})
136191

137192
manager.mu.Lock()
138-
manager.homeRuntimeAuths["home-auth-1"] = &Auth{
139-
ID: "home-auth-1",
140-
Provider: "test",
141-
Status: StatusActive,
193+
manager.homeRuntimeAuths["session-1"] = map[string]*Auth{
194+
"home-auth-1": &Auth{
195+
ID: "home-auth-1",
196+
Provider: "test",
197+
Status: StatusActive,
198+
},
142199
}
143200
manager.mu.Unlock()
144201

@@ -175,12 +232,12 @@ func TestHomeRuntimeAuthsClearWhenHomeDisabled(t *testing.T) {
175232
},
176233
})
177234

178-
if _, ok := manager.GetByID("home-auth-1"); !ok {
235+
if _, ok := manager.GetExecutionSessionAuthByID("session-1", "home-auth-1"); !ok {
179236
t.Fatal("expected remembered home auth before disabling home")
180237
}
181238

182239
manager.SetConfig(&internalconfig.Config{})
183-
if _, ok := manager.GetByID("home-auth-1"); ok {
240+
if _, ok := manager.GetExecutionSessionAuthByID("session-1", "home-auth-1"); ok {
184241
t.Fatal("remembered home auth was not cleared when home was disabled")
185242
}
186243
}
@@ -199,12 +256,15 @@ func TestCloseExecutionSessionClearsHomeRuntimeAuthForSession(t *testing.T) {
199256
manager.rememberHomeRuntimeAuth("session-2", auth)
200257

201258
manager.CloseExecutionSession("session-1")
202-
if _, ok := manager.GetByID("home-auth-1"); !ok {
203-
t.Fatal("shared home auth was cleared while another session still referenced it")
259+
if _, ok := manager.GetExecutionSessionAuthByID("session-1", "home-auth-1"); ok {
260+
t.Fatal("home auth for closed session was not cleared")
261+
}
262+
if _, ok := manager.GetExecutionSessionAuthByID("session-2", "home-auth-1"); !ok {
263+
t.Fatal("home auth for another session was cleared")
204264
}
205265

206266
manager.CloseExecutionSession("session-2")
207-
if _, ok := manager.GetByID("home-auth-1"); ok {
267+
if _, ok := manager.GetExecutionSessionAuthByID("session-2", "home-auth-1"); ok {
208268
t.Fatal("home auth was not cleared when its last session closed")
209269
}
210270
}

0 commit comments

Comments
 (0)