internal/pluginhost/affinity_callbacks_test.go
Rama main · réplica de lectura del repositorio original.
package pluginhost
import (
"context"
"encoding/json"
"strings"
"sync"
"testing"
coreauth "github.com/router-for-me/CLIProxyAPI/v8/sdk/cliproxy/auth"
cliproxyexecutor "github.com/router-for-me/CLIProxyAPI/v8/sdk/cliproxy/executor"
"github.com/router-for-me/CLIProxyAPI/v8/sdk/pluginabi"
"github.com/router-for-me/CLIProxyAPI/v8/sdk/pluginapi"
)
type mockPluginScheduler struct{}
func (m *mockPluginScheduler) PickAuth(ctx context.Context, req pluginapi.SchedulerPickRequest) (pluginapi.SchedulerPickResponse, bool, error) {
return pluginapi.SchedulerPickResponse{}, false, nil
}
func TestHostAffinityLookupCallback_Contract(t *testing.T) {
host := New()
manager := coreauth.NewManager(nil, nil, nil)
selector := coreauth.NewSessionAffinitySelector(nil)
manager.SetSelector(selector)
host.SetAuthManager(manager)
authA := &coreauth.Auth{
ID: "claude-owner-a@example.com.json",
Provider: "anthropic",
FileName: "claude-owner-a@example.com.json",
Label: "auth-a",
Status: coreauth.StatusActive,
Metadata: map[string]any{
"type": "anthropic",
"api_key": "secret-api-key-a",
"email": "owner-a@example.com",
},
Attributes: map[string]string{
"path": "/etc/secrets/claude-owner-a@example.com.json",
},
}
authA.EnsureIndex()
authB := &coreauth.Auth{
ID: "claude-owner-b@example.com.json",
Provider: "anthropic",
FileName: "claude-owner-b@example.com.json",
Label: "auth-b",
Status: coreauth.StatusActive,
Metadata: map[string]any{
"type": "anthropic",
"api_key": "secret-api-key-b",
"email": "owner-b@example.com",
},
Attributes: map[string]string{
"path": "/etc/secrets/claude-owner-b@example.com.json",
},
}
authB.EnsureIndex()
if _, errReg := manager.Register(context.Background(), authA); errReg != nil {
t.Fatalf("register authA: %v", errReg)
}
if _, errReg := manager.Register(context.Background(), authB); errReg != nil {
t.Fatalf("register authB: %v", errReg)
}
t.Run("absent binding returns unbound", func(t *testing.T) {
reqPayload, _ := json.Marshal(pluginapi.HostAffinityLookupRequest{
Provider: "anthropic",
Model: "claude-3-7-sonnet",
SessionID: "non-existent-session",
})
rawResp, errCall := host.callFromPlugin(context.Background(), pluginabi.MethodHostAffinityLookup, reqPayload)
if errCall != nil {
t.Fatalf("callFromPlugin error = %v", errCall)
}
resp, errDecode := decodeRPCEnvelope[pluginapi.HostAffinityLookupResponse](rawResp)
if errDecode != nil {
t.Fatalf("decode response: %v", errDecode)
}
if resp.Status != pluginapi.HostAffinityStatusUnbound {
t.Fatalf("status = %q, want %q", resp.Status, pluginapi.HostAffinityStatusUnbound)
}
if resp.AuthIndex != "" {
t.Fatalf("unexpected auth in unbound response: %#v", resp)
}
})
t.Run("bound session lookup and payload hygiene without email leak", func(t *testing.T) {
opts := cliproxyexecutor.Options{
Headers: map[string][]string{"X-Claude-Code-Session-Id": {"bound-session-1"}},
Metadata: make(map[string]any),
}
picked, errPick := selector.Pick(context.Background(), "anthropic", "claude-3-7-sonnet", opts, []*coreauth.Auth{authA})
if errPick != nil || picked.ID != authA.ID {
t.Fatalf("pick failed: auth=%v err=%v", picked, errPick)
}
reqPayload, _ := json.Marshal(pluginapi.HostAffinityLookupRequest{
Provider: "anthropic",
Model: "claude-3-7-sonnet",
SessionID: "bound-session-1",
})
rawResp, errCall := host.callFromPlugin(context.Background(), pluginabi.MethodHostAffinityLookup, reqPayload)
if errCall != nil {
t.Fatalf("callFromPlugin error = %v", errCall)
}
// Verify no secret, email, credential ID, or physical file path leaked in raw json
rawStr := string(rawResp)
for _, sensitive := range []string{"secret-api-key-a", "owner-a@example.com", "claude-owner-a@example.com.json", "/etc/secrets"} {
if strings.Contains(rawStr, sensitive) {
t.Fatalf("raw response leaked sensitive token %q: %s", sensitive, rawStr)
}
}
resp, errDecode := decodeRPCEnvelope[pluginapi.HostAffinityLookupResponse](rawResp)
if errDecode != nil {
t.Fatalf("decode response: %v", errDecode)
}
if resp.Status != pluginapi.HostAffinityStatusBound {
t.Fatalf("status = %q, want %q", resp.Status, pluginapi.HostAffinityStatusBound)
}
if resp.AuthIndex != authA.Index {
t.Fatalf("auth_index = %q, want %q", resp.AuthIndex, authA.Index)
}
if resp.Disabled || resp.Unavailable {
t.Fatalf("unexpected disabled/unavailable flags: disabled=%v unavailable=%v", resp.Disabled, resp.Unavailable)
}
if resp.ObservedAt.IsZero() {
t.Fatalf("observed_at is zero")
}
})
t.Run("thinking suffix normalization", func(t *testing.T) {
reqPayload, _ := json.Marshal(pluginapi.HostAffinityLookupRequest{
Provider: "anthropic",
Model: "claude-3-7-sonnet(high)",
SessionID: "bound-session-1",
})
rawResp, errCall := host.callFromPlugin(context.Background(), pluginabi.MethodHostAffinityLookup, reqPayload)
if errCall != nil {
t.Fatalf("callFromPlugin error = %v", errCall)
}
resp, errDecode := decodeRPCEnvelope[pluginapi.HostAffinityLookupResponse](rawResp)
if errDecode != nil {
t.Fatalf("decode response: %v", errDecode)
}
if resp.Status != pluginapi.HostAffinityStatusBound || resp.AuthIndex != authA.Index {
t.Fatalf("normalized thinking model lookup failed: %#v", resp)
}
})
t.Run("long session ID normalization (> 256 bytes)", func(t *testing.T) {
longSessionID := strings.Repeat("s", 252) // prefix "claude:" + 252 chars = 259 chars > 256
opts := cliproxyexecutor.Options{
Headers: map[string][]string{"X-Claude-Code-Session-Id": {longSessionID}},
Metadata: make(map[string]any),
}
picked, errPick := selector.Pick(context.Background(), "anthropic", "claude-3-7-sonnet", opts, []*coreauth.Auth{authA})
if errPick != nil || picked.ID != authA.ID {
t.Fatalf("pick failed: auth=%v err=%v", picked, errPick)
}
reqPayload, _ := json.Marshal(pluginapi.HostAffinityLookupRequest{
Provider: "anthropic",
Model: "claude-3-7-sonnet",
SessionID: longSessionID,
})
rawResp, errCall := host.callFromPlugin(context.Background(), pluginabi.MethodHostAffinityLookup, reqPayload)
if errCall != nil {
t.Fatalf("callFromPlugin error = %v", errCall)
}
resp, errDecode := decodeRPCEnvelope[pluginapi.HostAffinityLookupResponse](rawResp)
if errDecode != nil {
t.Fatalf("decode response: %v", errDecode)
}
if resp.Status != pluginapi.HostAffinityStatusBound || resp.AuthIndex != authA.Index {
t.Fatalf("long session lookup failed: %#v", resp)
}
})
t.Run("rebinding A to B reflects updated binding", func(t *testing.T) {
sessionID := "rebind-session"
opts := cliproxyexecutor.Options{
Headers: map[string][]string{"X-Claude-Code-Session-Id": {sessionID}},
Metadata: make(map[string]any),
}
picked, _ := selector.Pick(context.Background(), "anthropic", "claude-3-7-sonnet", opts, []*coreauth.Auth{authA})
if picked.ID != authA.ID {
t.Fatalf("initial pick = %s, want %s", picked.ID, authA.ID)
}
// Rebind to authB
picked2, _ := selector.Pick(context.Background(), "anthropic", "claude-3-7-sonnet", opts, []*coreauth.Auth{authB})
if picked2.ID != authB.ID {
t.Fatalf("second pick = %s, want %s", picked2.ID, authB.ID)
}
reqPayload, _ := json.Marshal(pluginapi.HostAffinityLookupRequest{
Provider: "anthropic",
Model: "claude-3-7-sonnet",
SessionID: sessionID,
})
rawResp, _ := host.callFromPlugin(context.Background(), pluginabi.MethodHostAffinityLookup, reqPayload)
resp, _ := decodeRPCEnvelope[pluginapi.HostAffinityLookupResponse](rawResp)
if resp.Status != pluginapi.HostAffinityStatusBound || resp.AuthIndex != authB.Index {
t.Fatalf("lookup after rebinding: got index=%q status=%q, want index=%q", resp.AuthIndex, resp.Status, authB.Index)
}
})
t.Run("unknown subagent does not guess parent binding", func(t *testing.T) {
// Parent is bound to authA
parentSession := "parent-main-thread-100"
opts := cliproxyexecutor.Options{
Headers: map[string][]string{"Session-Id": {parentSession}},
Metadata: make(map[string]any),
}
selector.Pick(context.Background(), "anthropic", "claude-3-7-sonnet", opts, []*coreauth.Auth{authA})
// Query unknown child subagent
reqPayload, _ := json.Marshal(pluginapi.HostAffinityLookupRequest{
Provider: "anthropic",
Model: "claude-3-7-sonnet",
SessionID: "codex:" + parentSession + ":subagent-child-unknown",
})
rawResp, _ := host.callFromPlugin(context.Background(), pluginabi.MethodHostAffinityLookup, reqPayload)
resp, _ := decodeRPCEnvelope[pluginapi.HostAffinityLookupResponse](rawResp)
if resp.Status != pluginapi.HostAffinityStatusUnbound {
t.Fatalf("unknown subagent lookup status = %q, want %q", resp.Status, pluginapi.HostAffinityStatusUnbound)
}
})
t.Run("ambiguous unqualified session resolves to ambiguous across namespaces", func(t *testing.T) {
// Bind "ambig-sess" under claude to authA, and under codex to authB
optsClaude := cliproxyexecutor.Options{
Headers: map[string][]string{"X-Claude-Code-Session-Id": {"ambig-sess"}},
Metadata: make(map[string]any),
}
selector.Pick(context.Background(), "anthropic", "claude-3-7-sonnet", optsClaude, []*coreauth.Auth{authA})
optsCodex := cliproxyexecutor.Options{
Headers: map[string][]string{"Session-Id": {"ambig-sess"}},
Metadata: make(map[string]any),
}
selector.Pick(context.Background(), "anthropic", "claude-3-7-sonnet", optsCodex, []*coreauth.Auth{authB})
// Query bare "ambig-sess" without namespace
reqPayload, _ := json.Marshal(pluginapi.HostAffinityLookupRequest{
Provider: "anthropic",
Model: "claude-3-7-sonnet",
SessionID: "ambig-sess",
})
rawResp, _ := host.callFromPlugin(context.Background(), pluginabi.MethodHostAffinityLookup, reqPayload)
resp, _ := decodeRPCEnvelope[pluginapi.HostAffinityLookupResponse](rawResp)
if resp.Status != pluginapi.HostAffinityStatusAmbiguous {
t.Fatalf("ambiguous session status = %q, want %q", resp.Status, pluginapi.HostAffinityStatusAmbiguous)
}
// Qualified queries resolve unambiguously
reqClaude, _ := json.Marshal(pluginapi.HostAffinityLookupRequest{
Provider: "anthropic",
Model: "claude-3-7-sonnet",
SessionID: "claude:ambig-sess",
})
rawClaude, _ := host.callFromPlugin(context.Background(), pluginabi.MethodHostAffinityLookup, reqClaude)
respClaude, _ := decodeRPCEnvelope[pluginapi.HostAffinityLookupResponse](rawClaude)
if respClaude.Status != pluginapi.HostAffinityStatusBound || respClaude.AuthIndex != authA.Index {
t.Fatalf("qualified claude session status = %q index = %q, want bound %q", respClaude.Status, respClaude.AuthIndex, authA.Index)
}
})
t.Run("extended namespaces resolve bare ID (task, clientreq, geminicache, execution)", func(t *testing.T) {
testCases := []struct {
name string
opts cliproxyexecutor.Options
bareID string
qualID string
auth *coreauth.Auth
}{
{
name: "task header",
opts: cliproxyexecutor.Options{
Headers: map[string][]string{"X-Task-ID": {"task-xyz-1"}},
Metadata: make(map[string]any),
},
bareID: "task-xyz-1",
qualID: "task:task-xyz-1",
auth: authA,
},
{
name: "client request id header",
opts: cliproxyexecutor.Options{
Headers: map[string][]string{"X-Client-Request-Id": {"crid-abc-2"}},
Metadata: make(map[string]any),
},
bareID: "crid-abc-2",
qualID: "clientreq:crid-abc-2",
auth: authB,
},
{
name: "gemini cache payload",
opts: cliproxyexecutor.Options{
OriginalRequest: []byte(`{"cachedContent":"cached-gemini-3"}`),
Metadata: make(map[string]any),
},
bareID: "cached-gemini-3",
qualID: "geminicache:cached-gemini-3",
auth: authA,
},
{
name: "execution session metadata",
opts: cliproxyexecutor.Options{
Metadata: map[string]any{
cliproxyexecutor.ExecutionSessionMetadataKey: "exec-run-4",
},
},
bareID: "exec-run-4",
qualID: "execution:exec-run-4",
auth: authB,
},
}
for _, tc := range testCases {
t.Run(tc.name, func(t *testing.T) {
picked, errPick := selector.Pick(context.Background(), "anthropic", "claude-3-7-sonnet", tc.opts, []*coreauth.Auth{authA, authB})
if errPick != nil || picked.ID != tc.auth.ID {
t.Fatalf("pick failed: auth=%v err=%v", picked, errPick)
}
// Query via bare ID
reqBare, _ := json.Marshal(pluginapi.HostAffinityLookupRequest{
Provider: "anthropic",
Model: "claude-3-7-sonnet",
SessionID: tc.bareID,
})
rawBare, errCall := host.callFromPlugin(context.Background(), pluginabi.MethodHostAffinityLookup, reqBare)
if errCall != nil {
t.Fatalf("callFromPlugin error = %v", errCall)
}
respBare, _ := decodeRPCEnvelope[pluginapi.HostAffinityLookupResponse](rawBare)
if respBare.Status != pluginapi.HostAffinityStatusBound || respBare.AuthIndex != tc.auth.Index {
t.Fatalf("bare lookup %q: got status=%q index=%q, want bound index=%q", tc.bareID, respBare.Status, respBare.AuthIndex, tc.auth.Index)
}
// Query via qualified ID
reqQual, _ := json.Marshal(pluginapi.HostAffinityLookupRequest{
Provider: "anthropic",
Model: "claude-3-7-sonnet",
SessionID: tc.qualID,
})
rawQual, _ := host.callFromPlugin(context.Background(), pluginabi.MethodHostAffinityLookup, reqQual)
respQual, _ := decodeRPCEnvelope[pluginapi.HostAffinityLookupResponse](rawQual)
if respQual.Status != pluginapi.HostAffinityStatusBound || respQual.AuthIndex != tc.auth.Index {
t.Fatalf("qual lookup %q: got status=%q index=%q, want bound index=%q", tc.qualID, respQual.Status, respQual.AuthIndex, tc.auth.Index)
}
})
}
// Conflict test: bind task:shared-clash to authA and claude:shared-clash to authB
optsClash1 := cliproxyexecutor.Options{
Headers: map[string][]string{"X-Task-ID": {"shared-clash"}},
Metadata: make(map[string]any),
}
selector.Pick(context.Background(), "anthropic", "claude-3-7-sonnet", optsClash1, []*coreauth.Auth{authA})
optsClash2 := cliproxyexecutor.Options{
Headers: map[string][]string{"X-Claude-Code-Session-Id": {"shared-clash"}},
Metadata: make(map[string]any),
}
selector.Pick(context.Background(), "anthropic", "claude-3-7-sonnet", optsClash2, []*coreauth.Auth{authB})
reqClash, _ := json.Marshal(pluginapi.HostAffinityLookupRequest{
Provider: "anthropic",
Model: "claude-3-7-sonnet",
SessionID: "shared-clash",
})
rawClash, _ := host.callFromPlugin(context.Background(), pluginabi.MethodHostAffinityLookup, reqClash)
respClash, _ := decodeRPCEnvelope[pluginapi.HostAffinityLookupResponse](rawClash)
if respClash.Status != pluginapi.HostAffinityStatusAmbiguous {
t.Fatalf("clash lookup status = %q, want %q", respClash.Status, pluginapi.HostAffinityStatusAmbiguous)
}
})
t.Run("disabled credential returns bound with disabled flag", func(t *testing.T) {
sessionID := "disabled-cred-session"
opts := cliproxyexecutor.Options{
Headers: map[string][]string{"X-Claude-Code-Session-Id": {sessionID}},
Metadata: make(map[string]any),
}
selector.Pick(context.Background(), "anthropic", "claude-3-7-sonnet", opts, []*coreauth.Auth{authA})
// Disable authA in manager
authA.Disabled = true
manager.Update(context.Background(), authA)
defer func() {
authA.Disabled = false
manager.Update(context.Background(), authA)
}()
reqPayload, _ := json.Marshal(pluginapi.HostAffinityLookupRequest{
Provider: "anthropic",
Model: "claude-3-7-sonnet",
SessionID: sessionID,
})
rawResp, _ := host.callFromPlugin(context.Background(), pluginabi.MethodHostAffinityLookup, reqPayload)
resp, _ := decodeRPCEnvelope[pluginapi.HostAffinityLookupResponse](rawResp)
if resp.Status != pluginapi.HostAffinityStatusBound {
t.Fatalf("status = %q, want bound", resp.Status)
}
if !resp.Disabled {
t.Fatalf("expected disabled = true, got %v", resp.Disabled)
}
})
t.Run("unsupported when plugin scheduler active", func(t *testing.T) {
managerWithPluginSched := coreauth.NewManager(nil, nil, nil)
managerWithPluginSched.SetSelector(selector)
managerWithPluginSched.SetPluginScheduler(&mockPluginScheduler{})
hostWithPluginSched := New()
hostWithPluginSched.SetAuthManager(managerWithPluginSched)
reqPayload, _ := json.Marshal(pluginapi.HostAffinityLookupRequest{
Provider: "anthropic",
Model: "claude-3-7-sonnet",
SessionID: "bound-session-1",
})
rawResp, _ := hostWithPluginSched.callFromPlugin(context.Background(), pluginabi.MethodHostAffinityLookup, reqPayload)
resp, _ := decodeRPCEnvelope[pluginapi.HostAffinityLookupResponse](rawResp)
if resp.Status != pluginapi.HostAffinityStatusUnsupported {
t.Fatalf("status with plugin scheduler = %q, want %q", resp.Status, pluginapi.HostAffinityStatusUnsupported)
}
})
t.Run("unsupported when selector is not session affinity", func(t *testing.T) {
managerNoAffinity := coreauth.NewManager(nil, nil, nil)
managerNoAffinity.SetSelector(&coreauth.RoundRobinSelector{})
hostNoAffinity := New()
hostNoAffinity.SetAuthManager(managerNoAffinity)
reqPayload, _ := json.Marshal(pluginapi.HostAffinityLookupRequest{
Provider: "anthropic",
Model: "claude-3-7-sonnet",
SessionID: "bound-session-1",
})
rawResp, _ := hostNoAffinity.callFromPlugin(context.Background(), pluginabi.MethodHostAffinityLookup, reqPayload)
resp, _ := decodeRPCEnvelope[pluginapi.HostAffinityLookupResponse](rawResp)
if resp.Status != pluginapi.HostAffinityStatusUnsupported {
t.Fatalf("status with round robin = %q, want %q", resp.Status, pluginapi.HostAffinityStatusUnsupported)
}
})
t.Run("invalid input validation", func(t *testing.T) {
for _, tc := range []struct {
name string
req pluginapi.HostAffinityLookupRequest
}{
{"empty provider", pluginapi.HostAffinityLookupRequest{Provider: "", Model: "m", SessionID: "s"}},
{"empty model", pluginapi.HostAffinityLookupRequest{Provider: "p", Model: "", SessionID: "s"}},
{"empty session", pluginapi.HostAffinityLookupRequest{Provider: "p", Model: "m", SessionID: ""}},
} {
t.Run(tc.name, func(t *testing.T) {
reqPayload, _ := json.Marshal(tc.req)
_, errCall := host.callFromPlugin(context.Background(), pluginabi.MethodHostAffinityLookup, reqPayload)
if errCall == nil {
t.Fatalf("expected error for %s, got nil", tc.name)
}
})
}
})
t.Run("LCP session binding observation", func(t *testing.T) {
lcpOpts := cliproxyexecutor.Options{
OriginalRequest: []byte(`{"messages":[{"role":"system","content":"host-lcp-sys"},{"role":"user","content":"host-lcp-user"}]}`),
Metadata: map[string]any{
cliproxyexecutor.CallerScopeMetadataKey: "plugin-caller",
},
}
picked, errPick := selector.Pick(context.Background(), "anthropic", "claude-3-7-sonnet", lcpOpts, []*coreauth.Auth{authA})
if errPick != nil || picked.ID != authA.ID {
t.Fatalf("LCP Pick failed: %v", errPick)
}
lcpSessionID, _ := lcpOpts.Metadata[cliproxyexecutor.LCPAffinitySessionIDMetadataKey].(string)
if lcpSessionID == "" {
t.Fatalf("missing LCP session ID in metadata")
}
reqPayload, _ := json.Marshal(pluginapi.HostAffinityLookupRequest{
Provider: "anthropic",
Model: "claude-3-7-sonnet",
SessionID: lcpSessionID,
})
rawResp, errCall := host.callFromPlugin(context.Background(), pluginabi.MethodHostAffinityLookup, reqPayload)
if errCall != nil {
t.Fatalf("callFromPlugin error = %v", errCall)
}
resp, errDecode := decodeRPCEnvelope[pluginapi.HostAffinityLookupResponse](rawResp)
if errDecode != nil {
t.Fatalf("decode response: %v", errDecode)
}
if resp.Status != pluginapi.HostAffinityStatusBound || resp.AuthIndex != authA.Index {
t.Fatalf("LCP observation failed: got status=%q index=%q, want bound %q", resp.Status, resp.AuthIndex, authA.Index)
}
})
t.Run("concurrent reads", func(t *testing.T) {
reqPayload, _ := json.Marshal(pluginapi.HostAffinityLookupRequest{
Provider: "anthropic",
Model: "claude-3-7-sonnet",
SessionID: "bound-session-1",
})
var wg sync.WaitGroup
for i := 0; i < 20; i++ {
wg.Add(1)
go func() {
defer wg.Done()
rawResp, errCall := host.callFromPlugin(context.Background(), pluginabi.MethodHostAffinityLookup, reqPayload)
if errCall != nil {
t.Errorf("concurrent callFromPlugin error: %v", errCall)
return
}
resp, errDecode := decodeRPCEnvelope[pluginapi.HostAffinityLookupResponse](rawResp)
if errDecode != nil || resp.Status != pluginapi.HostAffinityStatusBound {
t.Errorf("concurrent response error: %#v %v", resp, errDecode)
}
}()
}
wg.Wait()
})
}