ollama/cmd/cmd_launcher_test.go
2026-06-22 16:48:38 -07:00

581 lines
17 KiB
Go

package cmd
import (
"context"
"errors"
"net/http"
"net/http/httptest"
"testing"
"github.com/spf13/cobra"
"github.com/ollama/ollama/cmd/config"
"github.com/ollama/ollama/cmd/launch"
"github.com/ollama/ollama/cmd/tui"
"github.com/ollama/ollama/envconfig"
)
func setCmdTestHome(t *testing.T, dir string) {
t.Helper()
t.Setenv("HOME", dir)
t.Setenv("USERPROFILE", dir)
envconfig.ReloadServerConfig()
}
func unexpectedRunModelResolution(t *testing.T) func(context.Context, launch.RunModelRequest) (string, error) {
t.Helper()
return func(ctx context.Context, req launch.RunModelRequest) (string, error) {
t.Fatalf("did not expect run-model resolution: %+v", req)
return "", nil
}
}
func unexpectedIntegrationLaunch(t *testing.T) func(context.Context, launch.IntegrationLaunchRequest) error {
t.Helper()
return func(ctx context.Context, req launch.IntegrationLaunchRequest) error {
t.Fatalf("did not expect integration launch: %+v", req)
return nil
}
}
func unexpectedModelLaunch(t *testing.T) func(*cobra.Command, string) error {
t.Helper()
return func(cmd *cobra.Command, model string) error {
t.Fatalf("did not expect chat launch: %s", model)
return nil
}
}
func TestRunAgentModelPickerUsesSavedModelWhenAvailable(t *testing.T) {
setCmdTestHome(t, t.TempDir())
if err := config.SetAgentSignInPromptSeen(true); err != nil {
t.Fatal(err)
}
var gotReq launch.RunModelRequest
var launched string
prefetchedAccount := &launch.AccountState{}
accountUpdates := func(context.Context) <-chan *launch.AccountState { return nil }
deps := agentModelPickerDeps{
resolveRunModel: func(ctx context.Context, req launch.RunModelRequest) (string, error) {
gotReq = req
return "qwen3:8b", nil
},
runModel: func(cmd *cobra.Command, model string) error {
launched = model
return nil
},
accountState: func() *launch.AccountState {
return prefetchedAccount
},
accountStateUpdates: accountUpdates,
}
cmd := &cobra.Command{}
cmd.SetContext(context.Background())
if err := runAgentModelPickerWithDeps(cmd, deps); err != nil {
t.Fatalf("runAgentModelPickerWithDeps error: %v", err)
}
if gotReq.ForcePicker {
t.Fatal("expected root agent flow to reuse a saved model when available")
}
if gotReq.AccountState != prefetchedAccount {
t.Fatal("expected prefetched account state to be passed to model picker")
}
if gotReq.AccountStateProvider == nil || gotReq.AccountStateUpdates == nil {
t.Fatal("expected account state callbacks to be passed to model picker")
}
if launched != "qwen3:8b" {
t.Fatalf("launched model = %q, want qwen3:8b", launched)
}
}
func TestRunAgentModelPickerFallsBackToPickerWhenPlanVerificationFails(t *testing.T) {
setCmdTestHome(t, t.TempDir())
var requests []launch.RunModelRequest
var launched string
deps := agentModelPickerDeps{
resolveRunModel: func(ctx context.Context, req launch.RunModelRequest) (string, error) {
requests = append(requests, req)
if len(requests) == 1 {
return "", launch.ErrPlanVerificationUnavailable
}
return "llama3.2", nil
},
runModel: func(cmd *cobra.Command, model string) error {
launched = model
return nil
},
accountState: func() *launch.AccountState {
return &launch.AccountState{}
},
}
cmd := &cobra.Command{}
cmd.SetContext(context.Background())
if err := runAgentModelPickerWithDeps(cmd, deps); err != nil {
t.Fatalf("runAgentModelPickerWithDeps error: %v", err)
}
if len(requests) != 2 {
t.Fatalf("resolve calls = %d, want 2", len(requests))
}
if requests[0].ForcePicker {
t.Fatal("first request should try the saved model path")
}
if !requests[1].ForcePicker {
t.Fatal("second request should force the model picker")
}
if requests[1].AccountStateProvider != nil {
t.Fatal("retry should not keep using the stale account-state provider")
}
if requests[1].AccountState == nil {
t.Fatal("retry should pass an explicit unknown account state")
}
if launched != "llama3.2" {
t.Fatalf("launched model = %q, want llama3.2", launched)
}
}
func TestRunAgentModelPickerReturnsPlanVerificationErrorWhenPickerRetryFails(t *testing.T) {
setCmdTestHome(t, t.TempDir())
var calls int
deps := agentModelPickerDeps{
resolveRunModel: func(ctx context.Context, req launch.RunModelRequest) (string, error) {
calls++
return "", launch.ErrPlanVerificationUnavailable
},
runModel: unexpectedModelLaunch(t),
}
cmd := &cobra.Command{}
cmd.SetContext(context.Background())
err := runAgentModelPickerWithDeps(cmd, deps)
if !errors.Is(err, launch.ErrPlanVerificationUnavailable) {
t.Fatalf("error = %v, want ErrPlanVerificationUnavailable", err)
}
if calls != 2 {
t.Fatalf("resolve calls = %d, want 2", calls)
}
}
func TestMaybeRunAgentOnboarding(t *testing.T) {
t.Run("prompts once and saves seen state", func(t *testing.T) {
setCmdTestHome(t, t.TempDir())
oldPrompt := agentOnboardingPrompt
oldSignedIn := agentOnboardingSignedInStatus
t.Cleanup(func() {
agentOnboardingPrompt = oldPrompt
agentOnboardingSignedInStatus = oldSignedIn
})
var prompts int
agentOnboardingPrompt = func() (bool, error) {
prompts++
return false, nil
}
agentOnboardingSignedInStatus = func(context.Context) (bool, bool) {
return false, true
}
signIn, err := maybeRunAgentOnboarding(context.Background())
if err != nil {
t.Fatalf("maybeRunAgentOnboarding error: %v", err)
}
if signIn {
t.Fatal("signIn = true, want false")
}
if prompts != 1 {
t.Fatalf("prompts = %d, want 1", prompts)
}
if !config.AgentSignInPromptSeen() {
t.Fatal("expected onboarding state to be saved")
}
signIn, err = maybeRunAgentOnboarding(context.Background())
if err != nil {
t.Fatalf("second maybeRunAgentOnboarding error: %v", err)
}
if signIn {
t.Fatal("second signIn = true, want false")
}
if prompts != 1 {
t.Fatalf("prompt should not run again, prompts = %d", prompts)
}
})
t.Run("skips prompt when already signed in", func(t *testing.T) {
setCmdTestHome(t, t.TempDir())
oldPrompt := agentOnboardingPrompt
oldSignedIn := agentOnboardingSignedInStatus
t.Cleanup(func() {
agentOnboardingPrompt = oldPrompt
agentOnboardingSignedInStatus = oldSignedIn
})
var prompts int
agentOnboardingPrompt = func() (bool, error) {
prompts++
return false, nil
}
agentOnboardingSignedInStatus = func(context.Context) (bool, bool) {
return true, true
}
signIn, err := maybeRunAgentOnboarding(context.Background())
if err != nil {
t.Fatalf("maybeRunAgentOnboarding error: %v", err)
}
if signIn {
t.Fatal("signIn = true, want false")
}
if prompts != 0 {
t.Fatalf("prompts = %d, want 0", prompts)
}
if !config.AgentSignInPromptSeen() {
t.Fatal("expected onboarding state to be saved")
}
})
t.Run("skips prompt when signed-in check is unknown", func(t *testing.T) {
setCmdTestHome(t, t.TempDir())
oldPrompt := agentOnboardingPrompt
oldSignedIn := agentOnboardingSignedInStatus
t.Cleanup(func() {
agentOnboardingPrompt = oldPrompt
agentOnboardingSignedInStatus = oldSignedIn
})
var prompts int
agentOnboardingPrompt = func() (bool, error) {
prompts++
return false, nil
}
agentOnboardingSignedInStatus = func(context.Context) (bool, bool) {
return false, false
}
signIn, err := maybeRunAgentOnboarding(context.Background())
if err != nil {
t.Fatalf("maybeRunAgentOnboarding error: %v", err)
}
if signIn {
t.Fatal("signIn = true, want false")
}
if prompts != 0 {
t.Fatalf("prompts = %d, want 0", prompts)
}
if config.AgentSignInPromptSeen() {
t.Fatal("unknown auth state should not save onboarding state")
}
})
t.Run("cancel does not save seen state", func(t *testing.T) {
setCmdTestHome(t, t.TempDir())
oldPrompt := agentOnboardingPrompt
oldSignedIn := agentOnboardingSignedInStatus
t.Cleanup(func() {
agentOnboardingPrompt = oldPrompt
agentOnboardingSignedInStatus = oldSignedIn
})
agentOnboardingPrompt = func() (bool, error) {
return false, tui.ErrCancelled
}
agentOnboardingSignedInStatus = func(context.Context) (bool, bool) {
return false, true
}
_, err := maybeRunAgentOnboarding(context.Background())
if !errors.Is(err, launch.ErrCancelled) {
t.Fatalf("error = %v, want launch.ErrCancelled", err)
}
if config.AgentSignInPromptSeen() {
t.Fatal("cancel should not save onboarding state")
}
})
}
func TestRunAgentOnboardingSignInEmptyWhoamiDoesNotSilentlySucceed(t *testing.T) {
server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
if r.URL.Path != "/api/me" {
http.NotFound(w, r)
return
}
w.Header().Set("Content-Type", "application/json")
_, _ = w.Write([]byte(`{}`))
}))
defer server.Close()
t.Setenv("OLLAMA_HOST", server.URL)
err := runAgentOnboardingSignIn(context.Background())
if !errors.Is(err, errAgentOnboardingNotSignedIn) {
t.Fatalf("error = %v, want errAgentOnboardingNotSignedIn", err)
}
}
func TestRunInteractiveTUI_RunModelActionsUseResolveRunModel(t *testing.T) {
tests := []struct {
name string
action tui.TUIAction
wantForce bool
wantModel string
}{
{
name: "enter uses saved model flow",
action: tui.TUIAction{Kind: tui.TUIActionRunModel},
wantModel: "qwen3:8b",
},
{
name: "right forces picker",
action: tui.TUIAction{Kind: tui.TUIActionRunModel, ForceConfigure: true},
wantForce: true,
wantModel: "glm-5:cloud",
},
}
for _, tt := range tests {
t.Run(tt.name, func(t *testing.T) {
setCmdTestHome(t, t.TempDir())
var menuCalls int
runMenu := func(state *launch.LauncherState) (tui.TUIAction, error) {
menuCalls++
if menuCalls == 1 {
return tt.action, nil
}
return tui.TUIAction{Kind: tui.TUIActionNone}, nil
}
var gotReq launch.RunModelRequest
var launched string
prefetchedAccount := &launch.AccountState{}
accountUpdates := func(context.Context) <-chan *launch.AccountState { return nil }
deps := launcherDeps{
buildState: func(ctx context.Context) (*launch.LauncherState, error) {
return &launch.LauncherState{}, nil
},
runMenu: func(state *launch.LauncherState) (tui.TUIAction, error) {
if state.AccountState != prefetchedAccount {
t.Fatalf("prefetched account state was not piped to menu state")
}
return runMenu(state)
},
resolveRunModel: func(ctx context.Context, req launch.RunModelRequest) (string, error) {
gotReq = req
return tt.wantModel, nil
},
launchIntegration: unexpectedIntegrationLaunch(t),
runModel: func(cmd *cobra.Command, model string) error {
launched = model
return nil
},
accountState: func() *launch.AccountState {
return prefetchedAccount
},
accountStateUpdates: accountUpdates,
}
cmd := &cobra.Command{}
cmd.SetContext(context.Background())
for {
continueLoop, err := runInteractiveTUIStep(cmd, deps)
if err != nil {
t.Fatalf("unexpected step error: %v", err)
}
if !continueLoop {
break
}
}
if gotReq.ForcePicker != tt.wantForce {
t.Fatalf("expected ForcePicker=%v, got %v", tt.wantForce, gotReq.ForcePicker)
}
if gotReq.AccountState != prefetchedAccount {
t.Fatalf("expected prefetched account state to be passed to run model request")
}
if gotReq.AccountStateUpdates == nil {
t.Fatalf("expected account state updates to be passed to run model request")
}
if launched != tt.wantModel {
t.Fatalf("expected interactive launcher to run %q, got %q", tt.wantModel, launched)
}
if got := config.LastSelection(); got != "run" {
t.Fatalf("expected last selection to be run, got %q", got)
}
})
}
}
func TestRunInteractiveTUI_IntegrationActionsUseLaunchIntegration(t *testing.T) {
tests := []struct {
name string
action tui.TUIAction
wantForce bool
}{
{
name: "enter launches integration",
action: tui.TUIAction{Kind: tui.TUIActionLaunchIntegration, Integration: "claude"},
},
{
name: "right forces configure",
action: tui.TUIAction{Kind: tui.TUIActionLaunchIntegration, Integration: "claude", ForceConfigure: true},
wantForce: true,
},
}
for _, tt := range tests {
t.Run(tt.name, func(t *testing.T) {
setCmdTestHome(t, t.TempDir())
var menuCalls int
runMenu := func(state *launch.LauncherState) (tui.TUIAction, error) {
menuCalls++
if menuCalls == 1 {
return tt.action, nil
}
return tui.TUIAction{Kind: tui.TUIActionNone}, nil
}
var gotReq launch.IntegrationLaunchRequest
prefetchedAccount := &launch.AccountState{}
accountUpdates := func(context.Context) <-chan *launch.AccountState { return nil }
deps := launcherDeps{
buildState: func(ctx context.Context) (*launch.LauncherState, error) {
return &launch.LauncherState{}, nil
},
runMenu: func(state *launch.LauncherState) (tui.TUIAction, error) {
if state.AccountState != prefetchedAccount {
t.Fatalf("prefetched account state was not piped to menu state")
}
return runMenu(state)
},
resolveRunModel: unexpectedRunModelResolution(t),
launchIntegration: func(ctx context.Context, req launch.IntegrationLaunchRequest) error {
gotReq = req
return nil
},
runModel: unexpectedModelLaunch(t),
accountState: func() *launch.AccountState {
return prefetchedAccount
},
accountStateUpdates: accountUpdates,
}
cmd := &cobra.Command{}
cmd.SetContext(context.Background())
for {
continueLoop, err := runInteractiveTUIStep(cmd, deps)
if err != nil {
t.Fatalf("unexpected step error: %v", err)
}
if !continueLoop {
break
}
}
if gotReq.Name != "claude" {
t.Fatalf("expected integration name to be passed through, got %q", gotReq.Name)
}
if gotReq.ForceConfigure != tt.wantForce {
t.Fatalf("expected ForceConfigure=%v, got %v", tt.wantForce, gotReq.ForceConfigure)
}
if gotReq.AccountState != prefetchedAccount {
t.Fatalf("expected prefetched account state to be passed to integration request")
}
if gotReq.AccountStateUpdates == nil {
t.Fatalf("expected account state updates to be passed to integration request")
}
if got := config.LastSelection(); got != "claude" {
t.Fatalf("expected last selection to be claude, got %q", got)
}
})
}
}
func TestRunLauncherAction_RunModelContinuesAfterCancellation(t *testing.T) {
setCmdTestHome(t, t.TempDir())
cmd := &cobra.Command{}
cmd.SetContext(context.Background())
continueLoop, err := runLauncherAction(cmd, tui.TUIAction{Kind: tui.TUIActionRunModel}, launcherDeps{
buildState: nil,
runMenu: nil,
resolveRunModel: func(ctx context.Context, req launch.RunModelRequest) (string, error) {
return "", launch.ErrCancelled
},
launchIntegration: unexpectedIntegrationLaunch(t),
runModel: unexpectedModelLaunch(t),
})
if err != nil {
t.Fatalf("expected nil error on cancellation, got %v", err)
}
if !continueLoop {
t.Fatal("expected cancellation to continue the menu loop")
}
}
func TestRunLauncherAction_GUIAppsExitTUILoop(t *testing.T) {
setCmdTestHome(t, t.TempDir())
cmd := &cobra.Command{}
cmd.SetContext(context.Background())
for _, integration := range []string{"codex-app", "vscode"} {
continueLoop, err := runLauncherAction(cmd, tui.TUIAction{Kind: tui.TUIActionLaunchIntegration, Integration: integration}, launcherDeps{
resolveRunModel: unexpectedRunModelResolution(t),
launchIntegration: func(ctx context.Context, req launch.IntegrationLaunchRequest) error {
return nil
},
runModel: unexpectedModelLaunch(t),
})
if err != nil {
t.Fatalf("expected nil error for %s, got %v", integration, err)
}
if continueLoop {
t.Fatalf("expected %s launch to exit the TUI loop (return false)", integration)
}
}
// Other integrations should continue the TUI loop (return true).
continueLoop, err := runLauncherAction(cmd, tui.TUIAction{Kind: tui.TUIActionLaunchIntegration, Integration: "claude"}, launcherDeps{
resolveRunModel: unexpectedRunModelResolution(t),
launchIntegration: func(ctx context.Context, req launch.IntegrationLaunchRequest) error {
return nil
},
runModel: unexpectedModelLaunch(t),
})
if err != nil {
t.Fatalf("expected nil error, got %v", err)
}
if !continueLoop {
t.Fatal("expected non-vscode integration to continue the TUI loop (return true)")
}
}
func TestRunLauncherAction_IntegrationContinuesAfterCancellation(t *testing.T) {
setCmdTestHome(t, t.TempDir())
cmd := &cobra.Command{}
cmd.SetContext(context.Background())
continueLoop, err := runLauncherAction(cmd, tui.TUIAction{Kind: tui.TUIActionLaunchIntegration, Integration: "claude"}, launcherDeps{
buildState: nil,
runMenu: nil,
resolveRunModel: unexpectedRunModelResolution(t),
launchIntegration: func(ctx context.Context, req launch.IntegrationLaunchRequest) error {
return launch.ErrCancelled
},
runModel: unexpectedModelLaunch(t),
})
if err != nil {
t.Fatalf("expected nil error on cancellation, got %v", err)
}
if !continueLoop {
t.Fatal("expected cancellation to continue the menu loop")
}
}