mirror of
https://github.com/router-for-me/CLIProxyAPI.git
synced 2026-09-20 23:59:06 +08:00
601ca43090
Redefine request-retry as additional credential retry rounds and enforce max-retry-credentials per round. Home dispatch now carries excluded and pinned auth constraints, supports remote retry limits, and propagates cooldown retry-after metadata across exhausted rounds. Move Antigravity upstream retries under conductor ownership to avoid double-consuming retry attempts. Update configuration comments and add coverage for Home retry rounds, cooldown handling, pinned credentials, and legacy dispatcher compatibility.
1525 lines
61 KiB
Go
1525 lines
61 KiB
Go
package auth
|
|
|
|
import (
|
|
"context"
|
|
"encoding/json"
|
|
"net/http"
|
|
"strconv"
|
|
"sync"
|
|
"sync/atomic"
|
|
"testing"
|
|
"time"
|
|
|
|
internalconfig "github.com/router-for-me/CLIProxyAPI/v7/internal/config"
|
|
"github.com/router-for-me/CLIProxyAPI/v7/internal/home"
|
|
internallogging "github.com/router-for-me/CLIProxyAPI/v7/internal/logging"
|
|
"github.com/router-for-me/CLIProxyAPI/v7/sdk/cliproxy/executionregistry"
|
|
cliproxyexecutor "github.com/router-for-me/CLIProxyAPI/v7/sdk/cliproxy/executor"
|
|
coreusage "github.com/router-for-me/CLIProxyAPI/v7/sdk/cliproxy/usage"
|
|
log "github.com/sirupsen/logrus"
|
|
logtest "github.com/sirupsen/logrus/hooks/test"
|
|
)
|
|
|
|
type homeExecutionDispatcher struct{}
|
|
|
|
func (homeExecutionDispatcher) HeartbeatOK() bool { return true }
|
|
|
|
func (homeExecutionDispatcher) RPopAuth(context.Context, string, string, http.Header, int) ([]byte, error) {
|
|
return json.Marshal(homeAuthDispatchResponse{Auth: Auth{ID: "home-auth", Provider: "home-execution", Status: StatusActive}})
|
|
}
|
|
|
|
func (homeExecutionDispatcher) AbortAmbiguousDispatch() {}
|
|
|
|
type homeExecutionStreamExecutor struct {
|
|
chunks <-chan cliproxyexecutor.StreamChunk
|
|
}
|
|
|
|
type homeExecutionExecutor struct {
|
|
ctx context.Context
|
|
}
|
|
|
|
func (*homeExecutionExecutor) Identifier() string { return "home-execution" }
|
|
func (e *homeExecutionExecutor) Execute(ctx context.Context, _ *Auth, _ cliproxyexecutor.Request, _ cliproxyexecutor.Options) (cliproxyexecutor.Response, error) {
|
|
e.ctx = ctx
|
|
if errCtx := ctx.Err(); errCtx != nil {
|
|
return cliproxyexecutor.Response{}, errCtx
|
|
}
|
|
return cliproxyexecutor.Response{Payload: []byte("ok")}, nil
|
|
}
|
|
func (*homeExecutionExecutor) ExecuteStream(context.Context, *Auth, cliproxyexecutor.Request, cliproxyexecutor.Options) (*cliproxyexecutor.StreamResult, error) {
|
|
return nil, nil
|
|
}
|
|
func (*homeExecutionExecutor) Refresh(context.Context, *Auth) (*Auth, error) { return nil, nil }
|
|
func (*homeExecutionExecutor) CountTokens(context.Context, *Auth, cliproxyexecutor.Request, cliproxyexecutor.Options) (cliproxyexecutor.Response, error) {
|
|
return cliproxyexecutor.Response{}, nil
|
|
}
|
|
func (*homeExecutionExecutor) HttpRequest(context.Context, *Auth, *http.Request) (*http.Response, error) {
|
|
return nil, nil
|
|
}
|
|
|
|
func (*homeExecutionStreamExecutor) Identifier() string { return "home-execution" }
|
|
func (*homeExecutionStreamExecutor) Execute(context.Context, *Auth, cliproxyexecutor.Request, cliproxyexecutor.Options) (cliproxyexecutor.Response, error) {
|
|
return cliproxyexecutor.Response{}, nil
|
|
}
|
|
func (e *homeExecutionStreamExecutor) ExecuteStream(context.Context, *Auth, cliproxyexecutor.Request, cliproxyexecutor.Options) (*cliproxyexecutor.StreamResult, error) {
|
|
return &cliproxyexecutor.StreamResult{Chunks: e.chunks}, nil
|
|
}
|
|
func (*homeExecutionStreamExecutor) Refresh(context.Context, *Auth) (*Auth, error) { return nil, nil }
|
|
func (*homeExecutionStreamExecutor) CountTokens(context.Context, *Auth, cliproxyexecutor.Request, cliproxyexecutor.Options) (cliproxyexecutor.Response, error) {
|
|
return cliproxyexecutor.Response{}, nil
|
|
}
|
|
func (*homeExecutionStreamExecutor) HttpRequest(context.Context, *Auth, *http.Request) (*http.Response, error) {
|
|
return nil, nil
|
|
}
|
|
|
|
func TestHomeModeNeverAuthorizesLocalAuthFallback(t *testing.T) {
|
|
manager := NewManager(nil, nil, nil)
|
|
cfg := &internalconfig.Config{}
|
|
cfg.Home.Enabled = true
|
|
manager.runtimeConfig.Store(cfg)
|
|
manager.auths["local-antigravity"] = &Auth{ID: "local-antigravity", Provider: "antigravity", Status: StatusActive}
|
|
|
|
if manager.localExecutionAllowed() {
|
|
t.Fatal("local execution allowed in Home mode")
|
|
}
|
|
if selected := manager.localFallbackAuth("local-antigravity"); selected != nil {
|
|
t.Fatalf("local fallback auth = %#v", selected)
|
|
}
|
|
}
|
|
|
|
func TestHomeSelectionEndsAfterExecute(t *testing.T) {
|
|
manager := NewManager(nil, nil, nil)
|
|
manager.SetConfig(&internalconfig.Config{Home: internalconfig.HomeConfig{Enabled: true}})
|
|
manager.PublishHomeDispatch(homeExecutionDispatcher{}, executionregistry.New(), 1)
|
|
executor := &homeExecutionExecutor{}
|
|
manager.RegisterExecutor(executor)
|
|
|
|
if _, errExecute := manager.Execute(context.Background(), []string{"home-execution"}, cliproxyexecutor.Request{Model: "test"}, cliproxyexecutor.Options{}); errExecute != nil {
|
|
t.Fatalf("Execute() error = %v", errExecute)
|
|
}
|
|
if executor.ctx == nil {
|
|
t.Fatal("executor did not receive an attempt context")
|
|
}
|
|
if errCtx := executor.ctx.Err(); errCtx == nil {
|
|
t.Fatal("attempt context was not canceled after execution")
|
|
}
|
|
}
|
|
|
|
func TestHomeNonStreamingExecutionLogsSelectedOAuthAuth(t *testing.T) {
|
|
previousLevel := log.GetLevel()
|
|
log.SetLevel(log.DebugLevel)
|
|
hook := logtest.NewLocal(log.StandardLogger())
|
|
t.Cleanup(func() {
|
|
hook.Reset()
|
|
log.SetLevel(previousLevel)
|
|
})
|
|
|
|
tests := []struct {
|
|
name string
|
|
run func(*Manager, context.Context) error
|
|
}{
|
|
{
|
|
name: "execute",
|
|
run: func(manager *Manager, ctx context.Context) error {
|
|
_, errExecute := manager.Execute(ctx, []string{"home-execution"}, cliproxyexecutor.Request{Model: "model-a"}, cliproxyexecutor.Options{})
|
|
return errExecute
|
|
},
|
|
},
|
|
{
|
|
name: "count_tokens",
|
|
run: func(manager *Manager, ctx context.Context) error {
|
|
_, errCount := manager.ExecuteCount(ctx, []string{"home-execution"}, cliproxyexecutor.Request{Model: "model-a"}, cliproxyexecutor.Options{})
|
|
return errCount
|
|
},
|
|
},
|
|
}
|
|
|
|
for _, tt := range tests {
|
|
t.Run(tt.name, func(t *testing.T) {
|
|
hook.Reset()
|
|
manager := NewManager(nil, nil, nil)
|
|
manager.SetConfig(&internalconfig.Config{Home: internalconfig.HomeConfig{Enabled: true}})
|
|
manager.PublishHomeDispatch(homeOAuthLoggingDispatcher{}, executionregistry.New(), 1)
|
|
manager.RegisterExecutor(&homeExecutionExecutor{})
|
|
|
|
ctx := internallogging.WithRequestID(context.Background(), "req-home-log")
|
|
if errRun := tt.run(manager, ctx); errRun != nil {
|
|
t.Fatalf("execution error = %v", errRun)
|
|
}
|
|
|
|
const expected = "Use OAuth provider=home-execution auth_file=home-auth for model model-a via socks5 proxy"
|
|
for _, entry := range hook.AllEntries() {
|
|
if entry.Level == log.DebugLevel && entry.Message == expected {
|
|
if got := entry.Data["request_id"]; got != "req-home-log" {
|
|
t.Fatalf("request_id = %v, want req-home-log", got)
|
|
}
|
|
return
|
|
}
|
|
}
|
|
t.Fatalf("selected auth log %q not found", expected)
|
|
})
|
|
}
|
|
}
|
|
|
|
type homeOAuthLoggingDispatcher struct{}
|
|
|
|
func (homeOAuthLoggingDispatcher) HeartbeatOK() bool { return true }
|
|
|
|
func (homeOAuthLoggingDispatcher) RPopAuth(context.Context, string, string, http.Header, int) ([]byte, error) {
|
|
return json.Marshal(homeAuthDispatchResponse{Auth: Auth{
|
|
ID: "home-auth",
|
|
Provider: "home-execution",
|
|
ProxyURL: "socks5://127.0.0.1:1080",
|
|
Status: StatusActive,
|
|
Attributes: map[string]string{
|
|
AttributeAuthKind: AuthKindOAuth,
|
|
},
|
|
}})
|
|
}
|
|
|
|
func (homeOAuthLoggingDispatcher) AbortAmbiguousDispatch() {}
|
|
|
|
func TestHomeSelectionEndsOnMissingExecutor(t *testing.T) {
|
|
manager := NewManager(nil, nil, nil)
|
|
manager.SetConfig(&internalconfig.Config{Home: internalconfig.HomeConfig{Enabled: true}})
|
|
registry := executionregistry.New()
|
|
manager.PublishHomeDispatch(homeExecutionDispatcher{}, registry, 1)
|
|
|
|
if _, errExecute := manager.Execute(context.Background(), []string{"home-execution"}, cliproxyexecutor.Request{Model: "test"}, cliproxyexecutor.Options{}); errExecute == nil {
|
|
t.Fatal("Execute() error = nil, want missing executor")
|
|
}
|
|
if errDrain := registry.Drain(context.Background()); errDrain != nil {
|
|
t.Fatalf("Drain() error = %v", errDrain)
|
|
}
|
|
}
|
|
|
|
func TestHomeSelectionClosesAttemptAndWebSocketResources(t *testing.T) {
|
|
registry := executionregistry.New()
|
|
pending, errBegin := registry.BeginDispatch()
|
|
if errBegin != nil {
|
|
t.Fatal(errBegin)
|
|
}
|
|
scope, errInstall := registry.Install(pending, executionregistry.ScopeSpec{})
|
|
if errInstall != nil {
|
|
t.Fatal(errInstall)
|
|
}
|
|
selection, errSelection := newHomeDispatchSelection(&Auth{ID: "home-auth"}, nil, "home-execution", scope)
|
|
if errSelection != nil {
|
|
t.Fatal(errSelection)
|
|
}
|
|
attemptCtx, releaseAttempt, errBind := homeExecutionAttemptContext(context.Background(), selection)
|
|
if errBind != nil {
|
|
t.Fatal(errBind)
|
|
}
|
|
var closeCalls atomic.Int32
|
|
if errBind = selection.Bind(func() error {
|
|
closeCalls.Add(1)
|
|
return nil
|
|
}); errBind != nil {
|
|
t.Fatal(errBind)
|
|
}
|
|
selection.End("completed")
|
|
releaseAttempt()
|
|
if errCtx := attemptCtx.Err(); errCtx == nil {
|
|
t.Fatal("attempt context was not canceled")
|
|
}
|
|
if got := closeCalls.Load(); got != 1 {
|
|
t.Fatalf("resource close calls = %d, want 1", got)
|
|
}
|
|
}
|
|
|
|
func TestHomeStreamConsumerCancelEndsSelection(t *testing.T) {
|
|
ctx, cancel := context.WithCancel(context.Background())
|
|
defer cancel()
|
|
manager := NewManager(nil, nil, nil)
|
|
manager.SetConfig(&internalconfig.Config{Home: internalconfig.HomeConfig{Enabled: true}})
|
|
registry := executionregistry.New()
|
|
manager.PublishHomeDispatch(homeExecutionDispatcher{}, registry, 1)
|
|
chunks := make(chan cliproxyexecutor.StreamChunk, 1)
|
|
chunks <- cliproxyexecutor.StreamChunk{Payload: []byte("initial")}
|
|
manager.RegisterExecutor(&homeExecutionStreamExecutor{chunks: chunks})
|
|
|
|
result, errExecute := manager.ExecuteStream(ctx, []string{"home-execution"}, cliproxyexecutor.Request{Model: "test"}, cliproxyexecutor.Options{Stream: true})
|
|
if errExecute != nil {
|
|
t.Fatalf("ExecuteStream() error = %v", errExecute)
|
|
}
|
|
cancel()
|
|
for range result.Chunks {
|
|
}
|
|
if errDrain := registry.Drain(context.Background()); errDrain != nil {
|
|
t.Fatalf("Drain() error = %v", errDrain)
|
|
}
|
|
}
|
|
|
|
type retainingHomeExecutionDispatcher struct {
|
|
calls atomic.Int32
|
|
}
|
|
|
|
func (d *retainingHomeExecutionDispatcher) HeartbeatOK() bool { return true }
|
|
|
|
func (d *retainingHomeExecutionDispatcher) RPopAuth(context.Context, string, string, http.Header, int) ([]byte, error) {
|
|
d.calls.Add(1)
|
|
return json.Marshal(homeAuthDispatchResponse{Auth: Auth{
|
|
ID: "home-auth",
|
|
Provider: "home-execution",
|
|
Status: StatusActive,
|
|
Attributes: map[string]string{
|
|
"websockets": "true",
|
|
},
|
|
}})
|
|
}
|
|
|
|
func (*retainingHomeExecutionDispatcher) AbortAmbiguousDispatch() {}
|
|
|
|
type retainingHomeExecutionExecutor struct {
|
|
calls atomic.Int32
|
|
}
|
|
|
|
func (*retainingHomeExecutionExecutor) Identifier() string { return "home-execution" }
|
|
|
|
func (e *retainingHomeExecutionExecutor) Execute(_ context.Context, _ *Auth, _ cliproxyexecutor.Request, opts cliproxyexecutor.Options) (cliproxyexecutor.Response, error) {
|
|
e.calls.Add(1)
|
|
if lifecycle, ok := opts.ExecutionLifecycle.(interface{ Retain() }); ok {
|
|
lifecycle.Retain()
|
|
}
|
|
return cliproxyexecutor.Response{Payload: []byte("ok")}, nil
|
|
}
|
|
|
|
func (*retainingHomeExecutionExecutor) ExecuteStream(context.Context, *Auth, cliproxyexecutor.Request, cliproxyexecutor.Options) (*cliproxyexecutor.StreamResult, error) {
|
|
return nil, nil
|
|
}
|
|
func (*retainingHomeExecutionExecutor) Refresh(context.Context, *Auth) (*Auth, error) {
|
|
return nil, nil
|
|
}
|
|
func (*retainingHomeExecutionExecutor) CountTokens(context.Context, *Auth, cliproxyexecutor.Request, cliproxyexecutor.Options) (cliproxyexecutor.Response, error) {
|
|
return cliproxyexecutor.Response{}, nil
|
|
}
|
|
func (*retainingHomeExecutionExecutor) HttpRequest(context.Context, *Auth, *http.Request) (*http.Response, error) {
|
|
return nil, nil
|
|
}
|
|
|
|
func TestHomeWebsocketSessionReusesRetainedSelection(t *testing.T) {
|
|
dispatcher := &retainingHomeExecutionDispatcher{}
|
|
manager := NewManager(nil, nil, nil)
|
|
manager.SetConfig(&internalconfig.Config{Home: internalconfig.HomeConfig{Enabled: true}})
|
|
manager.PublishHomeDispatch(dispatcher, executionregistry.New(), 1)
|
|
executor := &retainingHomeExecutionExecutor{}
|
|
manager.RegisterExecutor(executor)
|
|
|
|
ctx := cliproxyexecutor.WithDownstreamWebsocket(context.Background())
|
|
opts := cliproxyexecutor.Options{Metadata: map[string]any{
|
|
cliproxyexecutor.ExecutionSessionMetadataKey: "session-1",
|
|
cliproxyexecutor.PinnedAuthMetadataKey: "home-auth",
|
|
}}
|
|
for range 2 {
|
|
if _, errExecute := manager.Execute(ctx, []string{"home-execution"}, cliproxyexecutor.Request{Model: "model-a"}, opts); errExecute != nil {
|
|
t.Fatalf("Execute() error = %v", errExecute)
|
|
}
|
|
}
|
|
if got := dispatcher.calls.Load(); got != 1 {
|
|
t.Fatalf("Home RPOP calls = %d, want 1 for one retained session target", got)
|
|
}
|
|
if got := executor.calls.Load(); got != 2 {
|
|
t.Fatalf("executor calls = %d, want 2", got)
|
|
}
|
|
}
|
|
|
|
type changingHomeTargetDispatcher struct {
|
|
calls atomic.Int32
|
|
firstSelection *HomeDispatchSelection
|
|
oldEndedBeforeRPop atomic.Bool
|
|
}
|
|
|
|
func (d *changingHomeTargetDispatcher) HeartbeatOK() bool { return true }
|
|
func (d *changingHomeTargetDispatcher) RPopAuth(context.Context, string, string, http.Header, int) ([]byte, error) {
|
|
if d.calls.Add(1) == 2 && d.firstSelection != nil {
|
|
d.oldEndedBeforeRPop.Store(!d.firstSelection.Active())
|
|
}
|
|
return json.Marshal(homeAuthDispatchResponse{Auth: Auth{ID: "home-auth", Provider: "home-execution", Status: StatusActive, Attributes: map[string]string{"websockets": "true"}}})
|
|
}
|
|
func (*changingHomeTargetDispatcher) AbortAmbiguousDispatch() {}
|
|
|
|
type selectionRecordingExecutor struct {
|
|
first *HomeDispatchSelection
|
|
}
|
|
|
|
func (*selectionRecordingExecutor) Identifier() string { return "home-execution" }
|
|
func (e *selectionRecordingExecutor) Execute(_ context.Context, _ *Auth, _ cliproxyexecutor.Request, opts cliproxyexecutor.Options) (cliproxyexecutor.Response, error) {
|
|
selection, _ := opts.ExecutionLifecycle.(*HomeDispatchSelection)
|
|
if e.first == nil {
|
|
e.first = selection
|
|
}
|
|
if selection != nil {
|
|
selection.Retain()
|
|
}
|
|
return cliproxyexecutor.Response{Payload: []byte("ok")}, nil
|
|
}
|
|
func (*selectionRecordingExecutor) ExecuteStream(context.Context, *Auth, cliproxyexecutor.Request, cliproxyexecutor.Options) (*cliproxyexecutor.StreamResult, error) {
|
|
return nil, nil
|
|
}
|
|
func (*selectionRecordingExecutor) Refresh(context.Context, *Auth) (*Auth, error) { return nil, nil }
|
|
func (*selectionRecordingExecutor) CountTokens(context.Context, *Auth, cliproxyexecutor.Request, cliproxyexecutor.Options) (cliproxyexecutor.Response, error) {
|
|
return cliproxyexecutor.Response{}, nil
|
|
}
|
|
func (*selectionRecordingExecutor) HttpRequest(context.Context, *Auth, *http.Request) (*http.Response, error) {
|
|
return nil, nil
|
|
}
|
|
|
|
func TestHomeWebsocketTargetChangeEndsSelectionBeforeRedispatch(t *testing.T) {
|
|
dispatcher := &changingHomeTargetDispatcher{}
|
|
manager := NewManager(nil, nil, nil)
|
|
manager.SetConfig(&internalconfig.Config{Home: internalconfig.HomeConfig{Enabled: true}})
|
|
manager.PublishHomeDispatch(dispatcher, executionregistry.New(), 1)
|
|
executor := &selectionRecordingExecutor{}
|
|
manager.RegisterExecutor(executor)
|
|
ctx := cliproxyexecutor.WithDownstreamWebsocket(context.Background())
|
|
opts := cliproxyexecutor.Options{Metadata: map[string]any{
|
|
cliproxyexecutor.ExecutionSessionMetadataKey: "session-1",
|
|
cliproxyexecutor.PinnedAuthMetadataKey: "home-auth",
|
|
}}
|
|
|
|
if _, errExecute := manager.Execute(ctx, []string{"home-execution"}, cliproxyexecutor.Request{Model: "model-a"}, opts); errExecute != nil {
|
|
t.Fatalf("first Execute() error = %v", errExecute)
|
|
}
|
|
dispatcher.firstSelection = executor.first
|
|
if _, errExecute := manager.Execute(ctx, []string{"home-execution"}, cliproxyexecutor.Request{Model: "model-b"}, opts); errExecute != nil {
|
|
t.Fatalf("second Execute() error = %v", errExecute)
|
|
}
|
|
if got := dispatcher.calls.Load(); got != 2 {
|
|
t.Fatalf("Home RPOP calls = %d, want 2 after target change", got)
|
|
}
|
|
if !dispatcher.oldEndedBeforeRPop.Load() {
|
|
t.Fatal("previous selection remained active when target-change RPOP started")
|
|
}
|
|
}
|
|
|
|
type unpinnedTargetChangeDispatcher struct {
|
|
calls atomic.Int32
|
|
first *HomeDispatchSelection
|
|
oldClosedBeforeDispatch atomic.Bool
|
|
closeCalls *atomic.Int32
|
|
}
|
|
|
|
func (d *unpinnedTargetChangeDispatcher) HeartbeatOK() bool { return true }
|
|
func (d *unpinnedTargetChangeDispatcher) RPopAuth(_ context.Context, _ string, _ string, _ http.Header, _ int) ([]byte, error) {
|
|
call := d.calls.Add(1)
|
|
if call == 2 && d.first != nil {
|
|
d.oldClosedBeforeDispatch.Store(!d.first.Active() && d.closeCalls.Load() == 1)
|
|
}
|
|
return json.Marshal(homeAuthDispatchResponse{Auth: Auth{
|
|
ID: "home-auth-" + strconv.Itoa(int(call)),
|
|
Provider: "home-execution",
|
|
Status: StatusActive,
|
|
Attributes: map[string]string{
|
|
"websockets": "true",
|
|
},
|
|
}})
|
|
}
|
|
func (*unpinnedTargetChangeDispatcher) AbortAmbiguousDispatch() {}
|
|
|
|
type bindingSelectionRecordingExecutor struct {
|
|
first *HomeDispatchSelection
|
|
closeCalls *atomic.Int32
|
|
}
|
|
|
|
func (*bindingSelectionRecordingExecutor) Identifier() string { return "home-execution" }
|
|
func (e *bindingSelectionRecordingExecutor) Execute(_ context.Context, _ *Auth, _ cliproxyexecutor.Request, opts cliproxyexecutor.Options) (cliproxyexecutor.Response, error) {
|
|
selection, _ := opts.ExecutionLifecycle.(*HomeDispatchSelection)
|
|
if e.first == nil {
|
|
e.first = selection
|
|
}
|
|
if selection != nil {
|
|
if errBind := selection.Bind(func() error {
|
|
e.closeCalls.Add(1)
|
|
return nil
|
|
}); errBind != nil {
|
|
return cliproxyexecutor.Response{}, errBind
|
|
}
|
|
selection.Retain()
|
|
}
|
|
return cliproxyexecutor.Response{Payload: []byte("ok")}, nil
|
|
}
|
|
func (*bindingSelectionRecordingExecutor) ExecuteStream(context.Context, *Auth, cliproxyexecutor.Request, cliproxyexecutor.Options) (*cliproxyexecutor.StreamResult, error) {
|
|
return nil, nil
|
|
}
|
|
func (*bindingSelectionRecordingExecutor) Refresh(context.Context, *Auth) (*Auth, error) {
|
|
return nil, nil
|
|
}
|
|
func (*bindingSelectionRecordingExecutor) CountTokens(context.Context, *Auth, cliproxyexecutor.Request, cliproxyexecutor.Options) (cliproxyexecutor.Response, error) {
|
|
return cliproxyexecutor.Response{}, nil
|
|
}
|
|
func (*bindingSelectionRecordingExecutor) HttpRequest(context.Context, *Auth, *http.Request) (*http.Response, error) {
|
|
return nil, nil
|
|
}
|
|
|
|
func TestHomeWebsocketUnpinnedModelChangeClosesSelectionBeforeRedispatch(t *testing.T) {
|
|
var closeCalls atomic.Int32
|
|
dispatcher := &unpinnedTargetChangeDispatcher{closeCalls: &closeCalls}
|
|
registry := executionregistry.New()
|
|
manager := NewManager(nil, nil, nil)
|
|
manager.SetConfig(&internalconfig.Config{Home: internalconfig.HomeConfig{Enabled: true}})
|
|
manager.PublishHomeDispatch(dispatcher, registry, 1)
|
|
executor := &bindingSelectionRecordingExecutor{closeCalls: &closeCalls}
|
|
manager.RegisterExecutor(executor)
|
|
ctx := cliproxyexecutor.WithDownstreamWebsocket(context.Background())
|
|
opts := cliproxyexecutor.Options{Metadata: map[string]any{
|
|
cliproxyexecutor.ExecutionSessionMetadataKey: "session-1",
|
|
}}
|
|
|
|
if _, errExecute := manager.Execute(ctx, []string{"home-execution"}, cliproxyexecutor.Request{Model: "model-a"}, opts); errExecute != nil {
|
|
t.Fatalf("first Execute() error = %v", errExecute)
|
|
}
|
|
dispatcher.first = executor.first
|
|
if _, errExecute := manager.Execute(ctx, []string{"home-execution"}, cliproxyexecutor.Request{Model: "model-b"}, opts); errExecute != nil {
|
|
t.Fatalf("second Execute() error = %v", errExecute)
|
|
}
|
|
if got := dispatcher.calls.Load(); got != 2 {
|
|
t.Fatalf("Home RPOP calls = %d, want 2", got)
|
|
}
|
|
if !dispatcher.oldClosedBeforeDispatch.Load() {
|
|
t.Fatal("old unpinned selection was not ended and closed before the second RPOP")
|
|
}
|
|
manager.CloseExecutionSession("session-1")
|
|
if errDrain := registry.Drain(context.Background()); errDrain != nil {
|
|
t.Fatalf("Drain() error = %v", errDrain)
|
|
}
|
|
}
|
|
|
|
type lifecycleRetryDispatcher struct {
|
|
calls atomic.Int32
|
|
executor *lifecycleRetryExecutor
|
|
firstEndedBeforeRedispatch atomic.Bool
|
|
}
|
|
|
|
func (d *lifecycleRetryDispatcher) HeartbeatOK() bool { return true }
|
|
func (d *lifecycleRetryDispatcher) RPopAuth(ctx context.Context, model string, sessionID string, headers http.Header, count int) ([]byte, error) {
|
|
return d.RPopAuthWithConstraints(ctx, model, sessionID, headers, count, nil, "")
|
|
}
|
|
func (d *lifecycleRetryDispatcher) RPopAuthWithConstraints(_ context.Context, _ string, _ string, _ http.Header, _ int, excludedAuthIDs []string, _ string) ([]byte, error) {
|
|
for _, authID := range excludedAuthIDs {
|
|
if authID == "home-auth" {
|
|
return nil, home.ErrAuthNotFound
|
|
}
|
|
}
|
|
if d.calls.Add(1) == 2 && d.executor.first != nil {
|
|
d.firstEndedBeforeRedispatch.Store(!d.executor.first.Active() && d.executor.firstCtx.Err() != nil)
|
|
}
|
|
return json.Marshal(homeAuthDispatchResponse{Auth: Auth{ID: "home-auth", Provider: "home-execution", Status: StatusActive, Attributes: map[string]string{"websockets": "true"}}})
|
|
}
|
|
func (*lifecycleRetryDispatcher) AbortAmbiguousDispatch() {}
|
|
|
|
type lifecycleRetryExecutor struct {
|
|
calls atomic.Int32
|
|
first *HomeDispatchSelection
|
|
firstCtx context.Context
|
|
}
|
|
|
|
func (*lifecycleRetryExecutor) Identifier() string { return "home-execution" }
|
|
func (*lifecycleRetryExecutor) Execute(context.Context, *Auth, cliproxyexecutor.Request, cliproxyexecutor.Options) (cliproxyexecutor.Response, error) {
|
|
return cliproxyexecutor.Response{}, nil
|
|
}
|
|
func (e *lifecycleRetryExecutor) ExecuteStream(ctx context.Context, _ *Auth, _ cliproxyexecutor.Request, opts cliproxyexecutor.Options) (*cliproxyexecutor.StreamResult, error) {
|
|
if e.calls.Add(1) == 1 {
|
|
e.first, _ = opts.ExecutionLifecycle.(*HomeDispatchSelection)
|
|
e.firstCtx = ctx
|
|
return nil, &Error{HTTPStatus: http.StatusUpgradeRequired, Message: "websocket upgrade required"}
|
|
}
|
|
chunks := make(chan cliproxyexecutor.StreamChunk, 1)
|
|
chunks <- cliproxyexecutor.StreamChunk{Payload: []byte(`{"type":"response.completed"}`)}
|
|
close(chunks)
|
|
return &cliproxyexecutor.StreamResult{Chunks: chunks}, nil
|
|
}
|
|
func (*lifecycleRetryExecutor) Refresh(context.Context, *Auth) (*Auth, error) { return nil, nil }
|
|
func (*lifecycleRetryExecutor) CountTokens(context.Context, *Auth, cliproxyexecutor.Request, cliproxyexecutor.Options) (cliproxyexecutor.Response, error) {
|
|
return cliproxyexecutor.Response{}, nil
|
|
}
|
|
func (*lifecycleRetryExecutor) HttpRequest(context.Context, *Auth, *http.Request) (*http.Response, error) {
|
|
return nil, nil
|
|
}
|
|
|
|
func TestHomeStreamLifecycleFailureEndsBeforeFreshDispatch(t *testing.T) {
|
|
executor := &lifecycleRetryExecutor{}
|
|
dispatcher := &lifecycleRetryDispatcher{executor: executor}
|
|
registry := executionregistry.New()
|
|
manager := NewManager(nil, nil, nil)
|
|
manager.SetConfig(&internalconfig.Config{Home: internalconfig.HomeConfig{Enabled: true}})
|
|
manager.SetRetryConfig(0, time.Second, 1)
|
|
manager.PublishHomeDispatch(dispatcher, registry, 1)
|
|
manager.RegisterExecutor(executor)
|
|
ctx := cliproxyexecutor.WithDownstreamWebsocket(context.Background())
|
|
opts := cliproxyexecutor.Options{Stream: true, Metadata: map[string]any{
|
|
cliproxyexecutor.ExecutionSessionMetadataKey: "session-426",
|
|
}}
|
|
|
|
result, errExecute := manager.ExecuteStream(ctx, []string{"home-execution"}, cliproxyexecutor.Request{Model: "model-a"}, opts)
|
|
if errExecute != nil {
|
|
t.Fatalf("ExecuteStream() error = %v", errExecute)
|
|
}
|
|
for range result.Chunks {
|
|
}
|
|
if got := executor.calls.Load(); got != 2 {
|
|
t.Fatalf("executor invocations = %d, want 2", got)
|
|
}
|
|
if got := dispatcher.calls.Load(); got != 2 {
|
|
t.Fatalf("Home RPOP calls = %d, want 2", got)
|
|
}
|
|
if !dispatcher.firstEndedBeforeRedispatch.Load() {
|
|
t.Fatal("failed stream attempt remained active when the fresh Home selection was dispatched")
|
|
}
|
|
if errDrain := registry.Drain(context.Background()); errDrain != nil {
|
|
t.Fatalf("Drain() error = %v", errDrain)
|
|
}
|
|
}
|
|
|
|
func TestHomeSelectionCancellationPreventsExecute(t *testing.T) {
|
|
manager := NewManager(nil, nil, nil)
|
|
manager.SetConfig(&internalconfig.Config{Home: internalconfig.HomeConfig{Enabled: true}})
|
|
manager.PublishHomeDispatch(homeExecutionDispatcher{}, executionregistry.New(), 1)
|
|
executor := &homeExecutionExecutor{}
|
|
manager.RegisterExecutor(executor)
|
|
|
|
ctx, cancel := context.WithCancel(context.Background())
|
|
cancel()
|
|
_, errExecute := manager.Execute(ctx, []string{"home-execution"}, cliproxyexecutor.Request{Model: "test"}, cliproxyexecutor.Options{})
|
|
if errExecute == nil {
|
|
t.Fatal("Execute() error = nil, want canceled context")
|
|
}
|
|
if executor.ctx != nil {
|
|
t.Fatal("executor was invoked after attempt context cancellation")
|
|
}
|
|
}
|
|
|
|
type freshHomeStreamSelectionDispatcher struct {
|
|
calls atomic.Int32
|
|
}
|
|
|
|
func (*freshHomeStreamSelectionDispatcher) HeartbeatOK() bool { return true }
|
|
|
|
func (d *freshHomeStreamSelectionDispatcher) RPopAuth(ctx context.Context, model string, sessionID string, headers http.Header, count int) ([]byte, error) {
|
|
return d.RPopAuthWithConstraints(ctx, model, sessionID, headers, count, nil, "")
|
|
}
|
|
|
|
func (d *freshHomeStreamSelectionDispatcher) RPopAuthWithConstraints(_ context.Context, _ string, _ string, _ http.Header, _ int, excludedAuthIDs []string, _ string) ([]byte, error) {
|
|
d.calls.Add(1)
|
|
excluded := make(map[string]struct{}, len(excludedAuthIDs))
|
|
for _, authID := range excludedAuthIDs {
|
|
excluded[authID] = struct{}{}
|
|
}
|
|
for _, authID := range []string{"home-auth-a", "home-auth-b"} {
|
|
if _, okExcluded := excluded[authID]; okExcluded {
|
|
continue
|
|
}
|
|
return json.Marshal(homeAuthDispatchResponse{Auth: Auth{
|
|
ID: authID,
|
|
Provider: "home-execution",
|
|
Status: StatusActive,
|
|
Attributes: map[string]string{
|
|
AttributeAuthKind: AuthKindAPIKey,
|
|
},
|
|
}})
|
|
}
|
|
return nil, home.ErrAuthNotFound
|
|
}
|
|
|
|
func (*freshHomeStreamSelectionDispatcher) AbortAmbiguousDispatch() {}
|
|
|
|
type retryingHomeStreamExecutor struct {
|
|
mu sync.Mutex
|
|
calls atomic.Int32
|
|
authIDs []string
|
|
}
|
|
|
|
func (*retryingHomeStreamExecutor) Identifier() string { return "home-execution" }
|
|
func (*retryingHomeStreamExecutor) Execute(context.Context, *Auth, cliproxyexecutor.Request, cliproxyexecutor.Options) (cliproxyexecutor.Response, error) {
|
|
return cliproxyexecutor.Response{}, nil
|
|
}
|
|
func (e *retryingHomeStreamExecutor) ExecuteStream(_ context.Context, auth *Auth, _ cliproxyexecutor.Request, _ cliproxyexecutor.Options) (*cliproxyexecutor.StreamResult, error) {
|
|
e.mu.Lock()
|
|
e.authIDs = append(e.authIDs, auth.ID)
|
|
e.mu.Unlock()
|
|
if e.calls.Add(1) == 1 {
|
|
return nil, &Error{HTTPStatus: http.StatusUnauthorized, Message: "expired"}
|
|
}
|
|
chunks := make(chan cliproxyexecutor.StreamChunk, 1)
|
|
chunks <- cliproxyexecutor.StreamChunk{Payload: []byte("data: {\"type\":\"response.completed\"}\n\n")}
|
|
close(chunks)
|
|
return &cliproxyexecutor.StreamResult{Chunks: chunks}, nil
|
|
}
|
|
func (*retryingHomeStreamExecutor) Refresh(_ context.Context, auth *Auth) (*Auth, error) {
|
|
return auth, nil
|
|
}
|
|
func (*retryingHomeStreamExecutor) CountTokens(context.Context, *Auth, cliproxyexecutor.Request, cliproxyexecutor.Options) (cliproxyexecutor.Response, error) {
|
|
return cliproxyexecutor.Response{}, nil
|
|
}
|
|
func (*retryingHomeStreamExecutor) HttpRequest(context.Context, *Auth, *http.Request) (*http.Response, error) {
|
|
return nil, nil
|
|
}
|
|
|
|
func (e *retryingHomeStreamExecutor) AuthIDs() []string {
|
|
e.mu.Lock()
|
|
defer e.mu.Unlock()
|
|
return append([]string(nil), e.authIDs...)
|
|
}
|
|
|
|
func TestHomeStreamRetryUsesFreshSelection(t *testing.T) {
|
|
dispatcher := &freshHomeStreamSelectionDispatcher{}
|
|
manager := NewManager(nil, nil, nil)
|
|
manager.SetConfig(&internalconfig.Config{Home: internalconfig.HomeConfig{Enabled: true}})
|
|
manager.SetRetryConfig(0, time.Second, 2)
|
|
manager.PublishHomeDispatch(dispatcher, executionregistry.New(), 1)
|
|
executor := &retryingHomeStreamExecutor{}
|
|
manager.RegisterExecutor(executor)
|
|
|
|
result, errExecute := manager.ExecuteStream(context.Background(), []string{"home-execution"}, cliproxyexecutor.Request{Model: "model-a"}, cliproxyexecutor.Options{Stream: true})
|
|
if errExecute != nil {
|
|
t.Fatalf("ExecuteStream() error = %v", errExecute)
|
|
}
|
|
for range result.Chunks {
|
|
}
|
|
if got := dispatcher.calls.Load(); got != 2 {
|
|
t.Fatalf("Home RPOP calls = %d, want 2 for retrying stream invocations", got)
|
|
}
|
|
if got := executor.AuthIDs(); len(got) != 2 || got[0] != "home-auth-a" || got[1] != "home-auth-b" {
|
|
t.Fatalf("executor auth IDs = %v, want [home-auth-a home-auth-b]", got)
|
|
}
|
|
}
|
|
|
|
type cancellationBarrierExecutor struct {
|
|
executeCalls atomic.Int32
|
|
countCalls atomic.Int32
|
|
streamCalls atomic.Int32
|
|
}
|
|
|
|
func (*cancellationBarrierExecutor) Identifier() string { return "home-execution" }
|
|
func (e *cancellationBarrierExecutor) Execute(context.Context, *Auth, cliproxyexecutor.Request, cliproxyexecutor.Options) (cliproxyexecutor.Response, error) {
|
|
e.executeCalls.Add(1)
|
|
return cliproxyexecutor.Response{}, nil
|
|
}
|
|
func (e *cancellationBarrierExecutor) CountTokens(context.Context, *Auth, cliproxyexecutor.Request, cliproxyexecutor.Options) (cliproxyexecutor.Response, error) {
|
|
e.countCalls.Add(1)
|
|
return cliproxyexecutor.Response{}, nil
|
|
}
|
|
func (e *cancellationBarrierExecutor) ExecuteStream(context.Context, *Auth, cliproxyexecutor.Request, cliproxyexecutor.Options) (*cliproxyexecutor.StreamResult, error) {
|
|
e.streamCalls.Add(1)
|
|
return nil, nil
|
|
}
|
|
func (*cancellationBarrierExecutor) Refresh(context.Context, *Auth) (*Auth, error) { return nil, nil }
|
|
func (*cancellationBarrierExecutor) HttpRequest(context.Context, *Auth, *http.Request) (*http.Response, error) {
|
|
return nil, nil
|
|
}
|
|
|
|
func TestHomeCancellationBarrierPreventsEveryExecutorInvocation(t *testing.T) {
|
|
manager := NewManager(nil, nil, nil)
|
|
manager.SetConfig(&internalconfig.Config{Home: internalconfig.HomeConfig{Enabled: true}})
|
|
manager.PublishHomeDispatch(homeExecutionDispatcher{}, executionregistry.New(), 1)
|
|
executor := &cancellationBarrierExecutor{}
|
|
manager.RegisterExecutor(executor)
|
|
ctx, cancel := context.WithCancel(context.Background())
|
|
cancel()
|
|
|
|
if _, errExecute := manager.Execute(ctx, []string{"home-execution"}, cliproxyexecutor.Request{Model: "test"}, cliproxyexecutor.Options{}); errExecute == nil {
|
|
t.Fatal("Execute() error = nil, want canceled context")
|
|
}
|
|
if _, errCount := manager.ExecuteCount(ctx, []string{"home-execution"}, cliproxyexecutor.Request{Model: "test"}, cliproxyexecutor.Options{}); errCount == nil {
|
|
t.Fatal("ExecuteCount() error = nil, want canceled context")
|
|
}
|
|
if _, errStream := manager.ExecuteStream(ctx, []string{"home-execution"}, cliproxyexecutor.Request{Model: "test"}, cliproxyexecutor.Options{Stream: true}); errStream == nil {
|
|
t.Fatal("ExecuteStream() error = nil, want canceled context")
|
|
}
|
|
if got := executor.executeCalls.Load(); got != 0 {
|
|
t.Fatalf("Execute calls = %d, want 0", got)
|
|
}
|
|
if got := executor.countCalls.Load(); got != 0 {
|
|
t.Fatalf("CountTokens calls = %d, want 0", got)
|
|
}
|
|
if got := executor.streamCalls.Load(); got != 0 {
|
|
t.Fatalf("ExecuteStream calls = %d, want 0", got)
|
|
}
|
|
}
|
|
|
|
func TestHomeStreamEndsOnTerminalChunk(t *testing.T) {
|
|
manager := NewManager(nil, nil, nil)
|
|
manager.SetConfig(&internalconfig.Config{Home: internalconfig.HomeConfig{Enabled: true}})
|
|
registry := executionregistry.New()
|
|
manager.PublishHomeDispatch(homeExecutionDispatcher{}, registry, 1)
|
|
chunks := make(chan cliproxyexecutor.StreamChunk, 1)
|
|
chunks <- cliproxyexecutor.StreamChunk{Payload: []byte("initial")}
|
|
manager.RegisterExecutor(&homeExecutionStreamExecutor{chunks: chunks})
|
|
|
|
result, errExecute := manager.ExecuteStream(context.Background(), []string{"home-execution"}, cliproxyexecutor.Request{Model: "test"}, cliproxyexecutor.Options{Stream: true})
|
|
if errExecute != nil {
|
|
t.Fatalf("ExecuteStream() error = %v", errExecute)
|
|
}
|
|
|
|
close(chunks)
|
|
for range result.Chunks {
|
|
}
|
|
if errDrain := registry.Drain(context.Background()); errDrain != nil {
|
|
t.Fatalf("Drain() error = %v", errDrain)
|
|
}
|
|
}
|
|
|
|
func TestHomeWebsocketSessionReusesSelectionWithoutPinnedMetadataAndCachesRuntimeAuth(t *testing.T) {
|
|
dispatcher := &retainingHomeExecutionDispatcher{}
|
|
manager := NewManager(nil, nil, nil)
|
|
manager.SetConfig(&internalconfig.Config{Home: internalconfig.HomeConfig{Enabled: true}})
|
|
manager.PublishHomeDispatch(dispatcher, executionregistry.New(), 1)
|
|
executor := &retainingHomeExecutionExecutor{}
|
|
manager.RegisterExecutor(executor)
|
|
|
|
ctx := cliproxyexecutor.WithDownstreamWebsocket(context.Background())
|
|
opts := cliproxyexecutor.Options{Metadata: map[string]any{
|
|
cliproxyexecutor.ExecutionSessionMetadataKey: "session-without-pin",
|
|
}}
|
|
for range 2 {
|
|
if _, errExecute := manager.Execute(ctx, []string{"home-execution"}, cliproxyexecutor.Request{Model: "model-a"}, opts); errExecute != nil {
|
|
t.Fatalf("Execute() error = %v", errExecute)
|
|
}
|
|
}
|
|
if got := dispatcher.calls.Load(); got != 1 {
|
|
t.Fatalf("Home RPOP calls = %d, want 1 for a retained session without a pin", got)
|
|
}
|
|
if auth, ok := manager.GetExecutionSessionAuthByID("session-without-pin", "home-auth"); !ok || auth == nil {
|
|
t.Fatal("retained selection did not populate the handler runtime auth cache")
|
|
}
|
|
}
|
|
|
|
func TestCloseExecutionSessionReclaimsHomeSessionLock(t *testing.T) {
|
|
manager := NewManager(nil, nil, nil)
|
|
ctx := cliproxyexecutor.WithDownstreamWebsocket(context.Background())
|
|
opts := cliproxyexecutor.Options{Metadata: map[string]any{
|
|
cliproxyexecutor.ExecutionSessionMetadataKey: "reclaim-lock",
|
|
}}
|
|
unlock := manager.lockHomeWebsocketSession(ctx, opts)
|
|
if unlock == nil {
|
|
t.Fatal("lockHomeWebsocketSession() = nil")
|
|
}
|
|
unlock()
|
|
if _, ok := manager.homeSessionLocks.Load("reclaim-lock"); !ok {
|
|
t.Fatal("session lock was not created")
|
|
}
|
|
|
|
manager.CloseExecutionSession("reclaim-lock")
|
|
if _, ok := manager.homeSessionLocks.Load("reclaim-lock"); ok {
|
|
t.Fatal("closed session retained its mutex entry")
|
|
}
|
|
}
|
|
|
|
type homePerSelectionDispatcher struct {
|
|
auths []Auth
|
|
calls atomic.Int32
|
|
first *HomeDispatchSelection
|
|
firstEndedBefore2 atomic.Bool
|
|
}
|
|
|
|
func (*homePerSelectionDispatcher) HeartbeatOK() bool { return true }
|
|
func (d *homePerSelectionDispatcher) RPopAuth(context.Context, string, string, http.Header, int) ([]byte, error) {
|
|
call := d.calls.Add(1)
|
|
if call == 2 && d.first != nil {
|
|
d.firstEndedBefore2.Store(!d.first.Active())
|
|
}
|
|
if int(call) > len(d.auths) {
|
|
return nil, home.ErrAuthNotFound
|
|
}
|
|
return json.Marshal(homeAuthDispatchResponse{Auth: d.auths[call-1]})
|
|
}
|
|
func (*homePerSelectionDispatcher) AbortAmbiguousDispatch() {}
|
|
|
|
type homePerSelectionFailureExecutor struct {
|
|
dispatcher *homePerSelectionDispatcher
|
|
selections []*HomeDispatchSelection
|
|
invocations []string
|
|
}
|
|
|
|
func (*homePerSelectionFailureExecutor) Identifier() string { return openAICompatPoolProviderKey }
|
|
func (e *homePerSelectionFailureExecutor) invoke(auth *Auth, opts cliproxyexecutor.Options) (cliproxyexecutor.Response, error) {
|
|
selection, _ := opts.ExecutionLifecycle.(*HomeDispatchSelection)
|
|
if e.selections == nil {
|
|
e.selections = append(e.selections, selection)
|
|
}
|
|
if selection != nil && len(e.selections) == 1 {
|
|
e.selections[0] = selection
|
|
if e.dispatcher != nil {
|
|
e.dispatcher.first = selection
|
|
}
|
|
}
|
|
e.invocations = append(e.invocations, auth.ID)
|
|
return cliproxyexecutor.Response{}, &Error{HTTPStatus: http.StatusBadGateway, Message: "upstream failed"}
|
|
}
|
|
func (e *homePerSelectionFailureExecutor) Execute(_ context.Context, auth *Auth, _ cliproxyexecutor.Request, opts cliproxyexecutor.Options) (cliproxyexecutor.Response, error) {
|
|
return e.invoke(auth, opts)
|
|
}
|
|
func (*homePerSelectionFailureExecutor) ExecuteStream(context.Context, *Auth, cliproxyexecutor.Request, cliproxyexecutor.Options) (*cliproxyexecutor.StreamResult, error) {
|
|
return nil, nil
|
|
}
|
|
func (*homePerSelectionFailureExecutor) Refresh(context.Context, *Auth) (*Auth, error) {
|
|
return nil, nil
|
|
}
|
|
func (e *homePerSelectionFailureExecutor) CountTokens(_ context.Context, auth *Auth, _ cliproxyexecutor.Request, opts cliproxyexecutor.Options) (cliproxyexecutor.Response, error) {
|
|
return e.invoke(auth, opts)
|
|
}
|
|
func (*homePerSelectionFailureExecutor) HttpRequest(context.Context, *Auth, *http.Request) (*http.Response, error) {
|
|
return nil, nil
|
|
}
|
|
|
|
func TestHomeNonstreamAndCountUseOneModelPerSelection(t *testing.T) {
|
|
for _, countTokens := range []bool{false, true} {
|
|
t.Run(map[bool]string{false: "Execute", true: "CountTokens"}[countTokens], func(t *testing.T) {
|
|
dispatcher := &homePerSelectionDispatcher{auths: []Auth{
|
|
{ID: "home-auth-a", Provider: "home-pool", Status: StatusActive, Attributes: map[string]string{"api_key": "test-key", "compat_name": "pool", "provider_key": "pool"}},
|
|
{ID: "home-auth-b", Provider: "home-pool", Status: StatusActive, Attributes: map[string]string{"api_key": "test-key", "compat_name": "pool", "provider_key": "pool"}},
|
|
}}
|
|
manager := NewManager(nil, nil, nil)
|
|
manager.SetConfig(&internalconfig.Config{
|
|
Home: internalconfig.HomeConfig{Enabled: true},
|
|
OpenAICompatibility: []internalconfig.OpenAICompatibility{{
|
|
Name: "pool",
|
|
Models: []internalconfig.OpenAICompatibilityModel{{Name: "upstream-a", Alias: "requested"}, {Name: "upstream-b", Alias: "requested"}},
|
|
}},
|
|
})
|
|
manager.PublishHomeDispatch(dispatcher, executionregistry.New(), 1)
|
|
executor := &homePerSelectionFailureExecutor{dispatcher: dispatcher}
|
|
manager.RegisterExecutor(executor)
|
|
|
|
var errExecute error
|
|
if countTokens {
|
|
_, errExecute = manager.ExecuteCount(context.Background(), []string{openAICompatPoolProviderKey}, cliproxyexecutor.Request{Model: "requested"}, cliproxyexecutor.Options{})
|
|
} else {
|
|
_, errExecute = manager.Execute(context.Background(), []string{openAICompatPoolProviderKey}, cliproxyexecutor.Request{Model: "requested"}, cliproxyexecutor.Options{})
|
|
}
|
|
if errExecute == nil {
|
|
t.Fatal("execution error = nil, want upstream failure")
|
|
}
|
|
if len(executor.invocations) != 2 {
|
|
t.Fatalf("execution error = %v; upstream invocations = %v, want one per Home selection", errExecute, executor.invocations)
|
|
}
|
|
if !dispatcher.firstEndedBefore2.Load() {
|
|
t.Fatal("first Home selection was not ended before the next dispatch")
|
|
}
|
|
})
|
|
}
|
|
}
|
|
|
|
func TestHomeStreamEndsOnErrorChunk(t *testing.T) {
|
|
manager := NewManager(nil, nil, nil)
|
|
manager.SetConfig(&internalconfig.Config{Home: internalconfig.HomeConfig{Enabled: true}})
|
|
registry := executionregistry.New()
|
|
manager.PublishHomeDispatch(homeExecutionDispatcher{}, registry, 1)
|
|
chunks := make(chan cliproxyexecutor.StreamChunk, 2)
|
|
chunks <- cliproxyexecutor.StreamChunk{Payload: []byte("initial")}
|
|
chunks <- cliproxyexecutor.StreamChunk{Err: &Error{HTTPStatus: http.StatusBadGateway, Message: "upstream failed"}}
|
|
close(chunks)
|
|
manager.RegisterExecutor(&homeExecutionStreamExecutor{chunks: chunks})
|
|
|
|
result, errExecute := manager.ExecuteStream(context.Background(), []string{"home-execution"}, cliproxyexecutor.Request{Model: "test"}, cliproxyexecutor.Options{Stream: true})
|
|
if errExecute != nil {
|
|
t.Fatalf("ExecuteStream() error = %v", errExecute)
|
|
}
|
|
sawError := false
|
|
for chunk := range result.Chunks {
|
|
if chunk.Err != nil {
|
|
sawError = true
|
|
}
|
|
}
|
|
if !sawError {
|
|
t.Fatal("stream did not preserve the upstream error chunk")
|
|
}
|
|
if errDrain := registry.Drain(context.Background()); errDrain != nil {
|
|
t.Fatalf("Drain() error = %v", errDrain)
|
|
}
|
|
}
|
|
|
|
type missingHomeStreamSourceExecutor struct{}
|
|
|
|
func (*missingHomeStreamSourceExecutor) Identifier() string { return "home-execution" }
|
|
func (*missingHomeStreamSourceExecutor) Execute(context.Context, *Auth, cliproxyexecutor.Request, cliproxyexecutor.Options) (cliproxyexecutor.Response, error) {
|
|
return cliproxyexecutor.Response{}, nil
|
|
}
|
|
func (*missingHomeStreamSourceExecutor) ExecuteStream(context.Context, *Auth, cliproxyexecutor.Request, cliproxyexecutor.Options) (*cliproxyexecutor.StreamResult, error) {
|
|
return nil, nil
|
|
}
|
|
func (*missingHomeStreamSourceExecutor) Refresh(context.Context, *Auth) (*Auth, error) {
|
|
return nil, nil
|
|
}
|
|
func (*missingHomeStreamSourceExecutor) CountTokens(context.Context, *Auth, cliproxyexecutor.Request, cliproxyexecutor.Options) (cliproxyexecutor.Response, error) {
|
|
return cliproxyexecutor.Response{}, nil
|
|
}
|
|
func (*missingHomeStreamSourceExecutor) HttpRequest(context.Context, *Auth, *http.Request) (*http.Response, error) {
|
|
return nil, nil
|
|
}
|
|
|
|
type accountedHomeExecutionDispatcher struct {
|
|
calls atomic.Int32
|
|
auths []Auth
|
|
}
|
|
|
|
func (*accountedHomeExecutionDispatcher) HeartbeatOK() bool { return true }
|
|
func (d *accountedHomeExecutionDispatcher) RPopAuth(_ context.Context, model string, _ string, _ http.Header, _ int) ([]byte, error) {
|
|
index := int(d.calls.Add(1)) - 1
|
|
if index >= len(d.auths) {
|
|
return nil, home.ErrAuthNotFound
|
|
}
|
|
auth := d.auths[index]
|
|
return json.Marshal(struct {
|
|
Concurrency homeConcurrencyTuple `json:"concurrency"`
|
|
Model string `json:"model"`
|
|
AuthIndex string `json:"auth_index"`
|
|
Auth Auth `json:"auth"`
|
|
}{
|
|
Concurrency: homeConcurrencyTuple{Accounted: true, CredentialID: auth.ID, Model: model},
|
|
Model: model,
|
|
AuthIndex: auth.ID,
|
|
Auth: auth,
|
|
})
|
|
}
|
|
func (*accountedHomeExecutionDispatcher) AbortAmbiguousDispatch() {}
|
|
|
|
func TestAccountedHomeExecuteAndCountReleaseOnce(t *testing.T) {
|
|
for _, countTokens := range []bool{false, true} {
|
|
t.Run(map[bool]string{false: "Execute", true: "Count"}[countTokens], func(t *testing.T) {
|
|
manager := NewManager(nil, nil, nil)
|
|
manager.SetConfig(&internalconfig.Config{Home: internalconfig.HomeConfig{Enabled: true}})
|
|
registry := executionregistry.New()
|
|
releases := make(chan executionregistry.ReleaseGroup, 2)
|
|
registry.SetReleaseSink(func(group executionregistry.ReleaseGroup, _ int64) { releases <- group })
|
|
manager.PublishHomeDispatch(&accountedHomeExecutionDispatcher{auths: []Auth{{
|
|
ID: "cred-1", Provider: "home-execution", Status: StatusActive,
|
|
}}}, registry, 1)
|
|
manager.RegisterExecutor(&homeExecutionExecutor{})
|
|
|
|
var errExecute error
|
|
if countTokens {
|
|
_, errExecute = manager.ExecuteCount(context.Background(), []string{"home-execution"}, cliproxyexecutor.Request{Model: "model-a"}, cliproxyexecutor.Options{})
|
|
} else {
|
|
_, errExecute = manager.Execute(context.Background(), []string{"home-execution"}, cliproxyexecutor.Request{Model: "model-a"}, cliproxyexecutor.Options{})
|
|
}
|
|
if errExecute != nil {
|
|
t.Fatalf("execution error = %v", errExecute)
|
|
}
|
|
select {
|
|
case group := <-releases:
|
|
if group != (executionregistry.ReleaseGroup{CredentialID: "cred-1", Model: "model-a"}) {
|
|
t.Fatalf("release group = %#v", group)
|
|
}
|
|
default:
|
|
t.Fatal("accounted selection did not release")
|
|
}
|
|
select {
|
|
case group := <-releases:
|
|
t.Fatalf("duplicate release = %#v", group)
|
|
default:
|
|
}
|
|
})
|
|
}
|
|
}
|
|
|
|
func TestAccountedHomeStreamEndsOnlyAfterSourceTerminates(t *testing.T) {
|
|
manager := NewManager(nil, nil, nil)
|
|
manager.SetConfig(&internalconfig.Config{Home: internalconfig.HomeConfig{Enabled: true}})
|
|
registry := executionregistry.New()
|
|
releases := make(chan executionregistry.ReleaseGroup, 1)
|
|
registry.SetReleaseSink(func(group executionregistry.ReleaseGroup, _ int64) { releases <- group })
|
|
manager.PublishHomeDispatch(&accountedHomeExecutionDispatcher{auths: []Auth{{
|
|
ID: "cred-1", Provider: "home-execution", Status: StatusActive,
|
|
}}}, registry, 1)
|
|
chunks := make(chan cliproxyexecutor.StreamChunk, 1)
|
|
chunks <- cliproxyexecutor.StreamChunk{Payload: []byte("initial")}
|
|
manager.RegisterExecutor(&homeExecutionStreamExecutor{chunks: chunks})
|
|
|
|
result, errExecute := manager.ExecuteStream(context.Background(), []string{"home-execution"}, cliproxyexecutor.Request{Model: "model-a"}, cliproxyexecutor.Options{Stream: true})
|
|
if errExecute != nil {
|
|
t.Fatalf("ExecuteStream() error = %v", errExecute)
|
|
}
|
|
if _, ok := <-result.Chunks; !ok {
|
|
t.Fatal("stream closed before initial chunk")
|
|
}
|
|
select {
|
|
case group := <-releases:
|
|
t.Fatalf("stream released before source termination: %#v", group)
|
|
default:
|
|
}
|
|
|
|
close(chunks)
|
|
for range result.Chunks {
|
|
}
|
|
select {
|
|
case group := <-releases:
|
|
if group != (executionregistry.ReleaseGroup{CredentialID: "cred-1", Model: "model-a"}) {
|
|
t.Fatalf("release group = %#v", group)
|
|
}
|
|
case <-time.After(time.Second):
|
|
t.Fatal("stream did not release after source termination")
|
|
}
|
|
}
|
|
|
|
func TestAccountedHomeStreamErrorDrainsUntilSourceClosesBeforeRelease(t *testing.T) {
|
|
manager := NewManager(nil, nil, nil)
|
|
manager.SetConfig(&internalconfig.Config{Home: internalconfig.HomeConfig{Enabled: true}})
|
|
registry := executionregistry.New()
|
|
releases := make(chan executionregistry.ReleaseGroup, 1)
|
|
registry.SetReleaseSink(func(group executionregistry.ReleaseGroup, _ int64) { releases <- group })
|
|
manager.PublishHomeDispatch(&accountedHomeExecutionDispatcher{auths: []Auth{{
|
|
ID: "cred-1", Provider: "home-execution", Status: StatusActive,
|
|
}}}, registry, 1)
|
|
chunks := make(chan cliproxyexecutor.StreamChunk, 1)
|
|
chunks <- cliproxyexecutor.StreamChunk{Payload: []byte("initial")}
|
|
manager.RegisterExecutor(&homeExecutionStreamExecutor{chunks: chunks})
|
|
|
|
result, errExecute := manager.ExecuteStream(context.Background(), []string{"home-execution"}, cliproxyexecutor.Request{Model: "model-a"}, cliproxyexecutor.Options{Stream: true})
|
|
if errExecute != nil {
|
|
t.Fatalf("ExecuteStream() error = %v", errExecute)
|
|
}
|
|
if chunk, ok := <-result.Chunks; !ok || string(chunk.Payload) != "initial" {
|
|
t.Fatalf("initial chunk = %#v, open = %v", chunk, ok)
|
|
}
|
|
chunks <- cliproxyexecutor.StreamChunk{Err: &Error{HTTPStatus: http.StatusBadGateway, Message: "upstream failed"}}
|
|
if chunk, ok := <-result.Chunks; !ok || chunk.Err == nil {
|
|
t.Fatalf("error chunk = %#v, open = %v", chunk, ok)
|
|
}
|
|
|
|
sent := make(chan struct{})
|
|
go func() {
|
|
chunks <- cliproxyexecutor.StreamChunk{Payload: []byte("after-error-1")}
|
|
chunks <- cliproxyexecutor.StreamChunk{Payload: []byte("after-error-2")}
|
|
close(sent)
|
|
}()
|
|
select {
|
|
case <-sent:
|
|
case <-time.After(time.Second):
|
|
t.Fatal("stream source was not drained after its error chunk")
|
|
}
|
|
select {
|
|
case group := <-releases:
|
|
t.Fatalf("stream released while source remained open: %#v", group)
|
|
default:
|
|
}
|
|
select {
|
|
case chunk, ok := <-result.Chunks:
|
|
t.Fatalf("chunk after error = %#v, open = %v", chunk, ok)
|
|
case <-time.After(50 * time.Millisecond):
|
|
}
|
|
|
|
close(chunks)
|
|
for range result.Chunks {
|
|
}
|
|
select {
|
|
case group := <-releases:
|
|
if group != (executionregistry.ReleaseGroup{CredentialID: "cred-1", Model: "model-a"}) {
|
|
t.Fatalf("release group = %#v", group)
|
|
}
|
|
case <-time.After(time.Second):
|
|
t.Fatal("stream did not release after the source closed")
|
|
}
|
|
}
|
|
|
|
func TestAccountedHomeStreamErrorCancellationReleasesSelection(t *testing.T) {
|
|
ctx, cancel := context.WithCancel(context.Background())
|
|
manager := NewManager(nil, nil, nil)
|
|
manager.SetConfig(&internalconfig.Config{Home: internalconfig.HomeConfig{Enabled: true}})
|
|
registry := executionregistry.New()
|
|
releases := make(chan executionregistry.ReleaseGroup, 1)
|
|
registry.SetReleaseSink(func(group executionregistry.ReleaseGroup, _ int64) { releases <- group })
|
|
manager.PublishHomeDispatch(&accountedHomeExecutionDispatcher{auths: []Auth{{
|
|
ID: "cred-1", Provider: "home-execution", Status: StatusActive,
|
|
}}}, registry, 1)
|
|
chunks := make(chan cliproxyexecutor.StreamChunk, 2)
|
|
chunks <- cliproxyexecutor.StreamChunk{Payload: []byte("initial")}
|
|
chunks <- cliproxyexecutor.StreamChunk{Err: &Error{HTTPStatus: http.StatusBadGateway, Message: "upstream failed"}}
|
|
manager.RegisterExecutor(&homeExecutionStreamExecutor{chunks: chunks})
|
|
|
|
result, errExecute := manager.ExecuteStream(ctx, []string{"home-execution"}, cliproxyexecutor.Request{Model: "model-a"}, cliproxyexecutor.Options{Stream: true})
|
|
if errExecute != nil {
|
|
t.Fatalf("ExecuteStream() error = %v", errExecute)
|
|
}
|
|
if _, ok := <-result.Chunks; !ok {
|
|
t.Fatal("stream closed before initial chunk")
|
|
}
|
|
if chunk, ok := <-result.Chunks; !ok || chunk.Err == nil {
|
|
t.Fatalf("error chunk = %#v, open = %v", chunk, ok)
|
|
}
|
|
select {
|
|
case group := <-releases:
|
|
t.Fatalf("stream released before cancellation: %#v", group)
|
|
default:
|
|
}
|
|
|
|
cancel()
|
|
for range result.Chunks {
|
|
}
|
|
select {
|
|
case group := <-releases:
|
|
if group != (executionregistry.ReleaseGroup{CredentialID: "cred-1", Model: "model-a"}) {
|
|
t.Fatalf("release group = %#v", group)
|
|
}
|
|
case <-time.After(time.Second):
|
|
t.Fatal("stream did not release after cancellation")
|
|
}
|
|
close(chunks)
|
|
}
|
|
|
|
func TestAccountedHomeStreamConsumerCancellationEndsSelection(t *testing.T) {
|
|
ctx, cancel := context.WithCancel(context.Background())
|
|
defer cancel()
|
|
manager := NewManager(nil, nil, nil)
|
|
manager.SetConfig(&internalconfig.Config{Home: internalconfig.HomeConfig{Enabled: true}})
|
|
registry := executionregistry.New()
|
|
releases := make(chan executionregistry.ReleaseGroup, 1)
|
|
registry.SetReleaseSink(func(group executionregistry.ReleaseGroup, _ int64) { releases <- group })
|
|
manager.PublishHomeDispatch(&accountedHomeExecutionDispatcher{auths: []Auth{{
|
|
ID: "cred-1", Provider: "home-execution", Status: StatusActive,
|
|
}}}, registry, 1)
|
|
chunks := make(chan cliproxyexecutor.StreamChunk, 1)
|
|
chunks <- cliproxyexecutor.StreamChunk{Payload: []byte("initial")}
|
|
manager.RegisterExecutor(&homeExecutionStreamExecutor{chunks: chunks})
|
|
|
|
result, errExecute := manager.ExecuteStream(ctx, []string{"home-execution"}, cliproxyexecutor.Request{Model: "model-a"}, cliproxyexecutor.Options{Stream: true})
|
|
if errExecute != nil {
|
|
t.Fatalf("ExecuteStream() error = %v", errExecute)
|
|
}
|
|
cancel()
|
|
for range result.Chunks {
|
|
}
|
|
select {
|
|
case group := <-releases:
|
|
if group != (executionregistry.ReleaseGroup{CredentialID: "cred-1", Model: "model-a"}) {
|
|
t.Fatalf("release group = %#v", group)
|
|
}
|
|
case <-time.After(time.Second):
|
|
t.Fatal("stream did not release after consumer cancellation")
|
|
}
|
|
}
|
|
|
|
type retryingAccountedHomeExecutor struct{ calls atomic.Int32 }
|
|
|
|
func (*retryingAccountedHomeExecutor) Identifier() string { return "home-execution" }
|
|
func (e *retryingAccountedHomeExecutor) Execute(context.Context, *Auth, cliproxyexecutor.Request, cliproxyexecutor.Options) (cliproxyexecutor.Response, error) {
|
|
if e.calls.Add(1) == 1 {
|
|
return cliproxyexecutor.Response{}, &Error{HTTPStatus: http.StatusBadGateway, Message: "upstream failed"}
|
|
}
|
|
return cliproxyexecutor.Response{Payload: []byte("ok")}, nil
|
|
}
|
|
func (*retryingAccountedHomeExecutor) ExecuteStream(context.Context, *Auth, cliproxyexecutor.Request, cliproxyexecutor.Options) (*cliproxyexecutor.StreamResult, error) {
|
|
return nil, nil
|
|
}
|
|
func (*retryingAccountedHomeExecutor) Refresh(context.Context, *Auth) (*Auth, error) { return nil, nil }
|
|
func (*retryingAccountedHomeExecutor) CountTokens(context.Context, *Auth, cliproxyexecutor.Request, cliproxyexecutor.Options) (cliproxyexecutor.Response, error) {
|
|
return cliproxyexecutor.Response{}, nil
|
|
}
|
|
func (*retryingAccountedHomeExecutor) HttpRequest(context.Context, *Auth, *http.Request) (*http.Response, error) {
|
|
return nil, nil
|
|
}
|
|
|
|
func TestAccountedHomeRetrySelectsAndReleasesEveryAttempt(t *testing.T) {
|
|
manager := NewManager(nil, nil, nil)
|
|
manager.SetConfig(&internalconfig.Config{Home: internalconfig.HomeConfig{Enabled: true}})
|
|
registry := executionregistry.New()
|
|
releases := make(chan executionregistry.ReleaseGroup, 2)
|
|
registry.SetReleaseSink(func(group executionregistry.ReleaseGroup, _ int64) { releases <- group })
|
|
dispatcher := &accountedHomeExecutionDispatcher{auths: []Auth{
|
|
{ID: "cred-1", Provider: "home-execution", Status: StatusActive},
|
|
{ID: "cred-2", Provider: "home-execution", Status: StatusActive},
|
|
}}
|
|
manager.PublishHomeDispatch(dispatcher, registry, 1)
|
|
executor := &retryingAccountedHomeExecutor{}
|
|
manager.RegisterExecutor(executor)
|
|
|
|
if _, errExecute := manager.Execute(context.Background(), []string{"home-execution"}, cliproxyexecutor.Request{Model: "model-a"}, cliproxyexecutor.Options{}); errExecute != nil {
|
|
t.Fatalf("Execute() error = %v", errExecute)
|
|
}
|
|
if got := dispatcher.calls.Load(); got != 2 {
|
|
t.Fatalf("Home selections = %d, want 2", got)
|
|
}
|
|
if got := executor.calls.Load(); got != 2 {
|
|
t.Fatalf("executor attempts = %d, want 2", got)
|
|
}
|
|
groups := map[executionregistry.ReleaseGroup]bool{}
|
|
for range 2 {
|
|
groups[<-releases] = true
|
|
}
|
|
for _, credentialID := range []string{"cred-1", "cred-2"} {
|
|
if !groups[executionregistry.ReleaseGroup{CredentialID: credentialID, Model: "model-a"}] {
|
|
t.Fatalf("missing release for %s: %#v", credentialID, groups)
|
|
}
|
|
}
|
|
}
|
|
|
|
func TestHomeStreamWithoutSourceEndsSelection(t *testing.T) {
|
|
manager := NewManager(nil, nil, nil)
|
|
manager.SetConfig(&internalconfig.Config{Home: internalconfig.HomeConfig{Enabled: true}})
|
|
registry := executionregistry.New()
|
|
manager.PublishHomeDispatch(&homePerSelectionDispatcher{auths: []Auth{{
|
|
ID: "home-auth", Provider: "home-execution", Status: StatusActive,
|
|
}}}, registry, 1)
|
|
manager.RegisterExecutor(&missingHomeStreamSourceExecutor{})
|
|
|
|
result, errExecute := manager.ExecuteStream(context.Background(), []string{"home-execution"}, cliproxyexecutor.Request{Model: "test"}, cliproxyexecutor.Options{Stream: true})
|
|
if errExecute == nil {
|
|
t.Fatalf("ExecuteStream() result = %#v, want error", result)
|
|
}
|
|
|
|
drainCtx, cancelDrain := context.WithTimeout(context.Background(), time.Second)
|
|
defer cancelDrain()
|
|
if errDrain := registry.Drain(drainCtx); errDrain != nil {
|
|
t.Fatalf("Drain() error = %v", errDrain)
|
|
}
|
|
}
|
|
|
|
// homeRequestMetadataSnapshot captures the client request metadata a context carries.
|
|
type homeRequestMetadataSnapshot struct {
|
|
requestedModel string
|
|
reasoningEffort string
|
|
serviceTier string
|
|
generate bool
|
|
}
|
|
|
|
func homeRequestMetadataFromContext(ctx context.Context) homeRequestMetadataSnapshot {
|
|
return homeRequestMetadataSnapshot{
|
|
requestedModel: coreusage.RequestedModelAliasFromContext(ctx),
|
|
reasoningEffort: coreusage.ReasoningEffortFromContext(ctx),
|
|
serviceTier: coreusage.ServiceTierFromContext(ctx),
|
|
generate: coreusage.GenerateFromContext(ctx),
|
|
}
|
|
}
|
|
|
|
// homeRequestMetadataExecutor records the metadata visible at auth preparation and execution.
|
|
type homeRequestMetadataExecutor struct {
|
|
mu sync.Mutex
|
|
prepareMetadata homeRequestMetadataSnapshot
|
|
executeMetadata homeRequestMetadataSnapshot
|
|
// prepareErrOnce fails only the first preparation so Home redispatch still terminates.
|
|
prepareErrOnce error
|
|
executeErr error
|
|
}
|
|
|
|
func (*homeRequestMetadataExecutor) Identifier() string { return "home-execution" }
|
|
|
|
func (*homeRequestMetadataExecutor) ShouldPrepareRequestAuth(*Auth) bool { return true }
|
|
|
|
func (e *homeRequestMetadataExecutor) PrepareRequestAuth(ctx context.Context, auth *Auth) (*Auth, error) {
|
|
e.mu.Lock()
|
|
defer e.mu.Unlock()
|
|
e.prepareMetadata = homeRequestMetadataFromContext(ctx)
|
|
if e.prepareErrOnce != nil {
|
|
errPrepare := e.prepareErrOnce
|
|
e.prepareErrOnce = nil
|
|
return nil, errPrepare
|
|
}
|
|
return auth, nil
|
|
}
|
|
|
|
func (e *homeRequestMetadataExecutor) recordExecution(ctx context.Context) error {
|
|
e.mu.Lock()
|
|
defer e.mu.Unlock()
|
|
e.executeMetadata = homeRequestMetadataFromContext(ctx)
|
|
return e.executeErr
|
|
}
|
|
|
|
func (e *homeRequestMetadataExecutor) snapshots() (homeRequestMetadataSnapshot, homeRequestMetadataSnapshot) {
|
|
e.mu.Lock()
|
|
defer e.mu.Unlock()
|
|
return e.prepareMetadata, e.executeMetadata
|
|
}
|
|
|
|
func (e *homeRequestMetadataExecutor) Execute(ctx context.Context, _ *Auth, _ cliproxyexecutor.Request, _ cliproxyexecutor.Options) (cliproxyexecutor.Response, error) {
|
|
if errExecute := e.recordExecution(ctx); errExecute != nil {
|
|
return cliproxyexecutor.Response{}, errExecute
|
|
}
|
|
return cliproxyexecutor.Response{Payload: []byte("ok")}, nil
|
|
}
|
|
|
|
func (e *homeRequestMetadataExecutor) ExecuteStream(ctx context.Context, _ *Auth, _ cliproxyexecutor.Request, _ cliproxyexecutor.Options) (*cliproxyexecutor.StreamResult, error) {
|
|
if errExecute := e.recordExecution(ctx); errExecute != nil {
|
|
return nil, errExecute
|
|
}
|
|
chunks := make(chan cliproxyexecutor.StreamChunk, 1)
|
|
chunks <- cliproxyexecutor.StreamChunk{Payload: []byte("ok")}
|
|
close(chunks)
|
|
return &cliproxyexecutor.StreamResult{Chunks: chunks}, nil
|
|
}
|
|
|
|
func (*homeRequestMetadataExecutor) Refresh(context.Context, *Auth) (*Auth, error) {
|
|
return nil, nil
|
|
}
|
|
|
|
func (e *homeRequestMetadataExecutor) CountTokens(ctx context.Context, _ *Auth, _ cliproxyexecutor.Request, _ cliproxyexecutor.Options) (cliproxyexecutor.Response, error) {
|
|
if errExecute := e.recordExecution(ctx); errExecute != nil {
|
|
return cliproxyexecutor.Response{}, errExecute
|
|
}
|
|
return cliproxyexecutor.Response{Payload: []byte("ok")}, nil
|
|
}
|
|
|
|
func (*homeRequestMetadataExecutor) HttpRequest(context.Context, *Auth, *http.Request) (*http.Response, error) {
|
|
return nil, nil
|
|
}
|
|
|
|
// homeRequestMetadataHook buffers every Home result so a synchronous OnResult never blocks execution.
|
|
type homeRequestMetadataHook struct {
|
|
results chan homeRequestMetadataSnapshot
|
|
}
|
|
|
|
func newHomeRequestMetadataHook() *homeRequestMetadataHook {
|
|
return &homeRequestMetadataHook{results: make(chan homeRequestMetadataSnapshot, 8)}
|
|
}
|
|
|
|
func (*homeRequestMetadataHook) OnAuthRegistered(context.Context, *Auth) {}
|
|
func (*homeRequestMetadataHook) OnAuthUpdated(context.Context, *Auth) {}
|
|
func (h *homeRequestMetadataHook) OnResult(ctx context.Context, _ Result) {
|
|
select {
|
|
case h.results <- homeRequestMetadataFromContext(ctx):
|
|
default:
|
|
}
|
|
}
|
|
|
|
func (h *homeRequestMetadataHook) awaitResult(t *testing.T) homeRequestMetadataSnapshot {
|
|
t.Helper()
|
|
select {
|
|
case snapshot := <-h.results:
|
|
return snapshot
|
|
case <-time.After(time.Second):
|
|
t.Fatal("Home result hook did not run")
|
|
return homeRequestMetadataSnapshot{}
|
|
}
|
|
}
|
|
|
|
func assertHomeRequestMetadata(t *testing.T, got homeRequestMetadataSnapshot, serviceTier string) {
|
|
t.Helper()
|
|
want := homeRequestMetadataSnapshot{
|
|
requestedModel: "client-model",
|
|
reasoningEffort: "high",
|
|
serviceTier: serviceTier,
|
|
generate: false,
|
|
}
|
|
if got != want {
|
|
t.Fatalf("request metadata = %#v, want %#v", got, want)
|
|
}
|
|
}
|
|
|
|
func newHomeRequestMetadataManager(t *testing.T, executor *homeRequestMetadataExecutor, hook Hook) *Manager {
|
|
t.Helper()
|
|
manager := NewManager(nil, nil, hook)
|
|
manager.SetConfig(&internalconfig.Config{Home: internalconfig.HomeConfig{Enabled: true}})
|
|
manager.PublishHomeDispatch(homeExecutionDispatcher{}, executionregistry.New(), 1)
|
|
manager.RegisterExecutor(executor)
|
|
return manager
|
|
}
|
|
|
|
// homeRequestMetadataOptions mirrors handler-populated metadata. Handlers already derive the
|
|
// OpenAI "auto" default for an omitted tier (see sdk/api/handlers metadata tests); this layer
|
|
// only has to carry whatever the handler resolved.
|
|
func homeRequestMetadataOptions(serviceTier string) cliproxyexecutor.Options {
|
|
return cliproxyexecutor.Options{Metadata: map[string]any{
|
|
cliproxyexecutor.RequestedModelMetadataKey: "client-model",
|
|
cliproxyexecutor.ReasoningEffortMetadataKey: "high",
|
|
cliproxyexecutor.ServiceTierMetadataKey: serviceTier,
|
|
cliproxyexecutor.GenerateMetadataKey: false,
|
|
}}
|
|
}
|
|
|
|
type homeRequestMetadataPath struct {
|
|
name string
|
|
run func(*Manager, cliproxyexecutor.Options) error
|
|
}
|
|
|
|
func homeExecuteMetadataPath() homeRequestMetadataPath {
|
|
return homeRequestMetadataPath{
|
|
name: "execute",
|
|
run: func(manager *Manager, opts cliproxyexecutor.Options) error {
|
|
_, errExecute := manager.Execute(context.Background(), []string{"home-execution"}, cliproxyexecutor.Request{Model: "route-model"}, opts)
|
|
return errExecute
|
|
},
|
|
}
|
|
}
|
|
|
|
func homeCountMetadataPath() homeRequestMetadataPath {
|
|
return homeRequestMetadataPath{
|
|
name: "count_tokens",
|
|
run: func(manager *Manager, opts cliproxyexecutor.Options) error {
|
|
_, errCount := manager.ExecuteCount(context.Background(), []string{"home-execution"}, cliproxyexecutor.Request{Model: "route-model"}, opts)
|
|
return errCount
|
|
},
|
|
}
|
|
}
|
|
|
|
func homeStreamMetadataPath() homeRequestMetadataPath {
|
|
return homeRequestMetadataPath{
|
|
name: "stream",
|
|
run: func(manager *Manager, opts cliproxyexecutor.Options) error {
|
|
opts.Stream = true
|
|
result, errStream := manager.ExecuteStream(context.Background(), []string{"home-execution"}, cliproxyexecutor.Request{Model: "route-model"}, opts)
|
|
if errStream != nil {
|
|
return errStream
|
|
}
|
|
for range result.Chunks {
|
|
}
|
|
return nil
|
|
},
|
|
}
|
|
}
|
|
|
|
// TestHomeExecutionPropagatesRequestMetadata covers the Home regression from issue #4791: the
|
|
// executor context must carry the client request metadata at auth preparation, at execution, and
|
|
// in the Home result usage record.
|
|
func TestHomeExecutionPropagatesRequestMetadata(t *testing.T) {
|
|
paths := []homeRequestMetadataPath{homeExecuteMetadataPath(), homeCountMetadataPath(), homeStreamMetadataPath()}
|
|
|
|
for _, path := range paths {
|
|
for _, serviceTier := range []string{"priority", coreusage.AutoServiceTier} {
|
|
t.Run(path.name+"/"+serviceTier, func(t *testing.T) {
|
|
executor := &homeRequestMetadataExecutor{}
|
|
hook := newHomeRequestMetadataHook()
|
|
manager := newHomeRequestMetadataManager(t, executor, hook)
|
|
|
|
if errRun := path.run(manager, homeRequestMetadataOptions(serviceTier)); errRun != nil {
|
|
t.Fatalf("execution error = %v", errRun)
|
|
}
|
|
prepareMetadata, executeMetadata := executor.snapshots()
|
|
assertHomeRequestMetadata(t, prepareMetadata, serviceTier)
|
|
assertHomeRequestMetadata(t, executeMetadata, serviceTier)
|
|
assertHomeRequestMetadata(t, hook.awaitResult(t), serviceTier)
|
|
})
|
|
}
|
|
}
|
|
}
|
|
|
|
// TestHomeExecutionFailureResultPreservesRequestMetadata keeps the requested tier authoritative in
|
|
// the failure usage record instead of falling back to the upstream or default tier.
|
|
func TestHomeExecutionFailureResultPreservesRequestMetadata(t *testing.T) {
|
|
paths := []homeRequestMetadataPath{homeExecuteMetadataPath(), homeCountMetadataPath()}
|
|
|
|
for _, path := range paths {
|
|
t.Run(path.name, func(t *testing.T) {
|
|
executor := &homeRequestMetadataExecutor{
|
|
executeErr: &Error{HTTPStatus: http.StatusBadRequest, Message: "invalid request"},
|
|
}
|
|
hook := newHomeRequestMetadataHook()
|
|
manager := newHomeRequestMetadataManager(t, executor, hook)
|
|
|
|
if errRun := path.run(manager, homeRequestMetadataOptions("priority")); errRun == nil {
|
|
t.Fatal("execution error = nil, want invalid request")
|
|
}
|
|
assertHomeRequestMetadata(t, hook.awaitResult(t), "priority")
|
|
})
|
|
}
|
|
}
|
|
|
|
// TestHomePrepareFailureResultPreservesRequestMetadata covers the prepare_failed Home result paths,
|
|
// which report usage before any executor call happens.
|
|
func TestHomePrepareFailureResultPreservesRequestMetadata(t *testing.T) {
|
|
paths := []homeRequestMetadataPath{homeExecuteMetadataPath(), homeCountMetadataPath(), homeStreamMetadataPath()}
|
|
|
|
for _, path := range paths {
|
|
t.Run(path.name, func(t *testing.T) {
|
|
executor := &homeRequestMetadataExecutor{
|
|
prepareErrOnce: &Error{Code: "prepare_failed", Message: "prepare failed"},
|
|
}
|
|
hook := newHomeRequestMetadataHook()
|
|
manager := newHomeRequestMetadataManager(t, executor, hook)
|
|
|
|
_ = path.run(manager, homeRequestMetadataOptions("priority"))
|
|
prepareMetadata, _ := executor.snapshots()
|
|
assertHomeRequestMetadata(t, prepareMetadata, "priority")
|
|
assertHomeRequestMetadata(t, hook.awaitResult(t), "priority")
|
|
})
|
|
}
|
|
}
|