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")
		})
	}
}