97 lines
4.1 KiB
Go
97 lines
4.1 KiB
Go
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"
|
|
)
|
|
|
|
type selectedAuthCallbackDispatcher struct {
|
|
calls atomic.Int32
|
|
}
|
|
|
|
func (*selectedAuthCallbackDispatcher) HeartbeatOK() bool { return true }
|
|
func (d *selectedAuthCallbackDispatcher) RPopAuth(context.Context, string, string, http.Header, int) ([]byte, error) {
|
|
if d.calls.Add(1) > 2 {
|
|
return json.Marshal(homeErrorEnvelope{Error: &homeErrorDetail{Code: homeRequestRetryExceededErrorCode, Message: "no more auths"}})
|
|
}
|
|
return json.Marshal(homeAuthDispatchResponse{Auth: Auth{ID: "home-auth", Provider: "home-execution", Status: StatusActive, Attributes: map[string]string{"websockets": "true"}}})
|
|
}
|
|
func (*selectedAuthCallbackDispatcher) AbortAmbiguousDispatch() {}
|
|
|
|
type callbackPinHomeExecutor struct {
|
|
manager *Manager
|
|
session string
|
|
calls atomic.Int32
|
|
}
|
|
|
|
func (*callbackPinHomeExecutor) Identifier() string { return "home-execution" }
|
|
func (e *callbackPinHomeExecutor) Execute(context.Context, *Auth, cliproxyexecutor.Request, cliproxyexecutor.Options) (cliproxyexecutor.Response, error) {
|
|
return cliproxyexecutor.Response{}, nil
|
|
}
|
|
func (e *callbackPinHomeExecutor) ExecuteStream(_ context.Context, _ *Auth, _ cliproxyexecutor.Request, opts cliproxyexecutor.Options) (*cliproxyexecutor.StreamResult, error) {
|
|
if e.calls.Add(1) == 2 {
|
|
return nil, errSelectedAuthCallbackFailure
|
|
}
|
|
if lifecycle, ok := opts.ExecutionLifecycle.(interface{ Retain() }); ok {
|
|
lifecycle.Retain()
|
|
}
|
|
chunks := make(chan cliproxyexecutor.StreamChunk, 1)
|
|
chunks <- cliproxyexecutor.StreamChunk{Payload: []byte(`{"type":"response.completed"}`)}
|
|
close(chunks)
|
|
return &cliproxyexecutor.StreamResult{Chunks: chunks}, nil
|
|
}
|
|
func (*callbackPinHomeExecutor) Refresh(context.Context, *Auth) (*Auth, error) { return nil, nil }
|
|
func (*callbackPinHomeExecutor) CountTokens(context.Context, *Auth, cliproxyexecutor.Request, cliproxyexecutor.Options) (cliproxyexecutor.Response, error) {
|
|
return cliproxyexecutor.Response{}, nil
|
|
}
|
|
func (*callbackPinHomeExecutor) HttpRequest(context.Context, *Auth, *http.Request) (*http.Response, error) {
|
|
return nil, nil
|
|
}
|
|
|
|
var errSelectedAuthCallbackFailure = &Error{HTTPStatus: 502, Message: "selected auth failed"}
|
|
|
|
func TestHomeSelectedAuthCallbackPinsFirstHandlerSelectionAndCleansFailure(t *testing.T) {
|
|
manager := NewManager(nil, nil, nil)
|
|
manager.SetConfig(&internalconfig.Config{Home: internalconfig.HomeConfig{Enabled: true}})
|
|
manager.PublishHomeDispatch(&selectedAuthCallbackDispatcher{}, executionregistry.New(), 1)
|
|
executor := &callbackPinHomeExecutor{manager: manager, session: "callback-session"}
|
|
manager.RegisterExecutor(executor)
|
|
|
|
ctx := cliproxyexecutor.WithDownstreamWebsocket(context.Background())
|
|
callbackSawRuntimeAuth := false
|
|
opts := cliproxyexecutor.Options{Stream: true, Metadata: map[string]any{
|
|
cliproxyexecutor.ExecutionSessionMetadataKey: executor.session,
|
|
cliproxyexecutor.SelectedAuthCallbackMetadataKey: func(authID string) {
|
|
_, callbackSawRuntimeAuth = manager.GetExecutionSessionAuthByID(executor.session, authID)
|
|
},
|
|
}}
|
|
result, errExecute := manager.ExecuteStream(ctx, []string{"home-execution"}, cliproxyexecutor.Request{Model: "model-a"}, opts)
|
|
if errExecute != nil {
|
|
t.Fatalf("first ExecuteStream() error = %v", errExecute)
|
|
}
|
|
for range result.Chunks {
|
|
}
|
|
if !callbackSawRuntimeAuth {
|
|
t.Fatal("first selected-auth callback could not resolve the Home runtime auth")
|
|
}
|
|
|
|
manager.CloseExecutionSession(executor.session)
|
|
callbackSawRuntimeAuth = false
|
|
_, errExecute = manager.ExecuteStream(ctx, []string{"home-execution"}, cliproxyexecutor.Request{Model: "model-b"}, opts)
|
|
if errExecute == nil {
|
|
t.Fatal("failed ExecuteStream() error = nil")
|
|
}
|
|
if !callbackSawRuntimeAuth {
|
|
t.Fatal("failed selected-auth callback could not resolve the Home runtime auth")
|
|
}
|
|
if _, ok := manager.GetExecutionSessionAuthByID(executor.session, "home-auth"); ok {
|
|
t.Fatal("failed selection retained Home runtime auth")
|
|
}
|
|
}
|