Add projects
This commit is contained in:
parent
2d3a9ad623
commit
8b607dd700
SSH key fingerprint:
SHA256:GBzT2DXvAuGV9XIV5W3WrzVpjU54FThmxHXdbz95J24
1802 changed files with 503346 additions and 2 deletions
340
backend/sdk/cliproxy/auth/home_unauthorized_refresh_test.go
Normal file
340
backend/sdk/cliproxy/auth/home_unauthorized_refresh_test.go
Normal file
|
|
@ -0,0 +1,340 @@
|
|||
package auth
|
||||
|
||||
import (
|
||||
"context"
|
||||
"encoding/json"
|
||||
"net/http"
|
||||
"sync/atomic"
|
||||
"testing"
|
||||
|
||||
internalconfig "github.com/router-for-me/CLIProxyAPI/v7/internal/config"
|
||||
"github.com/router-for-me/CLIProxyAPI/v7/sdk/cliproxy/executionregistry"
|
||||
cliproxyexecutor "github.com/router-for-me/CLIProxyAPI/v7/sdk/cliproxy/executor"
|
||||
)
|
||||
|
||||
const homeUnauthorizedRefreshProvider = "home-unauthorized-refresh"
|
||||
|
||||
type homeUnauthorizedRefreshDispatcher struct {
|
||||
calls atomic.Int32
|
||||
}
|
||||
|
||||
func (*homeUnauthorizedRefreshDispatcher) HeartbeatOK() bool { return true }
|
||||
|
||||
func (d *homeUnauthorizedRefreshDispatcher) RPopAuth(context.Context, string, string, http.Header, int) ([]byte, error) {
|
||||
d.calls.Add(1)
|
||||
return json.Marshal(homeAuthDispatchResponse{Auth: Auth{
|
||||
ID: "home-refresh-auth",
|
||||
Provider: homeUnauthorizedRefreshProvider,
|
||||
Status: StatusActive,
|
||||
Attributes: map[string]string{
|
||||
AttributeAuthKind: AuthKindOAuth,
|
||||
"websockets": "true",
|
||||
},
|
||||
Metadata: map[string]any{
|
||||
"access_token": "stale-access-token",
|
||||
},
|
||||
}})
|
||||
}
|
||||
|
||||
func (*homeUnauthorizedRefreshDispatcher) AbortAmbiguousDispatch() {}
|
||||
|
||||
type homeUnauthorizedRefreshExecutor struct {
|
||||
streamMode string
|
||||
refreshErr error
|
||||
keepStale bool
|
||||
retainSelection bool
|
||||
executeCalls atomic.Int32
|
||||
countCalls atomic.Int32
|
||||
streamCalls atomic.Int32
|
||||
refreshCalls atomic.Int32
|
||||
}
|
||||
|
||||
func (*homeUnauthorizedRefreshExecutor) Identifier() string { return homeUnauthorizedRefreshProvider }
|
||||
|
||||
func (e *homeUnauthorizedRefreshExecutor) Execute(_ context.Context, auth *Auth, _ cliproxyexecutor.Request, opts cliproxyexecutor.Options) (cliproxyexecutor.Response, error) {
|
||||
e.executeCalls.Add(1)
|
||||
if e.retainSelection {
|
||||
if lifecycle, ok := opts.ExecutionLifecycle.(interface{ Retain() }); ok {
|
||||
lifecycle.Retain()
|
||||
}
|
||||
}
|
||||
if authAccessToken(auth) == "stale-access-token" {
|
||||
return cliproxyexecutor.Response{}, &Error{HTTPStatus: http.StatusUnauthorized, Message: "expired access token"}
|
||||
}
|
||||
return cliproxyexecutor.Response{Payload: []byte("ok")}, nil
|
||||
}
|
||||
|
||||
func (e *homeUnauthorizedRefreshExecutor) ExecuteStream(_ context.Context, auth *Auth, _ cliproxyexecutor.Request, _ cliproxyexecutor.Options) (*cliproxyexecutor.StreamResult, error) {
|
||||
e.streamCalls.Add(1)
|
||||
if authAccessToken(auth) == "stale-access-token" {
|
||||
switch e.streamMode {
|
||||
case "bootstrap":
|
||||
chunks := make(chan cliproxyexecutor.StreamChunk, 1)
|
||||
chunks <- cliproxyexecutor.StreamChunk{Err: &Error{HTTPStatus: http.StatusUnauthorized, Message: "expired access token"}}
|
||||
close(chunks)
|
||||
return &cliproxyexecutor.StreamResult{Chunks: chunks}, nil
|
||||
case "started":
|
||||
chunks := make(chan cliproxyexecutor.StreamChunk, 2)
|
||||
chunks <- cliproxyexecutor.StreamChunk{Payload: []byte("started")}
|
||||
chunks <- cliproxyexecutor.StreamChunk{Err: &Error{HTTPStatus: http.StatusUnauthorized, Message: "expired access token"}}
|
||||
close(chunks)
|
||||
return &cliproxyexecutor.StreamResult{Chunks: chunks}, nil
|
||||
default:
|
||||
return nil, &Error{HTTPStatus: http.StatusUnauthorized, Message: "expired access token"}
|
||||
}
|
||||
}
|
||||
chunks := make(chan cliproxyexecutor.StreamChunk, 1)
|
||||
chunks <- cliproxyexecutor.StreamChunk{Payload: []byte("ok")}
|
||||
close(chunks)
|
||||
return &cliproxyexecutor.StreamResult{Chunks: chunks}, nil
|
||||
}
|
||||
|
||||
func (e *homeUnauthorizedRefreshExecutor) Refresh(_ context.Context, auth *Auth) (*Auth, error) {
|
||||
e.refreshCalls.Add(1)
|
||||
if e.refreshErr != nil {
|
||||
return nil, e.refreshErr
|
||||
}
|
||||
updated := auth.Clone()
|
||||
if e.keepStale {
|
||||
return updated, nil
|
||||
}
|
||||
if updated.Metadata == nil {
|
||||
updated.Metadata = make(map[string]any)
|
||||
}
|
||||
updated.Metadata["access_token"] = "fresh-access-token"
|
||||
return updated, nil
|
||||
}
|
||||
|
||||
func (e *homeUnauthorizedRefreshExecutor) CountTokens(_ context.Context, auth *Auth, _ cliproxyexecutor.Request, _ cliproxyexecutor.Options) (cliproxyexecutor.Response, error) {
|
||||
e.countCalls.Add(1)
|
||||
if authAccessToken(auth) == "stale-access-token" {
|
||||
return cliproxyexecutor.Response{}, &Error{HTTPStatus: http.StatusUnauthorized, Message: "expired access token"}
|
||||
}
|
||||
return cliproxyexecutor.Response{Payload: []byte("ok")}, nil
|
||||
}
|
||||
|
||||
func (*homeUnauthorizedRefreshExecutor) HttpRequest(context.Context, *Auth, *http.Request) (*http.Response, error) {
|
||||
return nil, nil
|
||||
}
|
||||
|
||||
func newHomeUnauthorizedRefreshManager(dispatcher *homeUnauthorizedRefreshDispatcher, executor *homeUnauthorizedRefreshExecutor) *Manager {
|
||||
manager := NewManager(nil, nil, nil)
|
||||
manager.SetConfig(&internalconfig.Config{Home: internalconfig.HomeConfig{Enabled: true}})
|
||||
manager.PublishHomeDispatch(dispatcher, executionregistry.New(), 1)
|
||||
manager.RegisterExecutor(executor)
|
||||
return manager
|
||||
}
|
||||
|
||||
func TestHomeUnauthorizedRefreshesSameSelectionBeforeRedispatch(t *testing.T) {
|
||||
for _, test := range []struct {
|
||||
name string
|
||||
run func(*Manager) error
|
||||
}{
|
||||
{
|
||||
name: "execute",
|
||||
run: func(manager *Manager) error {
|
||||
_, errExecute := manager.Execute(context.Background(), []string{homeUnauthorizedRefreshProvider}, cliproxyexecutor.Request{Model: "model-a"}, cliproxyexecutor.Options{})
|
||||
return errExecute
|
||||
},
|
||||
},
|
||||
{
|
||||
name: "count_tokens",
|
||||
run: func(manager *Manager) error {
|
||||
_, errCount := manager.ExecuteCount(context.Background(), []string{homeUnauthorizedRefreshProvider}, cliproxyexecutor.Request{Model: "model-a"}, cliproxyexecutor.Options{})
|
||||
return errCount
|
||||
},
|
||||
},
|
||||
} {
|
||||
t.Run(test.name, func(t *testing.T) {
|
||||
dispatcher := &homeUnauthorizedRefreshDispatcher{}
|
||||
executor := &homeUnauthorizedRefreshExecutor{}
|
||||
manager := newHomeUnauthorizedRefreshManager(dispatcher, executor)
|
||||
|
||||
if errRun := test.run(manager); errRun != nil {
|
||||
t.Fatalf("execution error = %v", errRun)
|
||||
}
|
||||
if got := dispatcher.calls.Load(); got != 1 {
|
||||
t.Fatalf("Home dispatch calls = %d, want 1", got)
|
||||
}
|
||||
if got := executor.refreshCalls.Load(); got != 1 {
|
||||
t.Fatalf("refresh calls = %d, want 1", got)
|
||||
}
|
||||
if test.name == "execute" && executor.executeCalls.Load() != 2 {
|
||||
t.Fatalf("execute calls = %d, want 2", executor.executeCalls.Load())
|
||||
}
|
||||
if test.name == "count_tokens" && executor.countCalls.Load() != 2 {
|
||||
t.Fatalf("count calls = %d, want 2", executor.countCalls.Load())
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestHomeUnauthorizedRefreshUpdatesRetainedSelection(t *testing.T) {
|
||||
dispatcher := &homeUnauthorizedRefreshDispatcher{}
|
||||
executor := &homeUnauthorizedRefreshExecutor{retainSelection: true}
|
||||
manager := newHomeUnauthorizedRefreshManager(dispatcher, executor)
|
||||
ctx := cliproxyexecutor.WithDownstreamWebsocket(context.Background())
|
||||
opts := cliproxyexecutor.Options{Metadata: map[string]any{
|
||||
cliproxyexecutor.ExecutionSessionMetadataKey: "refresh-session",
|
||||
cliproxyexecutor.PinnedAuthMetadataKey: "home-refresh-auth",
|
||||
}}
|
||||
|
||||
for range 2 {
|
||||
if _, errExecute := manager.Execute(ctx, []string{homeUnauthorizedRefreshProvider}, cliproxyexecutor.Request{Model: "model-a"}, opts); errExecute != nil {
|
||||
t.Fatalf("Execute() error = %v", errExecute)
|
||||
}
|
||||
}
|
||||
if got := dispatcher.calls.Load(); got != 1 {
|
||||
t.Fatalf("Home dispatch calls = %d, want one retained selection", got)
|
||||
}
|
||||
if got := executor.refreshCalls.Load(); got != 1 {
|
||||
t.Fatalf("refresh calls = %d, want refreshed token reused by retained selection", got)
|
||||
}
|
||||
if got := executor.executeCalls.Load(); got != 3 {
|
||||
t.Fatalf("execute calls = %d, want stale attempt, retry, and retained reuse", got)
|
||||
}
|
||||
}
|
||||
|
||||
func TestRefreshHomeSelectionReusesConcurrentNewerToken(t *testing.T) {
|
||||
executor := &homeUnauthorizedRefreshExecutor{}
|
||||
selection := &HomeDispatchSelection{
|
||||
Auth: &Auth{ID: "home-refresh-auth", Provider: homeUnauthorizedRefreshProvider, Attributes: map[string]string{AttributeAuthKind: AuthKindOAuth}, Metadata: map[string]any{"access_token": "fresh-access-token"}},
|
||||
Executor: executor,
|
||||
Provider: homeUnauthorizedRefreshProvider,
|
||||
}
|
||||
failed := &Auth{ID: "home-refresh-auth", Provider: homeUnauthorizedRefreshProvider, Attributes: map[string]string{AttributeAuthKind: AuthKindOAuth}, Metadata: map[string]any{"access_token": "stale-access-token"}}
|
||||
manager := NewManager(nil, nil, nil)
|
||||
|
||||
updated, reused, errRefresh := manager.RefreshHomeSelectionAfterUnauthorized(context.Background(), selection, failed)
|
||||
if errRefresh != nil || !reused || authAccessToken(updated) != "fresh-access-token" {
|
||||
t.Fatalf("RefreshHomeSelectionAfterUnauthorized() = %#v, %v, %v", updated, reused, errRefresh)
|
||||
}
|
||||
if got := executor.refreshCalls.Load(); got != 0 {
|
||||
t.Fatalf("refresh calls = %d, want 0 when selection already has a newer token", got)
|
||||
}
|
||||
}
|
||||
|
||||
func TestHomeUnauthorizedRefreshIsAttemptedAtMostOnce(t *testing.T) {
|
||||
dispatcher := &homeUnauthorizedRefreshDispatcher{}
|
||||
executor := &homeUnauthorizedRefreshExecutor{keepStale: true}
|
||||
manager := newHomeUnauthorizedRefreshManager(dispatcher, executor)
|
||||
|
||||
_, errExecute := manager.Execute(context.Background(), []string{homeUnauthorizedRefreshProvider}, cliproxyexecutor.Request{Model: "model-a"}, cliproxyexecutor.Options{})
|
||||
if statusCodeFromError(errExecute) != http.StatusUnauthorized {
|
||||
t.Fatalf("Execute() error = %v, want original 401", errExecute)
|
||||
}
|
||||
if got := executor.refreshCalls.Load(); got != 1 {
|
||||
t.Fatalf("refresh calls = %d, want exactly 1", got)
|
||||
}
|
||||
if got := executor.executeCalls.Load(); got != 2 {
|
||||
t.Fatalf("execute calls = %d, want initial attempt and one retry", got)
|
||||
}
|
||||
}
|
||||
|
||||
func TestHomeNoCandidateAfterRefreshFailurePreservesRefreshError(t *testing.T) {
|
||||
refreshErr := &Error{Code: "refresh_temporarily_unavailable", HTTPStatus: http.StatusServiceUnavailable, Message: "refresh unavailable"}
|
||||
noCandidate := &Error{Code: "auth_not_found", HTTPStatus: http.StatusServiceUnavailable, Message: "no auth available"}
|
||||
if !shouldReturnLastErrorOnPickFailure(true, refreshErr, noCandidate) {
|
||||
t.Fatal("Home no-candidate error would overwrite the original refresh error")
|
||||
}
|
||||
}
|
||||
|
||||
func TestHomeUnauthorizedTransientRefreshFailureIsReturned(t *testing.T) {
|
||||
dispatcher := &homeUnauthorizedRefreshDispatcher{}
|
||||
executor := &homeUnauthorizedRefreshExecutor{
|
||||
refreshErr: &Error{HTTPStatus: http.StatusServiceUnavailable, Message: "Home refresh temporarily unavailable"},
|
||||
}
|
||||
manager := newHomeUnauthorizedRefreshManager(dispatcher, executor)
|
||||
|
||||
_, errExecute := manager.Execute(context.Background(), []string{homeUnauthorizedRefreshProvider}, cliproxyexecutor.Request{Model: "model-a"}, cliproxyexecutor.Options{})
|
||||
if statusCodeFromError(errExecute) != http.StatusServiceUnavailable {
|
||||
t.Fatalf("Execute() error = %v, want transient 503", errExecute)
|
||||
}
|
||||
if got := executor.executeCalls.Load(); got != 1 {
|
||||
t.Fatalf("execute calls = %d, want 1", got)
|
||||
}
|
||||
if got := executor.refreshCalls.Load(); got != 1 {
|
||||
t.Fatalf("refresh calls = %d, want 1", got)
|
||||
}
|
||||
}
|
||||
|
||||
func TestHomeUnauthorizedStreamRefreshesAtMostOnceAcrossRedispatch(t *testing.T) {
|
||||
dispatcher := &homeUnauthorizedRefreshDispatcher{}
|
||||
executor := &homeUnauthorizedRefreshExecutor{keepStale: true}
|
||||
manager := newHomeUnauthorizedRefreshManager(dispatcher, executor)
|
||||
|
||||
_, errStream := manager.ExecuteStream(context.Background(), []string{homeUnauthorizedRefreshProvider}, cliproxyexecutor.Request{Model: "model-a"}, cliproxyexecutor.Options{Stream: true})
|
||||
if statusCodeFromError(errStream) != http.StatusUnauthorized {
|
||||
t.Fatalf("ExecuteStream() error = %v, want original 401", errStream)
|
||||
}
|
||||
if got := executor.refreshCalls.Load(); got != 1 {
|
||||
t.Fatalf("refresh calls = %d, want exactly 1", got)
|
||||
}
|
||||
if got := executor.streamCalls.Load(); got != 2 {
|
||||
t.Fatalf("stream calls = %d, want initial attempt and one retry", got)
|
||||
}
|
||||
}
|
||||
|
||||
func TestHomeUnauthorizedStartedStreamDoesNotReplay(t *testing.T) {
|
||||
dispatcher := &homeUnauthorizedRefreshDispatcher{}
|
||||
executor := &homeUnauthorizedRefreshExecutor{streamMode: "started"}
|
||||
manager := newHomeUnauthorizedRefreshManager(dispatcher, executor)
|
||||
|
||||
result, errStream := manager.ExecuteStream(context.Background(), []string{homeUnauthorizedRefreshProvider}, cliproxyexecutor.Request{Model: "model-a"}, cliproxyexecutor.Options{Stream: true})
|
||||
if errStream != nil {
|
||||
t.Fatalf("ExecuteStream() error = %v", errStream)
|
||||
}
|
||||
sawPayload := false
|
||||
sawUnauthorized := false
|
||||
for chunk := range result.Chunks {
|
||||
if string(chunk.Payload) == "started" {
|
||||
sawPayload = true
|
||||
}
|
||||
if statusCodeFromError(chunk.Err) == http.StatusUnauthorized {
|
||||
sawUnauthorized = true
|
||||
}
|
||||
}
|
||||
if !sawPayload || !sawUnauthorized {
|
||||
t.Fatalf("stream results = payload %v unauthorized %v, want both", sawPayload, sawUnauthorized)
|
||||
}
|
||||
if got := executor.refreshCalls.Load(); got != 0 {
|
||||
t.Fatalf("refresh calls = %d, want 0 after stream started", got)
|
||||
}
|
||||
if got := executor.streamCalls.Load(); got != 1 {
|
||||
t.Fatalf("stream calls = %d, want 1", got)
|
||||
}
|
||||
}
|
||||
|
||||
func TestHomeUnauthorizedStreamRefreshesBeforeRedispatch(t *testing.T) {
|
||||
for _, mode := range []string{"synchronous", "bootstrap"} {
|
||||
t.Run(mode, func(t *testing.T) {
|
||||
dispatcher := &homeUnauthorizedRefreshDispatcher{}
|
||||
executor := &homeUnauthorizedRefreshExecutor{streamMode: mode}
|
||||
manager := newHomeUnauthorizedRefreshManager(dispatcher, executor)
|
||||
|
||||
result, errStream := manager.ExecuteStream(context.Background(), []string{homeUnauthorizedRefreshProvider}, cliproxyexecutor.Request{Model: "model-a"}, cliproxyexecutor.Options{Stream: true})
|
||||
if errStream != nil {
|
||||
t.Fatalf("ExecuteStream() error = %v", errStream)
|
||||
}
|
||||
var payload string
|
||||
for chunk := range result.Chunks {
|
||||
if chunk.Err != nil {
|
||||
t.Fatalf("stream chunk error = %v", chunk.Err)
|
||||
}
|
||||
payload += string(chunk.Payload)
|
||||
}
|
||||
if payload != "ok" {
|
||||
t.Fatalf("stream payload = %q, want ok", payload)
|
||||
}
|
||||
if got := dispatcher.calls.Load(); got != 1 {
|
||||
t.Fatalf("Home dispatch calls = %d, want 1", got)
|
||||
}
|
||||
if got := executor.refreshCalls.Load(); got != 1 {
|
||||
t.Fatalf("refresh calls = %d, want 1", got)
|
||||
}
|
||||
if got := executor.streamCalls.Load(); got != 2 {
|
||||
t.Fatalf("stream calls = %d, want 2", got)
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
Loading…
Reference in a new issue