mirror of
https://github.com/ollama/ollama.git
synced 2026-11-12 03:46:14 -06:00
415 lines
17 KiB
Go
415 lines
17 KiB
Go
package launch
|
|
|
|
import (
|
|
"context"
|
|
"encoding/json"
|
|
"errors"
|
|
"fmt"
|
|
"net/http"
|
|
"net/http/httptest"
|
|
"os"
|
|
"path/filepath"
|
|
"runtime"
|
|
"slices"
|
|
"strings"
|
|
"testing"
|
|
"time"
|
|
|
|
"github.com/ollama/ollama/cmd/config"
|
|
"github.com/spf13/cobra"
|
|
)
|
|
|
|
func TestLaunchCmdPiPickerPolicy(t *testing.T) {
|
|
if runtime.GOOS == "windows" {
|
|
t.Skip("uses a POSIX shell test binary")
|
|
}
|
|
for _, tc := range []struct {
|
|
name string
|
|
saved []string
|
|
args []string
|
|
wantPicker int
|
|
wantModel string
|
|
wantError string
|
|
wantPiArgs []string
|
|
runtimeOverride bool
|
|
corruptConfig bool
|
|
}{
|
|
{name: "available saved model", saved: []string{"keep-model"}, wantModel: "keep-model"},
|
|
{name: "legacy selection migrates to one starting model", saved: []string{"keep-model", "remote:cloud"}, wantModel: "keep-model"},
|
|
{name: "first setup", wantPicker: 1, wantModel: "chosen"},
|
|
{name: "missing saved default", saved: []string{"missing", "keep-model"}, wantPicker: 1, wantModel: "chosen"},
|
|
{name: "explicit configuration", saved: []string{"keep-model"}, args: []string{"--config"}, wantPicker: 1, wantModel: "chosen"},
|
|
{name: "disabled saved cloud default", saved: []string{"remote:cloud", "keep-model"}, wantError: "cloud is disabled"},
|
|
{name: "disabled explicit cloud override", saved: []string{"keep-model"}, args: []string{"--model", "remote:cloud"}, wantError: "cloud is disabled"},
|
|
{name: "configure replaces disabled cloud default", saved: []string{"remote:cloud", "keep-model"}, args: []string{"--config"}, wantPicker: 1, wantModel: "chosen"},
|
|
{name: "configure replaces disabled explicit cloud target", saved: []string{"keep-model"}, args: []string{"--config", "--model", "remote:cloud"}, wantPicker: 1, wantModel: "chosen"},
|
|
{name: "saved embedding model requires a chat model", saved: []string{"embeddings", "keep-model"}, wantPicker: 1, wantModel: "chosen"},
|
|
{name: "short model flag", saved: []string{"keep-model"}, args: []string{"-m", "chosen"}, wantModel: "chosen", wantPiArgs: []string{"--provider", "ollama", "--model", "chosen"}},
|
|
{name: "explicit embedding model is rejected", saved: []string{"keep-model"}, args: []string{"--model", "embeddings"}, wantError: "does not support chat"},
|
|
{name: "legacy saved embedding model requires a chat model", saved: []string{"legacy-embedding"}, wantPicker: 1, wantModel: "chosen"},
|
|
{name: "disabled saved cloud alias", saved: []string{"remote-alias", "keep-model"}, wantError: "cloud is disabled"},
|
|
{name: "disabled explicit cloud alias", saved: []string{"keep-model"}, args: []string{"--model", "remote-alias"}, wantError: "cloud is disabled"},
|
|
{name: "cloud alias identified by show", saved: []string{"keep-model"}, args: []string{"--model", "hidden-alias"}, wantError: "cloud is disabled"},
|
|
{name: "forwarded qualified model", runtimeOverride: true, saved: []string{"remote:cloud", "keep-model"}, args: []string{"--", "--model", "fixture-user/user-choice"}, wantPiArgs: []string{"--model", "fixture-user/user-choice"}},
|
|
{name: "native override ignores unused malformed launch config", runtimeOverride: true, corruptConfig: true, saved: []string{"keep-model"}, args: []string{"--", "--model", "fixture-user/user-choice"}, wantPiArgs: []string{"--model", "fixture-user/user-choice"}},
|
|
{name: "forwarded provider and model", runtimeOverride: true, saved: []string{"remote:cloud", "keep-model"}, args: []string{"--", "--provider", "fixture-user", "--model", "user-choice"}, wantPiArgs: []string{"--provider", "fixture-user", "--model", "user-choice"}},
|
|
{name: "forwarded model replaces explicit cloud target", runtimeOverride: true, saved: []string{"keep-model"}, args: []string{"--model", "remote:cloud", "--", "--model", "fixture-user/user-choice"}, wantPiArgs: []string{"--model", "fixture-user/user-choice"}},
|
|
{name: "model flag inside a prompt is not an override", wantModel: "keep-model", saved: []string{"keep-model"}, args: []string{"--", "--system-prompt", "--model", "probe"}, wantPiArgs: []string{"--system-prompt", "--model", "probe"}},
|
|
} {
|
|
t.Run(tc.name, func(t *testing.T) {
|
|
home := t.TempDir()
|
|
setLaunchTestHome(t, home)
|
|
withLauncherHooks(t)
|
|
withInteractiveSession(t, true)
|
|
if tc.saved != nil {
|
|
if err := config.SaveIntegration("pi", tc.saved); err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
}
|
|
if tc.corruptConfig {
|
|
if err := os.WriteFile(filepath.Join(home, ".ollama", "config.json"), []byte("{invalid"), 0o644); err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
}
|
|
binDir := t.TempDir()
|
|
t.Setenv("PATH", binDir)
|
|
argsPath := filepath.Join(home, "pi-args")
|
|
for name, data := range map[string]string{
|
|
"pi": fmt.Sprintf("#!/bin/sh\nprintf '%%s\\n' \"$@\" > %q\n", argsPath),
|
|
"npm": "#!/bin/sh\nexit 1\n",
|
|
"package.json": `{"name":"@earendil-works/pi-coding-agent"}`,
|
|
} {
|
|
if err := os.WriteFile(filepath.Join(binDir, name), []byte(data), 0o755); err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
}
|
|
pickerCalls := 0
|
|
DefaultSingleSelector = func(_ string, items []SelectionItem, _ string) (string, error) {
|
|
pickerCalls++
|
|
for _, item := range items {
|
|
if item.Name == "embeddings" || item.Name == "remote-alias" {
|
|
t.Errorf("unusable model appears in Pi picker: %s", item.Name)
|
|
}
|
|
}
|
|
return "chosen", nil
|
|
}
|
|
DefaultMultiSelector = func(_ string, _ []SelectionItem, _ []string) ([]string, error) {
|
|
t.Fatal("Pi must select only its starting model")
|
|
return nil, ErrCancelled
|
|
}
|
|
DefaultConfirmPrompt = func(prompt string, _ ConfirmOptions) (bool, error) {
|
|
if prompt != "Launch Pi now?" || len(tc.args) == 0 {
|
|
t.Errorf("unexpected confirmation: %s", prompt)
|
|
}
|
|
return true, nil
|
|
}
|
|
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
|
switch r.URL.Path {
|
|
case "/api/status":
|
|
fmt.Fprint(w, `{"cloud":{"disabled":true}}`)
|
|
case "/api/experimental/model-recommendations":
|
|
fmt.Fprint(w, `{"recommendations":[{"model":"chosen"}]}`)
|
|
case "/api/tags":
|
|
fmt.Fprint(w, `{"models":[{"name":"chosen","details":{"context_length":65536},"capabilities":["completion","tools"]},{"name":"keep-model","details":{"context_length":32768},"capabilities":["completion","tools","vision"]},{"name":"embeddings","details":{"context_length":8192},"capabilities":["embedding"]},{"name":"remote-alias","remote_model":"remote","details":{"context_length":32768},"capabilities":["completion","tools"]}]}`)
|
|
case "/api/show":
|
|
var req apiShowRequest
|
|
_ = json.NewDecoder(r.Body).Decode(&req)
|
|
if req.Model == "remote:cloud" {
|
|
http.Error(w, `{"error":"cloud is disabled"}`, http.StatusForbidden)
|
|
} else if req.Model == "remote-alias" || req.Model == "hidden-alias" {
|
|
fmt.Fprint(w, `{"remote_model":"remote","capabilities":["completion","tools"]}`)
|
|
} else if req.Model == "embeddings" || req.Model == "legacy-embedding" {
|
|
fmt.Fprint(w, `{"capabilities":["embedding"]}`)
|
|
} else if req.Model == "missing" {
|
|
http.Error(w, `{"error":"missing"}`, http.StatusNotFound)
|
|
} else {
|
|
fmt.Fprint(w, `{}`)
|
|
}
|
|
default:
|
|
t.Errorf("unexpected request: %s", r.URL.Path)
|
|
http.NotFound(w, r)
|
|
}
|
|
}))
|
|
defer srv.Close()
|
|
t.Setenv("OLLAMA_HOST", srv.URL)
|
|
cmd := LaunchCmd(func(*cobra.Command, []string) error { return nil }, nil)
|
|
cmd.SetArgs(append([]string{"pi"}, tc.args...))
|
|
err := cmd.Execute()
|
|
if tc.wantError != "" {
|
|
if err == nil || !strings.Contains(err.Error(), tc.wantError) {
|
|
t.Fatalf("launch error = %v, want %q", err, tc.wantError)
|
|
}
|
|
if pickerCalls != 0 {
|
|
t.Fatal("disabled cloud target unexpectedly opened the picker")
|
|
}
|
|
if _, err := os.Stat(argsPath); !os.IsNotExist(err) {
|
|
t.Fatalf("Pi started despite unavailable cloud target: %v", err)
|
|
}
|
|
saved, err := config.LoadIntegration("pi")
|
|
if err != nil || !slices.Equal(saved.Models, tc.saved) {
|
|
t.Fatalf("saved choices changed on disabled cloud failure: %v, %v", saved, err)
|
|
}
|
|
return
|
|
}
|
|
if err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
if pickerCalls != tc.wantPicker {
|
|
t.Fatalf("picker called %d times, want %d", pickerCalls, tc.wantPicker)
|
|
}
|
|
wantSaved := []string{tc.wantModel}
|
|
if tc.runtimeOverride {
|
|
wantSaved = tc.saved
|
|
}
|
|
saved, err := config.LoadIntegration("pi")
|
|
if tc.corruptConfig {
|
|
data, readErr := os.ReadFile(filepath.Join(home, ".ollama", "config.json"))
|
|
if readErr != nil || string(data) != "{invalid" {
|
|
t.Fatalf("unused malformed config changed: %s, %v", data, readErr)
|
|
}
|
|
} else if err != nil || !slices.Equal(saved.Models, wantSaved) {
|
|
t.Fatalf("saved starting model = %v, %v; want %v", saved, err, wantSaved)
|
|
}
|
|
if !tc.runtimeOverride {
|
|
models := (&Pi{}).Models()
|
|
if !slices.Equal(models, []string{"chosen", "keep-model"}) {
|
|
t.Fatalf("Pi catalog = %v, want all available models", models)
|
|
}
|
|
}
|
|
|
|
data, err := os.ReadFile(argsPath)
|
|
var want []string
|
|
if tc.wantPicker > 0 {
|
|
want = []string{"--provider", "ollama", "--model", tc.wantModel}
|
|
}
|
|
if tc.wantPiArgs != nil {
|
|
want = tc.wantPiArgs
|
|
}
|
|
if err != nil || !slices.Equal(strings.Fields(string(data)), want) {
|
|
t.Fatalf("Pi arguments = %q, %v; want %v", data, err, want)
|
|
}
|
|
})
|
|
}
|
|
}
|
|
|
|
func TestLaunchIntegration_PiDownloadPromptCancellationPreservesConfig(t *testing.T) {
|
|
home := t.TempDir()
|
|
setLaunchTestHome(t, home)
|
|
withLauncherHooks(t)
|
|
withInteractiveSession(t, true)
|
|
if err := config.SaveIntegration("pi", []string{"previous"}); err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
DefaultSingleSelector = func(_ string, _ []SelectionItem, _ string) (string, error) {
|
|
return "missing-new", nil
|
|
}
|
|
DefaultConfirmPrompt = func(prompt string, _ ConfirmOptions) (bool, error) {
|
|
if prompt != "Download missing-new?" {
|
|
t.Errorf("unexpected prompt: %s", prompt)
|
|
}
|
|
return false, ErrCancelled
|
|
}
|
|
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
|
switch r.URL.Path {
|
|
case "/api/status":
|
|
fmt.Fprint(w, `{"cloud":{"disabled":true}}`)
|
|
case "/api/experimental/model-recommendations":
|
|
fmt.Fprint(w, `{"recommendations":[{"model":"missing-new"}]}`)
|
|
case "/api/tags":
|
|
fmt.Fprint(w, `{"models":[{"name":"previous"}]}`)
|
|
case "/api/show":
|
|
var req apiShowRequest
|
|
_ = json.NewDecoder(r.Body).Decode(&req)
|
|
if req.Model == "previous" {
|
|
fmt.Fprint(w, `{}`)
|
|
} else {
|
|
http.Error(w, `{"error":"missing"}`, http.StatusNotFound)
|
|
}
|
|
default:
|
|
t.Errorf("unexpected request: %s", r.URL.Path)
|
|
http.NotFound(w, r)
|
|
}
|
|
}))
|
|
defer srv.Close()
|
|
t.Setenv("OLLAMA_HOST", srv.URL)
|
|
err := LaunchIntegration(context.Background(), IntegrationLaunchRequest{
|
|
Name: "pi", ForceConfigure: true, ConfigureOnly: true,
|
|
})
|
|
if !errors.Is(err, ErrCancelled) {
|
|
t.Fatalf("launch error = %v, want cancellation", err)
|
|
}
|
|
saved, err := config.LoadIntegration("pi")
|
|
if err != nil || !slices.Equal(saved.Models, []string{"previous"}) {
|
|
t.Fatalf("saved config changed on cancellation: %v, %v", saved, err)
|
|
}
|
|
if _, err := os.Stat(filepath.Join(home, ".pi")); !os.IsNotExist(err) {
|
|
t.Fatalf("Pi config was written on cancellation: %v", err)
|
|
}
|
|
}
|
|
|
|
func TestPiRun_PassesSelectedModelToClient(t *testing.T) {
|
|
if runtime.GOOS == "windows" {
|
|
t.Skip("uses a POSIX shell test binary")
|
|
}
|
|
dir := t.TempDir()
|
|
setTestHome(t, dir)
|
|
t.Setenv("PATH", dir)
|
|
argsPath := filepath.Join(dir, "args")
|
|
for name, content := range map[string]string{
|
|
"pi": fmt.Sprintf("#!/bin/sh\nprintf '%%s\\n' \"$@\" > %q\n", argsPath),
|
|
"npm": "#!/bin/sh\nexit 1\n",
|
|
"package.json": `{"name":"@earendil-works/pi-coding-agent"}`,
|
|
} {
|
|
if err := os.WriteFile(filepath.Join(dir, name), []byte(content), 0o755); err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
}
|
|
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
|
fmt.Fprint(w, `{"cloud":{"disabled":true}}`)
|
|
}))
|
|
defer srv.Close()
|
|
t.Setenv("OLLAMA_HOST", srv.URL)
|
|
if err := (&Pi{}).Run("chosen:latest", nil, []string{"--offline", "--print", "probe"}); err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
data, err := os.ReadFile(argsPath)
|
|
if err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
want := []string{"--provider", "ollama", "--model", "chosen:latest", "--offline", "--print", "probe"}
|
|
if got := strings.Split(strings.TrimSpace(string(data)), "\n"); !slices.Equal(got, want) {
|
|
t.Fatalf("Pi process arguments = %v, want %v", got, want)
|
|
}
|
|
}
|
|
|
|
func TestLaunchIntegrationPiPreservesModelChoice(t *testing.T) {
|
|
if runtime.GOOS == "windows" {
|
|
t.Skip("uses a POSIX shell test binary")
|
|
}
|
|
for _, tc := range []struct {
|
|
name string
|
|
current string
|
|
override string
|
|
configure bool
|
|
stallCatalog bool
|
|
extra []string
|
|
want string
|
|
wantArgs []string
|
|
}{
|
|
{name: "saved Pi default", current: "model-b", want: "model-b"},
|
|
{name: "resume restores its own model", current: "model-a", extra: []string{"--continue"}, want: "model-a", wantArgs: []string{"--continue"}},
|
|
{name: "explicit override wins", current: "model-b", override: "model-a", extra: []string{"--continue"}, want: "model-a", wantArgs: []string{"--provider", "ollama", "--model", "model-a", "--continue"}},
|
|
{name: "explicit lookup precedes timed out catalog", current: "model-b", override: "model-a", stallCatalog: true, want: "model-a", wantArgs: []string{"--provider", "ollama", "--model", "model-a"}},
|
|
{name: "explicit configuration wins", current: "model-b", configure: true, want: "model-a", wantArgs: []string{"--provider", "ollama", "--model", "model-a"}},
|
|
} {
|
|
t.Run(tc.name, func(t *testing.T) {
|
|
home := t.TempDir()
|
|
setLaunchTestHome(t, home)
|
|
withLauncherHooks(t)
|
|
withInteractiveSession(t, true)
|
|
if err := config.SaveIntegration("pi", []string{"model-a", "stale"}); err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
dir := filepath.Join(home, ".pi", "agent")
|
|
if err := os.MkdirAll(dir, 0o755); err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
if tc.stallCatalog {
|
|
oldTimeout := piModelDiscoveryTimeout
|
|
piModelDiscoveryTimeout = 50 * time.Millisecond
|
|
t.Cleanup(func() { piModelDiscoveryTimeout = oldTimeout })
|
|
if err := os.WriteFile(filepath.Join(dir, "models.json"), []byte(`{"providers":{"ollama":{"models":[{"id":"model-b","_launch":true,"contextWindow":8192}]}}}`), 0o644); err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
}
|
|
firstLookup := make(chan string, 1)
|
|
settingsPath := filepath.Join(dir, "settings.json")
|
|
settings := fmt.Sprintf(`{"defaultProvider":"ollama","defaultModel":%q,"editorPaddingX":7}`, tc.current)
|
|
if err := os.WriteFile(settingsPath, []byte(settings), 0o644); err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
bin := t.TempDir()
|
|
argsPath := filepath.Join(home, "pi-args")
|
|
for name, content := range map[string]string{
|
|
"pi": fmt.Sprintf("#!/bin/sh\nprintf '%%s\\n' \"$@\" > %q\n", argsPath),
|
|
"npm": "#!/bin/sh\nexit 1\n",
|
|
"package.json": `{"name":"@earendil-works/pi-coding-agent"}`,
|
|
} {
|
|
if err := os.WriteFile(filepath.Join(bin, name), []byte(content), 0o755); err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
}
|
|
t.Setenv("PATH", bin)
|
|
DefaultSingleSelector = func(string, []SelectionItem, string) (string, error) {
|
|
if !tc.configure {
|
|
t.Fatal("ordinary launch unexpectedly opened the picker")
|
|
}
|
|
return "model-a", nil
|
|
}
|
|
server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
|
switch r.URL.Path {
|
|
case "/api/status":
|
|
fmt.Fprint(w, `{"cloud":{"disabled":true}}`)
|
|
case "/api/tags":
|
|
if tc.stallCatalog {
|
|
select {
|
|
case firstLookup <- "tags":
|
|
default:
|
|
}
|
|
<-r.Context().Done()
|
|
return
|
|
}
|
|
fmt.Fprint(w, `{"models":[{"name":"model-a","capabilities":["completion"],"details":{"context_length":32768}},{"name":"model-b","capabilities":["completion"],"details":{"context_length":32768}}]}`)
|
|
case "/api/show":
|
|
select {
|
|
case firstLookup <- "show":
|
|
default:
|
|
}
|
|
fmt.Fprint(w, `{"capabilities":["completion"],"model_info":{"test.context_length":32768}}`)
|
|
case "/api/experimental/model-recommendations":
|
|
fmt.Fprint(w, `{"recommendations":[{"model":"model-a"}]}`)
|
|
default:
|
|
t.Errorf("unexpected request: %s", r.URL.Path)
|
|
http.NotFound(w, r)
|
|
}
|
|
}))
|
|
defer server.Close()
|
|
t.Setenv("OLLAMA_HOST", server.URL)
|
|
for range 2 {
|
|
select {
|
|
case <-firstLookup:
|
|
default:
|
|
}
|
|
err := LaunchIntegration(t.Context(), IntegrationLaunchRequest{Name: "pi", ModelOverride: tc.override, ForceConfigure: tc.configure, ExtraArgs: tc.extra})
|
|
if err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
data, err := os.ReadFile(settingsPath)
|
|
if err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
var settings map[string]any
|
|
if err := json.Unmarshal(data, &settings); err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
if settings["defaultModel"] != tc.want || settings["editorPaddingX"] != float64(7) {
|
|
t.Errorf("Pi settings = %s, want default %s and preserved padding", data, tc.want)
|
|
}
|
|
if tc.stallCatalog {
|
|
if lookup := <-firstLookup; lookup != "show" {
|
|
t.Errorf("first model lookup = %s, want show", lookup)
|
|
}
|
|
data, err := os.ReadFile(filepath.Join(dir, "models.json"))
|
|
if err != nil || !strings.Contains(string(data), `"model-b"`) {
|
|
t.Errorf("partial discovery discarded known catalog: %s, %v", data, err)
|
|
}
|
|
}
|
|
args, err := os.ReadFile(argsPath)
|
|
if err != nil || !slices.Equal(strings.Fields(string(args)), tc.wantArgs) {
|
|
t.Errorf("Pi args = %q, %v; want %v", args, err, tc.wantArgs)
|
|
}
|
|
}
|
|
})
|
|
}
|
|
}
|