From 159b1c3331a176abf7e8484e44d517db726707db Mon Sep 17 00:00:00 2001 From: Eva H <63033505+hoyyeva@users.noreply.github.com> Date: Wed, 9 Sep 2026 18:02:42 -0400 Subject: [PATCH 01/24] app: fix ChatGPT model selector spacing (#18347) --- app/ui/app/src/components/CodexDesktopModelsSettings.tsx | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/app/ui/app/src/components/CodexDesktopModelsSettings.tsx b/app/ui/app/src/components/CodexDesktopModelsSettings.tsx index 5a4da77121d..678d0e67b7e 100644 --- a/app/ui/app/src/components/CodexDesktopModelsSettings.tsx +++ b/app/ui/app/src/components/CodexDesktopModelsSettings.tsx @@ -566,7 +566,7 @@ export const CodexDesktopModelsSettings = forwardRef<
Choose ChatGPT models -
+
{selected.map((model) => ( Date: Thu, 10 Sep 2026 08:09:59 -0700 Subject: [PATCH 02/24] llama.cpp: version bump b10864 (#18317) --- LLAMA_CPP_VERSION | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/LLAMA_CPP_VERSION b/LLAMA_CPP_VERSION index 322774ed9da..f533d35b8f5 100644 --- a/LLAMA_CPP_VERSION +++ b/LLAMA_CPP_VERSION @@ -1 +1 @@ -b10760 +b10864 From ea8d65004affdae6464af1c73a55eb34f5a67eb1 Mon Sep 17 00:00:00 2001 From: Daniel Hiltgen Date: Thu, 10 Sep 2026 11:16:37 -0700 Subject: [PATCH 03/24] server: extract GGUF metadata and unify capabilities (#17858) Loading GGUF metadata is an expensive operation. Two caches had evolved to mitigate this, and the two capability implementations produced inconsistent results for some models. This PR now extracts the metadata once per blob into a file at /metadata/sha256-.json. Only arrays over 4096 elements and non-finite floats are left out. Direct Capabilities() discovery now costs us instead of ms. /api/tags can build directly from manifests and the extracted metadata. --- fs/gguf/gguf.go | 38 +- fs/gguf/keyvalue.go | 5 + manifest/manifest.go | 15 +- manifest/paths.go | 23 +- server/create.go | 6 +- server/gguf_metadata.go | 338 ++++++++++ server/gguf_metadata_test.go | 736 ++++++++++++++++++++++ server/images.go | 172 +++--- server/images_test.go | 38 +- server/model_caches.go | 7 - server/model_inference_cache.go | 121 ---- server/model_inference_cache_test.go | 121 ---- server/model_list.go | 133 ++++ server/model_list_cache.go | 888 --------------------------- server/model_list_cache_test.go | 371 ----------- server/model_list_test.go | 168 +++++ server/routes.go | 25 +- server/routes_list_test.go | 17 +- server/routes_test.go | 7 +- 19 files changed, 1550 insertions(+), 1679 deletions(-) create mode 100644 server/gguf_metadata.go create mode 100644 server/gguf_metadata_test.go delete mode 100644 server/model_inference_cache.go delete mode 100644 server/model_inference_cache_test.go create mode 100644 server/model_list.go delete mode 100644 server/model_list_cache.go delete mode 100644 server/model_list_cache_test.go create mode 100644 server/model_list_test.go diff --git a/fs/gguf/gguf.go b/fs/gguf/gguf.go index 4468a314061..d9bb876ecd8 100644 --- a/fs/gguf/gguf.go +++ b/fs/gguf/gguf.go @@ -51,12 +51,19 @@ type File struct { bts []byte } -func Open(path string) (f *File, err error) { - f = &File{bts: make([]byte, 4096)} +func Open(path string) (_ *File, err error) { + f := &File{bts: make([]byte, 4096)} f.file, err = os.Open(path) if err != nil { return nil, err } + defer func() { + if err != nil { + if closeErr := f.close(); closeErr != nil { + err = errors.Join(err, closeErr) + } + } + }() f.reader = newBufferedReader(f.file, 32<<10) @@ -273,14 +280,14 @@ func readArrayData[T any](f *File, n uint64) (s []T, err error) { return nil, err } - s = make([]T, size) - for i := range size { + s = make([]T, 0, min(size, 4096)) + for range size { e, err := read[T](f) if err != nil { return nil, err } - s[i] = e + s = append(s, e) } return s, nil @@ -292,14 +299,14 @@ func readArrayString(f *File, n uint64) (s []string, err error) { return nil, err } - s = make([]string, size) - for i := range size { + s = make([]string, 0, min(size, 4096)) + for range size { e, err := readString(f) if err != nil { return nil, err } - s[i] = e + s = append(s, e) } return s, nil @@ -324,8 +331,19 @@ func maxInt64() uint64 { } func (f *File) Close() error { - f.keyValues.stop() - f.tensors.stop() + return f.close() +} + +func (f *File) close() error { + if f.keyValues != nil { + f.keyValues.stop() + } + if f.tensors != nil { + f.tensors.stop() + } + if f.file == nil { + return nil + } return f.file.Close() } diff --git a/fs/gguf/keyvalue.go b/fs/gguf/keyvalue.go index 47420116591..f2f2b7f24d4 100644 --- a/fs/gguf/keyvalue.go +++ b/fs/gguf/keyvalue.go @@ -18,6 +18,11 @@ type Value struct { value any } +// Any returns Value as stored, without conversion. If it is not set, it returns nil. +func (v Value) Any() any { + return v.value +} + func value[T any](v Value, kinds ...reflect.Kind) (t T) { vv := reflect.ValueOf(v.value) if slices.Contains(kinds, vv.Kind()) { diff --git a/manifest/manifest.go b/manifest/manifest.go index c0277e9a572..d010ead73c3 100644 --- a/manifest/manifest.go +++ b/manifest/manifest.go @@ -71,10 +71,13 @@ func (m *Manifest) Remove() error { return PruneDirectory(manifests) } -func (m *Manifest) RemoveLayers() error { +// RemoveLayers deletes the layers no other manifest references. The digests it +// removed are returned even alongside an error, so callers can clean up +// anything derived from them. +func (m *Manifest) RemoveLayers() ([]string, error) { ms, err := Manifests(true) if err != nil { - return err + return nil, err } // Build set of digests still in use by other manifests @@ -88,6 +91,7 @@ func (m *Manifest) RemoveLayers() error { } // Remove layers not used by any other manifest + var removed []string for _, layer := range append(m.Layers, m.Config) { if layer.Digest == "" { continue @@ -97,16 +101,17 @@ func (m *Manifest) RemoveLayers() error { } blob, err := BlobsPath(layer.Digest) if err != nil { - return err + return removed, err } if err := os.Remove(blob); os.IsNotExist(err) { slog.Debug("layer does not exist", "digest", layer.Digest) } else if err != nil { - return err + return removed, err } + removed = append(removed, layer.Digest) } - return nil + return removed, nil } func ParseNamedManifest(n model.Name) (*Manifest, error) { diff --git a/manifest/paths.go b/manifest/paths.go index 4451c81aa12..df452794e3d 100644 --- a/manifest/paths.go +++ b/manifest/paths.go @@ -14,6 +14,18 @@ import ( var ErrInvalidDigestFormat = errors.New("invalid digest format") +// a manifest spells a digest with ":", a filename with "-" +var digestPattern = regexp.MustCompile(`^sha256[:-][0-9a-fA-F]{64}$`) + +// ValidateDigest reports whether digest names a blob. +func ValidateDigest(digest string) error { + if !digestPattern.MatchString(digest) { + return ErrInvalidDigestFormat + } + + return nil +} + func Path() (string, error) { path := filepath.Join(envconfig.Models(), "manifests") if err := os.MkdirAll(path, 0o755); err != nil { @@ -38,12 +50,11 @@ func PathForName(n model.Name) (string, error) { } func BlobsPath(digest string) (string, error) { - // only accept actual sha256 digests - pattern := "^sha256[:-][0-9a-fA-F]{64}$" - re := regexp.MustCompile(pattern) - - if digest != "" && !re.MatchString(digest) { - return "", ErrInvalidDigestFormat + // the empty digest names the blobs directory + if digest != "" { + if err := ValidateDigest(digest); err != nil { + return "", err + } } digest = strings.ReplaceAll(digest, ":", "-") diff --git a/server/create.go b/server/create.go index fddbe5ae0b6..f0ad013d153 100644 --- a/server/create.go +++ b/server/create.go @@ -337,13 +337,13 @@ func (s *Server) CreateHandler(c *gin.Context) { } if !envconfig.NoPrune() && oldManifest != nil { - if err := oldManifest.RemoveLayers(); err != nil { + removed, err := oldManifest.RemoveLayers() + removeGGUFMetadata(removed...) + if err != nil { ch <- gin.H{"error": err.Error()} } } - s.refreshModelListCache(name) - ch <- api.ProgressResponse{Status: "success"} }() diff --git a/server/gguf_metadata.go b/server/gguf_metadata.go new file mode 100644 index 00000000000..118fb74457d --- /dev/null +++ b/server/gguf_metadata.go @@ -0,0 +1,338 @@ +package server + +import ( + "bytes" + "encoding/json" + "errors" + "fmt" + "log/slog" + "math" + "os" + "path/filepath" + "strings" + + "github.com/ollama/ollama/envconfig" + fsgguf "github.com/ollama/ollama/fs/gguf" + "github.com/ollama/ollama/manifest" + "github.com/ollama/ollama/version" +) + +// A blob's metadata block is extracted once into a file beside the model +// store, keyed by blob digest. Values are held verbatim rather than derived, so +// reading them differently later needs no invalidation, and values left out are +// recorded by key so an absent key stays distinct from an uncopied one. + +// Comfortably above per-layer arrays, which scale with block count, and far +// below any tokenizer vocabulary. +const ggufMetadataMaxArray = 4096 + +type ggufMetadata struct { + OllamaVersion string `json:"ollama_version,omitempty"` + + // Values keyed exactly as they appear in the file. Never omitempty: absent + // on load means the file is not ours, and an all-omitted file is still + // usable. + KV map[string]any `json:"kv"` + // Keys present in the file whose values were too large to copy. + Omitted []string `json:"omitted,omitempty"` +} + +// Valid reports whether the key is present. Keys are architecture qualified the +// same way fs/gguf does it. +func (m ggufMetadata) Valid(key string) bool { + _, ok := m.lookup(key) + return ok +} + +func (m ggufMetadata) String(key string) string { + v, _ := m.lookup(key) + s, _ := v.(string) + return s +} + +func (m ggufMetadata) Int(key string) int64 { + n, ok := m.number(key) + if !ok { + return 0 + } + i, err := n.Int64() + if err != nil { + return 0 + } + return i +} + +func (m ggufMetadata) number(key string) (json.Number, bool) { + v, _ := m.lookup(key) + n, ok := v.(json.Number) + return n, ok +} + +// Keys returns every key the file carried, including omitted ones. +func (m ggufMetadata) Keys() []string { + keys := make([]string, 0, len(m.KV)+len(m.Omitted)) + for k := range m.KV { + keys = append(keys, k) + } + return append(keys, m.Omitted...) +} + +func (m ggufMetadata) lookup(key string) (any, bool) { + if !strings.HasPrefix(key, "general.") && !strings.HasPrefix(key, "tokenizer.") { + arch, _ := m.KV["general.architecture"].(string) + key = arch + "." + key + } + v, ok := m.KV[key] + return v, ok +} + +func ggufMetadataPath(digest string) (string, error) { + if err := manifest.ValidateDigest(digest); err != nil { + return "", fmt.Errorf("%w: %q", err, digest) + } + return filepath.Join(ggufMetadataDir(), strings.ReplaceAll(digest, ":", "-")+".json"), nil +} + +func ggufMetadataDir() string { + return filepath.Join(envconfig.Models(), "metadata") +} + +// readGGUFMetadata extracts the blob when there is no usable metadata file. Best +// effort throughout: a bad file is re-extracted, a failed write is dropped. +func readGGUFMetadata(digest string) (ggufMetadata, error) { + path, err := ggufMetadataPath(digest) + if err != nil { + return ggufMetadata{}, err + } + if md, ok := loadGGUFMetadata(path); ok { + return md, nil + } + + blob, err := manifest.BlobsPath(digest) + if err != nil { + return ggufMetadata{}, err + } + md, err := extractGGUFMetadata(blob) + if err != nil { + return ggufMetadata{}, err + } + + writeGGUFMetadata(path, blob, md) + return md, nil +} + +func loadGGUFMetadata(path string) (ggufMetadata, bool) { + data, err := os.ReadFile(path) + if err != nil { + if !errors.Is(err, os.ErrNotExist) { + slog.Debug("could not read gguf metadata file", "path", path, "error", err) + } + return ggufMetadata{}, false + } + + md, err := decodeGGUFMetadata(data) + if err != nil { + slog.Debug("ignoring unusable gguf metadata file", "path", path, "error", err) + return ggufMetadata{}, false + } + return md, true +} + +func writeGGUFMetadata(path, blob string, md ggufMetadata) { + data, err := json.Marshal(md) + if err != nil { + slog.Debug("could not encode gguf metadata", "path", path, "error", err) + return + } + if err := os.MkdirAll(filepath.Dir(path), 0o755); err != nil { + slog.Debug("could not create gguf metadata dir", "path", path, "error", err) + return + } + + tmp, err := os.CreateTemp(filepath.Dir(path), ".gguf-metadata-*.tmp") + if err != nil { + slog.Debug("could not create gguf metadata temp file", "path", path, "error", err) + return + } + defer os.Remove(tmp.Name()) + + if _, err = tmp.Write(data); err == nil { + err = tmp.Sync() + } + if cerr := tmp.Close(); err == nil { + err = cerr + } + if err == nil { + // Can fail transiently on Windows while another process holds the path. + err = os.Rename(tmp.Name(), path) + } + if err == nil { + // Blob removal precedes metadata removal. Rechecking after publication + // ensures either this writer or a concurrent remover cleans up the file. + if _, statErr := os.Stat(blob); errors.Is(statErr, os.ErrNotExist) { + if removeErr := os.Remove(path); removeErr != nil && !errors.Is(removeErr, os.ErrNotExist) { + slog.Debug("could not remove gguf metadata file", "path", path, "error", removeErr) + } + return + } + } + if err != nil { + slog.Debug("could not write gguf metadata file", "path", path, "error", err) + } +} + +// removeGGUFMetadata drops metadata for deleted blobs. Pass only digests whose +// last reference is gone. +func removeGGUFMetadata(digests ...string) { + for _, digest := range digests { + path, err := ggufMetadataPath(digest) + if err != nil { + continue + } + if err := os.Remove(path); err != nil && !errors.Is(err, os.ErrNotExist) { + slog.Debug("could not remove gguf metadata file", "path", path, "error", err) + } + } +} + +func pruneGGUFMetadata() { + dir := ggufMetadataDir() + entries, err := os.ReadDir(dir) + if errors.Is(err, os.ErrNotExist) { + return + } + if err != nil { + slog.Debug("could not read gguf metadata dir", "path", dir, "error", err) + return + } + + for _, entry := range entries { + name := entry.Name() + if entry.IsDir() || !strings.HasSuffix(name, ".json") { + continue + } + + digest := strings.TrimSuffix(name, ".json") + blob, err := manifest.BlobsPath(digest) + if err != nil { + continue + } + if _, err := os.Stat(blob); err == nil || !errors.Is(err, os.ErrNotExist) { + continue + } + path := filepath.Join(dir, name) + if err := os.Remove(path); err != nil && !errors.Is(err, os.ErrNotExist) { + slog.Debug("could not remove gguf metadata file", "path", path, "error", err) + } + } +} + +// extractGGUFMetadata reads the whole block; proving a key absent means reaching +// the end. Normalized through JSON so extracted and loaded metadata hold the +// same value types. +func extractGGUFMetadata(path string) (ggufMetadata, error) { + md, err := scanGGUFMetadata(path) + if err != nil { + return ggufMetadata{}, err + } + md.OllamaVersion = version.Version + + data, err := json.Marshal(md) + if err != nil { + return ggufMetadata{}, err + } + return decodeGGUFMetadata(data) +} + +func decodeGGUFMetadata(data []byte) (ggufMetadata, error) { + // UseNumber keeps integers exact; the default would widen them to float64. + dec := json.NewDecoder(bytes.NewReader(data)) + dec.UseNumber() + var md ggufMetadata + if err := dec.Decode(&md); err != nil { + return ggufMetadata{}, err + } + if md.KV == nil { + return ggufMetadata{}, fmt.Errorf("no metadata") + } + return md, nil +} + +func scanGGUFMetadata(path string) (ggufMetadata, error) { + f, err := fsgguf.Open(path) + if err != nil { + return ggufMetadata{}, err + } + defer f.Close() + + md := ggufMetadata{KV: make(map[string]any)} + for _, kv := range f.KeyValues() { + if omitValue(kv.Any()) { + md.Omitted = append(md.Omitted, kv.Key) + continue + } + md.KV[kv.Key] = kv.Any() + } + if err := f.Err(); err != nil { + return ggufMetadata{}, err + } + return md, nil +} + +// omitValue reports values left out of the metadata file: long arrays, which are the +// tokenizer, and non-finite floats, which JSON cannot represent. +func omitValue(v any) bool { + switch v := v.(type) { + case float32: + return !finite(float64(v)) + case float64: + return !finite(v) + case []float32: + for _, f := range v { + if !finite(float64(f)) { + return true + } + } + case []float64: + for _, f := range v { + if !finite(f) { + return true + } + } + } + return oversized(v) +} + +func finite(f float64) bool { + return !math.IsInf(f, 0) && !math.IsNaN(f) +} + +func oversized(v any) bool { + switch v := v.(type) { + case []string: + return len(v) > ggufMetadataMaxArray + case []int8: + return len(v) > ggufMetadataMaxArray + case []int16: + return len(v) > ggufMetadataMaxArray + case []int32: + return len(v) > ggufMetadataMaxArray + case []int64: + return len(v) > ggufMetadataMaxArray + case []uint8: + return len(v) > ggufMetadataMaxArray + case []uint16: + return len(v) > ggufMetadataMaxArray + case []uint32: + return len(v) > ggufMetadataMaxArray + case []uint64: + return len(v) > ggufMetadataMaxArray + case []float32: + return len(v) > ggufMetadataMaxArray + case []float64: + return len(v) > ggufMetadataMaxArray + case []bool: + return len(v) > ggufMetadataMaxArray + } + return false +} diff --git a/server/gguf_metadata_test.go b/server/gguf_metadata_test.go new file mode 100644 index 00000000000..f62951e35f5 --- /dev/null +++ b/server/gguf_metadata_test.go @@ -0,0 +1,736 @@ +package server + +import ( + "encoding/binary" + "encoding/json" + "errors" + "math" + "net/http" + "os" + "path/filepath" + "reflect" + "runtime" + "slices" + "strings" + "testing" + "time" + + "github.com/gin-gonic/gin" + + "github.com/ollama/ollama/api" + "github.com/ollama/ollama/envconfig" + "github.com/ollama/ollama/fs/ggml" + "github.com/ollama/ollama/manifest" + "github.com/ollama/ollama/types/model" +) + +func TestGGUFMetadataExtraction(t *testing.T) { + t.Setenv("OLLAMA_MODELS", t.TempDir()) + + big := make([]string, 5000) + for i := range big { + big[i] = "tok" + } + perLayer := make([]int32, 64) + for i := range perLayer { + perLayer[i] = int32(i) + } + + path, _ := createBinFile(t, ggml.KV{ + "general.architecture": "bert", + "bert.attention.softcap": float32(math.Inf(-1)), + "general.file_type": uint32(2), + "bert.pooling_type": uint32(1), + "bert.context_length": uint32(512), + "bert.attention.head_count_kv": perLayer, + "tokenizer.chat_template": "{{ messages }}", + "tokenizer.ggml.tokens": big, + "tokenizer.ggml.merges": big, + }, nil) + + md, err := extractGGUFMetadata(path) + if err != nil { + t.Fatal(err) + } + + if got := md.String("general.architecture"); got != "bert" { + t.Errorf("architecture = %q", got) + } + // Architecture-prefixed lookup, and the key order in the file must not matter. + if !md.Valid("pooling_type") { + t.Error("pooling_type not found") + } + if got := md.Int("context_length"); got != 512 { + t.Errorf("context_length = %d, want 512", got) + } + if got := md.String("tokenizer.chat_template"); got != "{{ messages }}" { + t.Errorf("chat_template = %q", got) + } + if md.Valid("vision.block_count") { + t.Error("absent key reported present") + } + + // A 64-element per-layer array is real metadata and must survive. + if v, ok := md.KV["bert.attention.head_count_kv"].([]any); !ok || len(v) != 64 { + t.Errorf("per-layer array = %#v, want 64 elements", md.KV["bert.attention.head_count_kv"]) + } + + // The tokenizer is recorded as omitted, not silently missing. + want := map[string]bool{"tokenizer.ggml.tokens": true, "tokenizer.ggml.merges": true} + for _, key := range md.Omitted { + delete(want, key) + } + if len(want) != 0 { + t.Errorf("not reported omitted: %v (omitted=%v)", want, md.Omitted) + } + if _, ok := md.KV["tokenizer.ggml.tokens"]; ok { + t.Error("oversized array was copied into the metadata file") + } + + // JSON cannot represent an infinity, and gemma3n ships one. Dropping that + // key must not cost the rest of the file. + if !slices.Contains(md.Omitted, "bert.attention.softcap") { + t.Errorf("non-finite float not omitted; omitted=%v", md.Omitted) + } +} + +func TestGGUFMetadataFileRoundTrip(t *testing.T) { + t.Setenv("OLLAMA_MODELS", t.TempDir()) + + _, digest := createBinFile(t, ggml.KV{ + "general.architecture": "llama", + "llama.block_count": uint32(32), + }, nil) + + first, err := readGGUFMetadata(digest) + if err != nil { + t.Fatal(err) + } + if got := first.Int("block_count"); got != 32 { + t.Fatalf("block_count = %d, want 32", got) + } + + // With the blob gone, only the metadata file can answer. + if err := removeBlob(t, digest); err != nil { + t.Fatal(err) + } + second, err := readGGUFMetadata(digest) + if err != nil { + t.Fatalf("metadata file read failed: %v", err) + } + if got := second.Int("block_count"); got != 32 { + t.Fatalf("block_count from metadata file = %d, want 32", got) + } + + // And it goes away with the blob. + removeGGUFMetadata(digest) + path, err := ggufMetadataPath(digest) + if err != nil { + t.Fatal(err) + } + if _, err := os.Stat(path); !os.IsNotExist(err) { + t.Fatalf("metadata file still present after removal: %v", err) + } +} + +// The metadata file is only ever an optimization: anything unusable must be re-extracted. +func TestGGUFMetadataUnusableFile(t *testing.T) { + t.Setenv("OLLAMA_MODELS", t.TempDir()) + + _, digest := createBinFile(t, ggml.KV{"general.architecture": "llama"}, nil) + if _, err := readGGUFMetadata(digest); err != nil { + t.Fatal(err) + } + path, err := ggufMetadataPath(digest) + if err != nil { + t.Fatal(err) + } + + for _, tt := range []struct { + name string + content []byte + }{ + {"truncated", []byte(`{"kv":{"general.arch`)}, + {"not json", []byte("\x00\x01\x02")}, + {"empty", []byte{}}, + {"wrong shape", []byte(`[1,2,3]`)}, + {"no kv", []byte(`{"ollama_version":"1.2.3"}`)}, + {"kv not an object", []byte(`{"kv":[]}`)}, + } { + t.Run(tt.name, func(t *testing.T) { + if err := os.WriteFile(path, tt.content, 0o644); err != nil { + t.Fatal(err) + } + if _, ok := loadGGUFMetadata(path); ok { + t.Fatal("unusable metadata file reported usable") + } + md, err := readGGUFMetadata(digest) + if err != nil { + t.Fatal(err) + } + if got := md.String("general.architecture"); got != "llama" { + t.Fatalf("architecture = %q, want llama", got) + } + }) + } +} + +func TestGGUFMetadataRejectsMalformed(t *testing.T) { + t.Setenv("OLLAMA_MODELS", t.TempDir()) + dir := t.TempDir() + + for _, tt := range []struct { + name string + data []byte + }{ + {"empty file", nil}, + {"bad magic", []byte("NOPE")}, + {"truncated header", []byte("GGUF\x03\x00\x00\x00")}, + } { + t.Run(tt.name, func(t *testing.T) { + path := filepath.Join(dir, tt.name) + if err := os.WriteFile(path, tt.data, 0o600); err != nil { + t.Fatal(err) + } + if _, err := extractGGUFMetadata(path); err == nil { + t.Fatal("expected an error") + } + }) + } +} + +// The omission rule decides what never reaches the metadata file, so every array type +// has to agree on it. A miss here silently drops real metadata. +func TestGGUFMetadataOmitRule(t *testing.T) { + short, long := 64, ggufMetadataMaxArray+1 + + for _, tt := range []struct { + name string + keep any + omit any + }{ + {"string", make([]string, short), make([]string, long)}, + {"int8", make([]int8, short), make([]int8, long)}, + {"int16", make([]int16, short), make([]int16, long)}, + {"int32", make([]int32, short), make([]int32, long)}, + {"int64", make([]int64, short), make([]int64, long)}, + {"uint8", make([]uint8, short), make([]uint8, long)}, + {"uint16", make([]uint16, short), make([]uint16, long)}, + {"uint32", make([]uint32, short), make([]uint32, long)}, + {"uint64", make([]uint64, short), make([]uint64, long)}, + {"float32", make([]float32, short), make([]float32, long)}, + {"float64", make([]float64, short), make([]float64, long)}, + {"bool", make([]bool, short), make([]bool, long)}, + } { + t.Run(tt.name, func(t *testing.T) { + if omitValue(tt.keep) { + t.Errorf("%d-element %s array omitted; real per-layer metadata would be lost", short, tt.name) + } + if !omitValue(tt.omit) { + t.Errorf("%d-element %s array kept; the tokenizer would be copied", long, tt.name) + } + }) + } + + // Scalars are always kept, except the floats JSON cannot encode. + for _, keep := range []any{"s", int32(0), uint32(0), float32(1.5), float64(1.5), true} { + if omitValue(keep) { + t.Errorf("scalar %#v omitted", keep) + } + } + for _, omit := range []any{ + float32(math.Inf(1)), float32(math.Inf(-1)), float64(math.NaN()), + []float32{1, float32(math.Inf(1))}, + []float64{math.NaN()}, + } { + if !omitValue(omit) { + t.Errorf("non-finite %#v kept; json.Marshal would fail the whole file", omit) + } + } +} + +func TestGGUFMetadataKeysIncludesOmitted(t *testing.T) { + t.Setenv("OLLAMA_MODELS", t.TempDir()) + + big := make([]string, ggufMetadataMaxArray+1) + path, _ := createBinFile(t, ggml.KV{ + "general.architecture": "llama", + "llama.block_count": uint32(4), + "tokenizer.ggml.tokens": big, + }, nil) + + md, err := extractGGUFMetadata(path) + if err != nil { + t.Fatal(err) + } + keys := md.Keys() + for _, want := range []string{"general.architecture", "llama.block_count", "tokenizer.ggml.tokens"} { + if !slices.Contains(keys, want) { + t.Errorf("Keys() = %v, missing %s", keys, want) + } + } +} + +// Projector audio detection reads through Keys(), including prefixed forms. +func TestProjectorAudioFromMetadata(t *testing.T) { + t.Setenv("OLLAMA_MODELS", t.TempDir()) + + for _, tt := range []struct { + name string + kv ggml.KV + audio bool + }{ + {"bare key", ggml.KV{"general.architecture": "clip", "has_audio_encoder": true}, true}, + {"prefixed key", ggml.KV{"general.architecture": "clip", "clip.has_audio_encoder": true}, true}, + {"present but false", ggml.KV{"general.architecture": "clip", "has_audio_encoder": false}, false}, + {"absent", ggml.KV{"general.architecture": "clip"}, false}, + } { + t.Run(tt.name, func(t *testing.T) { + path, _ := createBinFile(t, tt.kv, nil) + md, err := extractGGUFMetadata(path) + if err != nil { + t.Fatal(err) + } + if got := projectorHasAudio(md); got != tt.audio { + t.Errorf("projectorHasAudio = %v, want %v", got, tt.audio) + } + }) + } + + path, _ := createBinFile(t, ggml.KV{ + "general.architecture": "clip", + "has_audio_encoder": true, + "clip.vision.projector_type": "gemma3nv", + }, nil) + md, err := extractGGUFMetadata(path) + if err != nil { + t.Fatal(err) + } + if !projectorSuppressesAudioCapability(md) { + t.Error("gemma3nv projector should suppress audio") + } +} + +func TestGGUFMetadataPathRejectsBadDigest(t *testing.T) { + t.Setenv("OLLAMA_MODELS", t.TempDir()) + hex := "0123456789abcdef0123456789abcdef0123456789abcdef0123456789abcdef" + + for _, digest := range []string{"", "sha256-short", "md5:" + hex, "../escape", "sha256:" + hex + "x"} { + if _, err := ggufMetadataPath(digest); err == nil { + t.Errorf("ggufMetadataPath(%q) accepted", digest) + } + } + for _, digest := range []string{"sha256:" + hex, "sha256-" + hex} { + if _, err := ggufMetadataPath(digest); err != nil { + t.Errorf("ggufMetadataPath(%q) rejected: %v", digest, err) + } + } +} + +// A file can carry no values at all, and every value it does carry can be one +// the omission rule drops. Either way what is left still says which keys the +// file had, so it has to be storable: treating it as corrupt would rescan the +// blob on every request forever. +func TestGGUFMetadataWithNothingToCopy(t *testing.T) { + t.Setenv("OLLAMA_MODELS", t.TempDir()) + + path := filepath.Join(t.TempDir(), "empty.gguf") + if err := os.WriteFile(path, ggufBytes(0, nil), 0o600); err != nil { + t.Fatal(err) + } + + md, err := extractGGUFMetadata(path) + if err != nil { + t.Fatalf("a file with no metadata is still a file: %v", err) + } + if len(md.KV) != 0 { + t.Errorf("KV = %v, want empty", md.KV) + } + + encoded, err := json.Marshal(md) + if err != nil { + t.Fatal(err) + } + if _, err := decodeGGUFMetadata(encoded); err != nil { + t.Errorf("written metadata cannot be read back, so it would be rescanned forever: %v", err) + } +} + +const allocBudget = 1 << 20 + +func ggufBytes(kvCount uint64, extra []byte) []byte { + b := make([]byte, 0, 24+len(extra)) + b = binary.LittleEndian.AppendUint32(b, 0x46554747) // GGUF + b = binary.LittleEndian.AppendUint32(b, 3) // version + b = binary.LittleEndian.AppendUint64(b, 0) // tensor count + b = binary.LittleEndian.AppendUint64(b, kvCount) // kv count + return append(b, extra...) +} + +func declaredGGUF(t *testing.T, kvCount uint64, extra []byte) string { + t.Helper() + b := ggufBytes(kvCount, extra) + + path := filepath.Join(t.TempDir(), "declared.gguf") + if err := os.WriteFile(path, b, 0o600); err != nil { + t.Fatal(err) + } + t.Logf("file is %d bytes, declares %d KV entries", len(b), kvCount) + return path +} + +func allocDelta(t *testing.T, fn func()) uint64 { + t.Helper() + var before, after runtime.MemStats + runtime.GC() + runtime.ReadMemStats(&before) + fn() + runtime.ReadMemStats(&after) + return after.TotalAlloc - before.TotalAlloc +} + +func TestGGUFMetadataBoundsDeclaredKVCount(t *testing.T) { + path := declaredGGUF(t, 1<<20, nil) + got := allocDelta(t, func() { + if _, err := extractGGUFMetadata(path); err == nil { + t.Error("expected an error from a 24-byte file") + } + }) + t.Logf("allocated %d bytes before failing", got) + if got > allocBudget { + t.Errorf("allocated %d bytes from a 24-byte file, want at most %d", got, allocBudget) + } +} + +// Cover every array element type because strings and fixed-width values use +// separate readers. +func TestGGUFMetadataBoundsDeclaredArray(t *testing.T) { + for _, tt := range []struct { + name string + elemType uint32 + }{ + {"uint8", 0}, + {"int8", 1}, + {"uint16", 2}, + {"int16", 3}, + {"uint32", 4}, + {"int32", 5}, + {"float32", 6}, + {"bool", 7}, + {"string", 8}, + {"uint64", 10}, + {"int64", 11}, + {"float64", 12}, + } { + t.Run(tt.name, func(t *testing.T) { + var extra []byte + extra = binary.LittleEndian.AppendUint64(extra, 3) // key length + extra = append(extra, "big"...) + extra = binary.LittleEndian.AppendUint32(extra, 9) // type: array + extra = binary.LittleEndian.AppendUint32(extra, tt.elemType) + extra = binary.LittleEndian.AppendUint64(extra, 8_000_000) // element count + + path := declaredGGUF(t, 1, extra) + got := allocDelta(t, func() { + if _, err := extractGGUFMetadata(path); err == nil { + t.Error("expected an error") + } + }) + if got > allocBudget { + t.Errorf("allocated %d bytes for a declared array of 8000000, want at most %d", got, allocBudget) + } + }) + } +} + +func removeBlob(t *testing.T, digest string) error { + t.Helper() + path, err := manifest.BlobsPath(digest) + if err != nil { + t.Fatal(err) + } + return os.Remove(path) +} + +func metadataCount(t *testing.T) int { + t.Helper() + entries, err := os.ReadDir(filepath.Join(envconfig.Models(), "metadata")) + if err != nil { + return 0 + } + return len(entries) +} + +// createModelFromBlob and deleteModelNamed drive the handlers rather than +// writing manifests, so the metadata hooks under test are the ones the server +// actually runs. +func createModelFromBlob(t *testing.T, name, digest, tmpl string) { + t.Helper() + var s Server + w := createRequest(t, s.CreateHandler, api.CreateRequest{ + Model: name, Files: map[string]string{"model.gguf": digest}, + Template: tmpl, Stream: &stream, + }) + if w.Code != http.StatusOK { + t.Fatalf("create %s: %d %s", name, w.Code, w.Body.String()) + } +} + +func deleteModelNamed(t *testing.T, name string) { + t.Helper() + var s Server + w := createRequest(t, s.DeleteHandler, api.DeleteRequest{Model: name}) + if w.Code != http.StatusOK { + t.Fatalf("delete %s: %d %s", name, w.Code, w.Body.String()) + } +} + +// Metadata is derived from a blob, so it lives exactly as long as the blob: it +// survives while any manifest still references it and goes with the last one. +func TestGGUFMetadataRemovedWithLastReference(t *testing.T) { + gin.SetMode(gin.TestMode) + t.Setenv("OLLAMA_MODELS", t.TempDir()) + + _, digest := createBinFile(t, ggml.KV{ + "general.architecture": "llama", + "llama.block_count": uint32(1), + }, nil) + + createModelFromBlob(t, "alias-one", digest, "") + createModelFromBlob(t, "alias-two", digest, "") + + if _, err := readGGUFMetadata(digest); err != nil { + t.Fatal(err) + } + if got := metadataCount(t); got != 1 { + t.Fatalf("metadata files = %d, want 1 (both aliases share one blob)", got) + } + + deleteModelNamed(t, "alias-one") + if got := metadataCount(t); got != 1 { + t.Fatalf("metadata files after removing one alias = %d, want 1", got) + } + + deleteModelNamed(t, "alias-two") + if got := metadataCount(t); got != 0 { + t.Fatalf("metadata files after removing the last reference = %d, want 0", got) + } +} + +// Recreating a model over the same name replaces its blob; the old blob's +// metadata must go with it. +func TestGGUFMetadataRemovedOnReplacement(t *testing.T) { + gin.SetMode(gin.TestMode) + t.Setenv("OLLAMA_MODELS", t.TempDir()) + + _, first := createBinFile(t, ggml.KV{ + "general.architecture": "llama", + "llama.block_count": uint32(1), + }, nil) + createModelFromBlob(t, "replace-me", first, "") + if _, err := readGGUFMetadata(first); err != nil { + t.Fatal(err) + } + + _, second := createBinFile(t, ggml.KV{ + "general.architecture": "llama", + "llama.block_count": uint32(2), + }, nil) + createModelFromBlob(t, "replace-me", second, "") + if _, err := readGGUFMetadata(second); err != nil { + t.Fatal(err) + } + + if got := metadataCount(t); got != 1 { + t.Fatalf("metadata files after replacement = %d, want 1 (the old blob's should be gone)", got) + } + old, err := ggufMetadataPath(first) + if err != nil { + t.Fatal(err) + } + if _, err := os.Stat(old); !os.IsNotExist(err) { + t.Errorf("metadata for the replaced blob still present: %v", err) + } +} + +// A blob missing from disk is still a blob whose metadata has to go. It no +// longer appears in the blob scan, so pruning has to find its metadata directly. +func TestGGUFMetadataRemovedForMissingBlob(t *testing.T) { + t.Setenv("OLLAMA_MODELS", t.TempDir()) + + _, orphan := createBinFile(t, ggml.KV{"general.architecture": "llama"}, nil) + if _, err := readGGUFMetadata(orphan); err != nil { + t.Fatal(err) + } + if err := removeBlob(t, orphan); err != nil { + t.Fatal(err) + } + + _, live := createBinFile(t, ggml.KV{"general.architecture": "bert"}, nil) + if _, err := readGGUFMetadata(live); err != nil { + t.Fatal(err) + } + if got := metadataCount(t); got != 2 { + t.Fatalf("metadata files before the sweep = %d, want 2", got) + } + + if err := PruneLayers(); err != nil { + t.Fatal(err) + } + if got := metadataCount(t); got != 1 { + t.Errorf("metadata files = %d, want 1", got) + } + path, err := ggufMetadataPath(live) + if err != nil { + t.Fatal(err) + } + if _, err := os.Stat(path); err != nil { + t.Errorf("metadata for existing blob was removed: %v", err) + } +} + +func TestGGUFMetadataNotPublishedAfterDelete(t *testing.T) { + gin.SetMode(gin.TestMode) + t.Setenv("OLLAMA_MODELS", t.TempDir()) + + _, digest := createBinFile(t, ggml.KV{ + "general.architecture": "llama", + "general.description": strings.Repeat("x", 16<<20), + }, nil) + createModelFromBlob(t, "delete-during-load", digest, "") + + loaded := make(chan error, 1) + go func() { + _, err := GetModel("delete-during-load") + loaded <- err + }() + + deadline := time.Now().Add(5 * time.Second) + for { + entries, err := os.ReadDir(ggufMetadataDir()) + if err != nil && !errors.Is(err, os.ErrNotExist) { + t.Fatal(err) + } + if slices.ContainsFunc(entries, func(entry os.DirEntry) bool { + return strings.HasPrefix(entry.Name(), ".gguf-metadata-") + }) { + break + } + select { + case err := <-loaded: + t.Fatalf("GetModel completed before its metadata write could overlap deletion: %v", err) + default: + } + if time.Now().After(deadline) { + t.Fatal("timed out waiting for metadata write") + } + runtime.Gosched() + } + + deleteModelNamed(t, "delete-during-load") + select { + case err := <-loaded: + if err != nil { + t.Fatalf("GetModel: %v", err) + } + case <-time.After(5 * time.Second): + t.Fatal("timed out waiting for GetModel") + } + + path, err := ggufMetadataPath(digest) + if err != nil { + t.Fatal(err) + } + if _, err := os.Stat(path); !errors.Is(err, os.ErrNotExist) { + t.Fatalf("metadata published after model deletion: %v", err) + } +} + +// The point of the metadata file is that loading a model no longer reads its +// blob. With the blob deleted, GetModel must still describe the model exactly +// as before - which is also the guard against anything reintroducing a +// per-request read of the model file. +func TestGetModelReadsNoBlob(t *testing.T) { + gin.SetMode(gin.TestMode) + t.Setenv("OLLAMA_MODELS", t.TempDir()) + + _, digest := createBinFile(t, ggml.KV{ + "general.architecture": "bert", + "bert.pooling_type": uint32(1), + "bert.context_length": uint32(512), + "bert.embedding_length": uint32(384), + }, nil) + createModelFromBlob(t, "embedder", digest, "") + + want, err := GetModel("embedder") + if err != nil { + t.Fatal(err) + } + if !slices.Contains(want.Capabilities(), model.CapabilityEmbedding) { + t.Fatalf("capabilities = %v, want embedding", want.Capabilities()) + } + + if err := removeBlob(t, digest); err != nil { + t.Fatal(err) + } + got, err := GetModel("embedder") + if err != nil { + t.Fatalf("GetModel needed the blob: %v", err) + } + if !slices.Equal(got.Capabilities(), want.Capabilities()) { + t.Errorf("capabilities without the blob = %v, want %v", got.Capabilities(), want.Capabilities()) + } + if got.metadata.Int("context_length") != 512 { + t.Errorf("context_length = %d, want 512", got.metadata.Int("context_length")) + } +} + +// Extraction either fails or produces metadata that survives a metadata-file +// round trip unchanged. +func FuzzGGUFMetadata(f *testing.F) { + f.Add(ggufBytes(0, nil)) + f.Add(ggufBytes(1<<20, nil)) + + var str []byte + str = binary.LittleEndian.AppendUint64(str, 20) + str = append(str, "general.architecture"...) + str = binary.LittleEndian.AppendUint32(str, 8) // type: string + str = binary.LittleEndian.AppendUint64(str, 5) + str = append(str, "llama"...) + f.Add(ggufBytes(1, str)) + + var arr []byte + arr = binary.LittleEndian.AppendUint64(arr, 3) + arr = append(arr, "big"...) + arr = binary.LittleEndian.AppendUint32(arr, 9) // type: array + arr = binary.LittleEndian.AppendUint32(arr, 8) // element type: string + arr = binary.LittleEndian.AppendUint64(arr, 8_000_000) + f.Add(ggufBytes(1, arr)) + + path := filepath.Join(f.TempDir(), "fuzz.gguf") + f.Fuzz(func(t *testing.T, data []byte) { + if err := os.WriteFile(path, data, 0o600); err != nil { + t.Fatal(err) + } + + md, err := extractGGUFMetadata(path) + if err != nil { + return + } + if md.KV == nil { + t.Fatal("extraction reported success with no metadata") + } + + encoded, err := json.Marshal(md) + if err != nil { + t.Fatalf("extracted metadata cannot be written: %v", err) + } + loaded, err := decodeGGUFMetadata(encoded) + if err != nil { + t.Fatalf("extracted metadata cannot be read back: %v", err) + } + if !reflect.DeepEqual(md, loaded) { + t.Errorf("metadata changed on the trip through a file:\n%#v\n%#v", md, loaded) + } + }) +} diff --git a/server/images.go b/server/images.go index 9be4ce51dd1..43503059063 100644 --- a/server/images.go +++ b/server/images.go @@ -10,7 +10,6 @@ import ( "io" "log" "log/slog" - "math" "net" "net/http" "net/url" @@ -24,7 +23,6 @@ import ( "github.com/ollama/ollama/api" "github.com/ollama/ollama/envconfig" - "github.com/ollama/ollama/fs/gguf" "github.com/ollama/ollama/manifest" "github.com/ollama/ollama/model/parsers" "github.com/ollama/ollama/parser" @@ -83,8 +81,10 @@ type Model struct { Template *template.Template - capabilities []model.Capability - capabilitiesCached bool + // Metadata of the model blob and of each projector, read from their + // metadata files when the model is loaded. + metadata ggufMetadata + projectorMetadata []ggufMetadata } func (m *Model) IsMLX() bool { @@ -95,47 +95,33 @@ func (m *Model) isGGUF() bool { return m.Config.ModelFormat == "" || m.Config.ModelFormat == "gguf" } -func generationDefaultsFromGGUF(f *gguf.File) model.GenerationDefaults { +func generationDefaultsFromMetadata(md ggufMetadata) model.GenerationDefaults { return model.ParseGGUFGenerationDefaults( func(key string) (int64, bool) { - return ggufIntGenerationDefault(f.KeyValue(key)) + n, ok := md.number(key) + if !ok { + return 0, false + } + if value, err := n.Int64(); err == nil { + return value, true + } + value, err := n.Float64() + if err != nil { + return 0, false + } + return int64(value), true }, func(key string) (float64, bool) { - return ggufFloatGenerationDefault(f.KeyValue(key)) + n, ok := md.number(key) + if !ok { + return 0, false + } + value, err := n.Float64() + return value, err == nil }, ) } -func ggufIntGenerationDefault(kv gguf.KeyValue) (int64, bool) { - if value, ok := kv.IntOK(); ok { - return value, true - } - if value, ok := kv.UintOK(); ok { - if value > math.MaxInt64 { - return 0, false - } - return int64(value), true - } - if value, ok := kv.FloatOK(); ok { - // Match api.Options.FromMap; rounding may be better for near-integers. - return int64(value), true - } - return 0, false -} - -func ggufFloatGenerationDefault(kv gguf.KeyValue) (float64, bool) { - if value, ok := kv.FloatOK(); ok { - return value, true - } - if value, ok := kv.IntOK(); ok { - return float64(value), true - } - if value, ok := kv.UintOK(); ok { - return float64(value), true - } - return 0, false -} - func appendCapability(capabilities []model.Capability, capability model.Capability) []model.Capability { if slices.Contains(capabilities, capability) { return capabilities @@ -153,11 +139,7 @@ const ( // Capabilities returns the capabilities that the model supports func (m *Model) Capabilities() []model.Capability { - if m.capabilitiesCached { - return slices.Clone(m.capabilities) - } - - capabilities := m.capabilitiesForTemplate(templateCapabilitySelected, nil) + capabilities := m.capabilitiesForTemplate(templateCapabilitySelected) if len(capabilities) == 0 { slog.Warn("unknown capabilities for model", "model", m.Name) } @@ -165,12 +147,12 @@ func (m *Model) Capabilities() []model.Capability { return capabilities } -func (m *Model) capabilitiesForTemplate(source templateCapabilitySource, f *gguf.File) []model.Capability { +func (m *Model) capabilitiesForTemplate(source templateCapabilitySource) []model.Capability { capabilities := []model.Capability{} var modelArch string capabilities = m.configCapabilities(capabilities) - capabilities, modelArch = m.ggufCapabilities(capabilities, source, f) + capabilities, modelArch = m.ggufCapabilities(capabilities, source) capabilities = m.projectorCapabilities(capabilities) capabilities = m.templateCapabilities(capabilities, source) capabilities = m.parserCapabilities(capabilities) @@ -187,44 +169,33 @@ func (m *Model) configCapabilities(capabilities []model.Capability) []model.Capa return capabilities } -func (m *Model) ggufCapabilities(capabilities []model.Capability, source templateCapabilitySource, f *gguf.File) ([]model.Capability, string) { +func (m *Model) ggufCapabilities(capabilities []model.Capability, source templateCapabilitySource) ([]model.Capability, string) { if m.ModelPath == "" || !m.isGGUF() { return capabilities, "" } - if f == nil { - var err error - f, err = gguf.Open(m.ModelPath) - if err != nil { - slog.Error("couldn't open model file", "error", err) - return capabilities, "" - } - defer f.Close() - } - - modelArch := f.KeyValue("general.architecture").String() switch source { case templateCapabilitySelected: if !usesOllamaRenderedChat(m) { - capabilities = chatTemplateCapabilities(capabilities, f.KeyValue("tokenizer.chat_template").String()) + capabilities = chatTemplateCapabilities(capabilities, m.metadata.String("tokenizer.chat_template")) } case templateCapabilityChat: - capabilities = chatTemplateCapabilities(capabilities, f.KeyValue("tokenizer.chat_template").String()) + capabilities = chatTemplateCapabilities(capabilities, m.metadata.String("tokenizer.chat_template")) } - if f.KeyValue("pooling_type").Valid() { + if m.metadata.Valid("pooling_type") { capabilities = appendCapability(capabilities, model.CapabilityEmbedding) } else { // If no embedding is specified, we assume the model supports completion. capabilities = appendCapability(capabilities, model.CapabilityCompletion) } - if f.KeyValue("vision.block_count").Valid() { + if m.metadata.Valid("vision.block_count") { capabilities = appendCapability(capabilities, model.CapabilityVision) } - if f.KeyValue("audio.block_count").Valid() { + if m.metadata.Valid("audio.block_count") { capabilities = appendCapability(capabilities, model.CapabilityAudio) } - return capabilities, modelArch + return capabilities, m.metadata.String("general.architecture") } func chatTemplateCapabilities(capabilities []model.Capability, chatTemplate string) []model.Capability { @@ -390,28 +361,17 @@ func capabilityLogValue(present bool, capabilities []model.Capability) any { } func (m *Model) templateSelectionCapabilities(usesHarmony bool) (goTemplate, chatTemplate, harmony, rendererParser []model.Capability) { - var f *gguf.File - if m.ModelPath != "" && m.isGGUF() { - var err error - f, err = gguf.Open(m.ModelPath) - if err != nil { - slog.Error("couldn't open model file", "error", err) - } else { - defer f.Close() - } - } - if m.HasGoTemplate { - goTemplate = m.capabilitiesForTemplate(templateCapabilityGo, f) + goTemplate = m.capabilitiesForTemplate(templateCapabilityGo) } if m.HasChatTemplate { - chatTemplate = m.capabilitiesForTemplate(templateCapabilityChat, f) + chatTemplate = m.capabilitiesForTemplate(templateCapabilityChat) } if usesHarmony { - harmony = m.capabilitiesForTemplate(templateCapabilitySelected, f) + harmony = m.capabilitiesForTemplate(templateCapabilitySelected) } if m.Config.Renderer != "" || m.Config.Parser != "" { - rendererParser = m.capabilitiesForTemplate(templateCapabilitySelected, f) + rendererParser = m.capabilitiesForTemplate(templateCapabilitySelected) } return goTemplate, chatTemplate, harmony, rendererParser @@ -439,16 +399,10 @@ func (m *Model) projectorCapabilities(capabilities []model.Capability) []model.C } capabilities = appendCapability(capabilities, model.CapabilityVision) - for _, projectorPath := range m.ProjectorPaths { - f, err := gguf.Open(projectorPath) - if err != nil { - slog.Error("couldn't open projector file", "error", err) - continue - } - if projectorHasAudio(f) && !projectorSuppressesAudioCapability(f) { + for _, md := range m.projectorMetadata { + if projectorHasAudio(md) && !projectorSuppressesAudioCapability(md) { capabilities = appendCapability(capabilities, model.CapabilityAudio) } - f.Close() } return capabilities @@ -551,22 +505,21 @@ func isNemotron3NanoSafetensorsConfig(cfg model.ConfigV2) bool { slices.Contains(cfg.ModelFamilies, "nemotron_h_omni")) } -func projectorHasAudio(f *gguf.File) bool { - if f.KeyValue("has_audio_encoder").Bool() { - return true - } - - for _, kv := range f.KeyValues() { - if strings.HasSuffix(kv.Key, ".has_audio_encoder") && kv.Bool() { - return true +func projectorHasAudio(md ggufMetadata) bool { + // read directly: Keys reports qualified keys, the accessors qualify theirs + for _, key := range md.Keys() { + if key == "has_audio_encoder" || strings.HasSuffix(key, ".has_audio_encoder") { + if b, ok := md.KV[key].(bool); ok && b { + return true + } } } return false } -func projectorSuppressesAudioCapability(f *gguf.File) bool { - switch f.KeyValue("vision.projector_type").String() { +func projectorSuppressesAudioCapability(md ggufMetadata) bool { + switch md.String("vision.projector_type") { case "gemma3nv": return true } @@ -745,6 +698,13 @@ func GetModel(name string) (*Model, error) { modelHasPooling := false ggufChatTemplate := "" for _, layer := range mf.Layers { + // Nothing below reads a tensor layer, and resolving a path costs a + // syscall each. Named rather than allowlisting the types below, so a new + // layer type is slower here instead of silently unread. + if layer.MediaType == manifest.MediaTypeImageTensor { + continue + } + filename, err := manifest.BlobsPath(layer.Digest) if err != nil { return nil, err @@ -755,16 +715,16 @@ func GetModel(name string) (*Model, error) { m.ModelPath = filename m.ParentModel = layer.From if m.isGGUF() { - f, err := gguf.Open(filename) + md, err := readGGUFMetadata(layer.Digest) if err != nil { - slog.Error("couldn't open model file", "error", err) + slog.Error("couldn't read model metadata", "error", err) break } - ggufChatTemplate = f.KeyValue("tokenizer.chat_template").String() + m.metadata = md + ggufChatTemplate = md.String("tokenizer.chat_template") m.HasChatTemplate = ggufChatTemplate != "" - modelHasPooling = f.KeyValue("pooling_type").Valid() - m.GenerationDefaults = generationDefaultsFromGGUF(f) - f.Close() + modelHasPooling = md.Valid("pooling_type") + m.GenerationDefaults = generationDefaultsFromMetadata(md) } case manifest.MediaTypeImageDraft: m.DraftPath = filename @@ -776,6 +736,11 @@ func GetModel(name string) (*Model, error) { m.AdapterPaths = append(m.AdapterPaths, filename) case "application/vnd.ollama.image.projector": m.ProjectorPaths = append(m.ProjectorPaths, filename) + if md, err := readGGUFMetadata(layer.Digest); err != nil { + slog.Error("couldn't read projector metadata", "error", err) + } else { + m.projectorMetadata = append(m.projectorMetadata, md) + } case "application/vnd.ollama.image.prompt", "application/vnd.ollama.image.template": m.HasGoTemplate = true @@ -900,10 +865,11 @@ func deleteUnusedLayers(deleteMap map[string]struct{}) error { slog.Info(fmt.Sprintf("couldn't get file path for '%s': %v", k, err)) continue } - if err := os.Remove(fp); err != nil { + if err := os.Remove(fp); err != nil && !errors.Is(err, os.ErrNotExist) { slog.Info(fmt.Sprintf("couldn't remove file '%s': %v", fp, err)) continue } + removeGGUFMetadata(k) } return nil @@ -960,6 +926,7 @@ func PruneLayers() error { slog.Error(fmt.Sprintf("couldn't remove unused layers: %v", err)) return nil } + pruneGGUFMetadata() slog.Info(fmt.Sprintf("total unused blobs removed: %d", len(deleteMap))) @@ -1122,6 +1089,7 @@ func PullModel(ctx context.Context, name string, regOpts *registryOptions, fn fu if err := os.Remove(fp); err != nil { slog.Info(fmt.Sprintf("couldn't remove file with digest mismatch '%s': %v", fp, err)) } + removeGGUFMetadata(layer.Digest) } return err } diff --git a/server/images_test.go b/server/images_test.go index fd521dcbd79..a0302bbee84 100644 --- a/server/images_test.go +++ b/server/images_test.go @@ -14,7 +14,6 @@ import ( "github.com/ollama/ollama/api" "github.com/ollama/ollama/fs/ggml" - fsgguf "github.com/ollama/ollama/fs/gguf" "github.com/ollama/ollama/manifest" "github.com/ollama/ollama/template" "github.com/ollama/ollama/types/model" @@ -61,7 +60,7 @@ func TestPruneLayersSkipsRecentOrphans(t *testing.T) { } } -func TestGenerationDefaultsFromGGUF(t *testing.T) { +func TestGenerationDefaultsFromMetadata(t *testing.T) { file, err := os.CreateTemp(t.TempDir(), "model-*.gguf") if err != nil { t.Fatal(err) @@ -87,13 +86,12 @@ func TestGenerationDefaultsFromGGUF(t *testing.T) { t.Fatal(err) } - f, err := fsgguf.Open(file.Name()) + md, err := extractGGUFMetadata(file.Name()) if err != nil { t.Fatal(err) } - defer f.Close() - defaults := generationDefaultsFromGGUF(f) + defaults := generationDefaultsFromMetadata(md) check := func(key string, want any) { t.Helper() if got := defaults[key]; got != want { @@ -104,10 +102,10 @@ func TestGenerationDefaultsFromGGUF(t *testing.T) { check("top_k", int64(40)) check("top_p", float64(1)) check("min_p", float64(0)) - check("typical_p", float64(float32(0.95))) + check("typical_p", float64(0.95)) check("temperature", float64(1)) check("repeat_last_n", int64(64)) - check("repeat_penalty", float64(float32(1.05))) + check("repeat_penalty", float64(1.05)) check("frequency_penalty", float64(0)) check("presence_penalty", float64(0)) if _, ok := defaults["mirostat_tau"]; ok { @@ -354,6 +352,28 @@ func writeTestModelManifest(t *testing.T, name, digest, tmpl string) { } } +// loadTestMetadata fills in what GetModel would have read from the metadata +// files, so +// hand-built models resolve capabilities the same way loaded ones do. +func loadTestMetadata(t *testing.T, m *Model) { + t.Helper() + if m.ModelPath != "" { + md, err := extractGGUFMetadata(m.ModelPath) + if err != nil { + t.Fatalf("metadata for %s: %v", m.ModelPath, err) + } + m.metadata = md + } + m.projectorMetadata = nil + for _, path := range m.ProjectorPaths { + md, err := extractGGUFMetadata(path) + if err != nil { + t.Fatalf("projector metadata for %s: %v", path, err) + } + m.projectorMetadata = append(m.projectorMetadata, md) + } +} + func TestModelCapabilities(t *testing.T) { // Create completion model (llama architecture without vision) completionModelPath, _ := createBinFile(t, ggml.KV{ @@ -618,6 +638,8 @@ func TestModelCapabilities(t *testing.T) { for _, tt := range testModels { t.Run(tt.name, func(t *testing.T) { + loadTestMetadata(t, &tt.model) + // Test Capabilities method caps := tt.model.Capabilities() if !compareCapabilities(caps, tt.expectedCaps) { @@ -748,6 +770,8 @@ func TestModelCheckCapabilities(t *testing.T) { for _, tt := range tests { t.Run(tt.name, func(t *testing.T) { + loadTestMetadata(t, &tt.model) + // Test CheckCapabilities method err := tt.model.CheckCapabilities(tt.checkCaps...) if tt.expectedErrMsg == "" { diff --git a/server/model_caches.go b/server/model_caches.go index 06f9e52b972..720fa703d49 100644 --- a/server/model_caches.go +++ b/server/model_caches.go @@ -5,16 +5,12 @@ import "context" type modelCaches struct { recommendations *modelRecommendationsCache show *modelShowCache - modelList *modelListCache - inference *inferenceModelCache } func newModelCaches() *modelCaches { return &modelCaches{ recommendations: newModelRecommendationsCache(), show: newModelShowCache(), - modelList: newModelListCache(), - inference: newInferenceModelCache(), } } @@ -28,7 +24,4 @@ func (c *modelCaches) Start(ctx context.Context) { if c.show != nil { c.show.Start(ctx) } - if c.modelList != nil { - c.modelList.Start(ctx) - } } diff --git a/server/model_inference_cache.go b/server/model_inference_cache.go deleted file mode 100644 index a5edcbd44c3..00000000000 --- a/server/model_inference_cache.go +++ /dev/null @@ -1,121 +0,0 @@ -package server - -import ( - "maps" - "slices" - "strconv" - "sync" - - "github.com/ollama/ollama/envconfig" - "github.com/ollama/ollama/manifest" - "github.com/ollama/ollama/types/model" - "golang.org/x/sync/singleflight" -) - -// inferenceModelCache stores fully resolved model metadata and capabilities. -// Model blobs are content-addressed, and the manifest digest is the freshness -// boundary, so cached entries remain valid until the model is recreated or -// pulled with different content. -type inferenceModelCache struct { - mu sync.RWMutex - entries map[inferenceModelCacheKey]inferenceModelCacheEntry - loads singleflight.Group - - loadModel func(string) (*Model, error) -} - -type inferenceModelCacheKey struct { - name string - goTemplate bool - goTemplateSet bool -} - -type inferenceModelCacheEntry struct { - digest string - model *Model -} - -func newInferenceModelCache() *inferenceModelCache { - return &inferenceModelCache{ - entries: make(map[inferenceModelCacheKey]inferenceModelCacheEntry), - loadModel: GetModel, - } -} - -func (c *inferenceModelCache) Get(name string) (*Model, error) { - n := model.ParseName(name) - mf, err := manifest.ParseNamedManifest(n) - if err != nil { - return nil, err - } - - key := inferenceModelCacheKey{ - name: n.String(), - goTemplate: envconfig.GoTemplate(true), - goTemplateSet: goTemplateEnvSet(), - } - digest := mf.Digest() - - c.mu.RLock() - entry, ok := c.entries[key] - c.mu.RUnlock() - if ok && entry.digest == digest { - return cloneInferenceModel(entry.model), nil - } - - loadKey := key.name + "\x00" + digest + "\x00" + strconv.FormatBool(key.goTemplate) + "\x00" + strconv.FormatBool(key.goTemplateSet) - v, err, _ := c.loads.Do(loadKey, func() (any, error) { - c.mu.RLock() - entry, ok := c.entries[key] - c.mu.RUnlock() - if ok && entry.digest == digest { - return entry.model, nil - } - - m, err := c.loadModel(name) - if err != nil { - return nil, err - } - m.capabilities = m.Capabilities() - m.capabilitiesCached = true - - c.mu.Lock() - c.entries[key] = inferenceModelCacheEntry{digest: m.Digest, model: m} - c.mu.Unlock() - return m, nil - }) - if err != nil { - return nil, err - } - - return cloneInferenceModel(v.(*Model)), nil -} - -func cloneInferenceModel(src *Model) *Model { - if src == nil { - return nil - } - - dst := *src - dst.Config.ModelFamilies = slices.Clone(src.Config.ModelFamilies) - dst.Config.Capabilities = slices.Clone(src.Config.Capabilities) - if src.Config.Draft != nil { - draft := *src.Config.Draft - dst.Config.Draft = &draft - } - dst.AdapterPaths = slices.Clone(src.AdapterPaths) - dst.ProjectorPaths = slices.Clone(src.ProjectorPaths) - dst.License = slices.Clone(src.License) - dst.Options = maps.Clone(src.Options) - dst.Messages = slices.Clone(src.Messages) - dst.capabilities = slices.Clone(src.capabilities) - - return &dst -} - -func (s *Server) getModel(name string) (*Model, error) { - if s != nil && s.modelCaches != nil && s.modelCaches.inference != nil { - return s.modelCaches.inference.Get(name) - } - return GetModel(name) -} diff --git a/server/model_inference_cache_test.go b/server/model_inference_cache_test.go deleted file mode 100644 index fb53202e2a8..00000000000 --- a/server/model_inference_cache_test.go +++ /dev/null @@ -1,121 +0,0 @@ -package server - -import ( - "slices" - "sync" - "sync/atomic" - "testing" - "time" - - "github.com/ollama/ollama/fs/ggml" - "github.com/ollama/ollama/types/model" -) - -func TestInferenceModelCache(t *testing.T) { - t.Setenv("OLLAMA_MODELS", t.TempDir()) - t.Setenv("OLLAMA_GO_TEMPLATE", "") - - _, completionDigest := createBinFile(t, ggml.KV{ - "general.architecture": "llama", - }, nil) - writeTestModelManifest(t, "inference-cache", completionDigest, "{{ .Prompt }}") - - cache := newInferenceModelCache() - loadCount := 0 - cache.loadModel = func(name string) (*Model, error) { - loadCount++ - return GetModel(name) - } - - first, err := cache.Get("inference-cache") - if err != nil { - t.Fatal(err) - } - if loadCount != 1 { - t.Fatalf("load count = %d, want 1", loadCount) - } - if !first.capabilitiesCached { - t.Fatal("capabilities were not cached") - } - if got := first.Capabilities(); !slices.Contains(got, model.CapabilityCompletion) { - t.Fatalf("capabilities = %v, want completion", got) - } - - // Returned models are request-local clones. Mutating one must not alter the - // cached model or a later request. - first.Config.Parser = "mutated" - first.Config.ModelFamilies[0] = "mutated" - first.capabilities[0] = model.CapabilityImage - - second, err := cache.Get("inference-cache") - if err != nil { - t.Fatal(err) - } - if loadCount != 1 { - t.Fatalf("cache hit load count = %d, want 1", loadCount) - } - if second.Config.Parser == "mutated" || slices.Contains(second.Config.ModelFamilies, "mutated") { - t.Fatalf("cached model was mutated: %#v", second.Config) - } - if got := second.Capabilities(); !slices.Contains(got, model.CapabilityCompletion) || slices.Contains(got, model.CapabilityImage) { - t.Fatalf("cached capabilities = %v, want completion only", got) - } - - // Recreating the manifest changes its digest and invalidates the entry. - _, embeddingDigest := createBinFile(t, ggml.KV{ - "general.architecture": "bert", - "bert.pooling_type": uint32(1), - }, nil) - writeTestModelManifest(t, "inference-cache", embeddingDigest, "{{ .Prompt }}") - - third, err := cache.Get("inference-cache") - if err != nil { - t.Fatal(err) - } - if loadCount != 2 { - t.Fatalf("invalidated load count = %d, want 2", loadCount) - } - if got := third.Capabilities(); !slices.Contains(got, model.CapabilityEmbedding) || slices.Contains(got, model.CapabilityCompletion) { - t.Fatalf("refreshed capabilities = %v, want embedding only", got) - } -} - -func TestInferenceModelCacheConcurrentMiss(t *testing.T) { - t.Setenv("OLLAMA_MODELS", t.TempDir()) - t.Setenv("OLLAMA_GO_TEMPLATE", "") - - _, digest := createBinFile(t, ggml.KV{ - "general.architecture": "llama", - }, nil) - writeTestModelManifest(t, "inference-cache-concurrent", digest, "{{ .Prompt }}") - - cache := newInferenceModelCache() - var loadCount atomic.Int32 - cache.loadModel = func(name string) (*Model, error) { - loadCount.Add(1) - time.Sleep(10 * time.Millisecond) - return GetModel(name) - } - - var wg sync.WaitGroup - errs := make(chan error, 8) - for range 8 { - wg.Add(1) - go func() { - defer wg.Done() - _, err := cache.Get("inference-cache-concurrent") - errs <- err - }() - } - wg.Wait() - close(errs) - - for err := range errs { - if err != nil { - t.Fatal(err) - } - } - if got := loadCount.Load(); got != 1 { - t.Fatalf("load count = %d, want 1", got) - } -} diff --git a/server/model_list.go b/server/model_list.go new file mode 100644 index 00000000000..237e1b63cf6 --- /dev/null +++ b/server/model_list.go @@ -0,0 +1,133 @@ +package server + +import ( + "cmp" + "context" + "encoding/json" + "log/slog" + "slices" + "time" + + "github.com/ollama/ollama/api" + "github.com/ollama/ollama/fs/ggml" + "github.com/ollama/ollama/manifest" + "github.com/ollama/ollama/types/model" +) + +// listModels builds /api/tags from the manifests and the per-blob metadata +// files, extracting for any blob that has none yet. +func listModels(ctx context.Context) ([]api.ListModelResponse, error) { + manifests, err := manifest.Manifests(true) + if err != nil { + return nil, err + } + + models := make([]api.ListModelResponse, 0, len(manifests)) + for name, mf := range manifests { + if ctx != nil { + if err := ctx.Err(); err != nil { + return nil, err + } + } + + summary, err := describeModel(name, mf) + if err != nil { + slog.Warn("failed to describe model", "model", name.String(), "error", err) + continue + } + models = append(models, summary) + } + + sortListModelResponses(models) + return models, nil +} + +// describeModel describes one model for /api/tags. Capabilities come from the +// same Model.Capabilities() the inference path uses, so the two cannot drift. +func describeModel(name model.Name, mf *manifest.Manifest) (api.ListModelResponse, error) { + cfg, err := readModelListConfig(mf) + if err != nil { + return api.ListModelResponse{}, err + } + + var modified time.Time + if fi := mf.FileInfo(); fi != nil { + modified = fi.ModTime() + } + + summary := api.ListModelResponse{ + Model: name.DisplayShortest(), + Name: name.DisplayShortest(), + RemoteModel: cfg.RemoteModel, + RemoteHost: cfg.RemoteHost, + Size: mf.Size(), + Digest: mf.Digest(), + ModifiedAt: modified, + Details: api.ModelDetails{ + Format: cfg.ModelFormat, + Family: cfg.ModelFamily, + Families: append([]string(nil), cfg.ModelFamilies...), + ParameterSize: cfg.ModelType, + QuantizationLevel: cfg.FileType, + ContextLength: cfg.ContextLen, + EmbeddingLength: cfg.EmbedLen, + }, + } + + m, err := GetModel(name.String()) + if err != nil { + // A model that will not load is the one a user most needs to see, in + // order to remove it. Report what the manifest says. + slog.Warn("could not load model to describe it", "model", name.String(), "error", err) + return summary, nil + } + summary.Details.ParentModel = m.ParentModel + summary.Capabilities = m.Capabilities() + + if m.ModelPath != "" && m.isGGUF() { + if summary.Details.ContextLength == 0 { + summary.Details.ContextLength = int(m.metadata.Int("context_length")) + } + if summary.Details.EmbeddingLength == 0 { + summary.Details.EmbeddingLength = int(m.metadata.Int("embedding_length")) + } + if m.metadata.Valid("general.file_type") { + fileType := ggml.FileType(m.metadata.Int("general.file_type")).String() + if isUnknownQuantization(summary.Details.QuantizationLevel) && !isUnknownQuantization(fileType) { + summary.Details.QuantizationLevel = fileType + } + } + } + + return summary, nil +} + +func readModelListConfig(mf *manifest.Manifest) (model.ConfigV2, error) { + var cfg model.ConfigV2 + if mf == nil || mf.Config.Digest == "" { + return cfg, nil + } + + f, err := mf.Config.Open() + if err != nil { + return cfg, err + } + defer f.Close() + + if err := json.NewDecoder(f).Decode(&cfg); err != nil { + return cfg, err + } + + return cfg, nil +} + +func isUnknownQuantization(quantization string) bool { + return quantization == "" || quantization == "unknown" +} + +func sortListModelResponses(models []api.ListModelResponse) { + slices.SortStableFunc(models, func(i, j api.ListModelResponse) int { + // Preserve the existing /api/tags order: most recently modified first. + return cmp.Compare(j.ModifiedAt.Unix(), i.ModifiedAt.Unix()) + }) +} diff --git a/server/model_list_cache.go b/server/model_list_cache.go deleted file mode 100644 index 9017aae7059..00000000000 --- a/server/model_list_cache.go +++ /dev/null @@ -1,888 +0,0 @@ -package server - -import ( - "bufio" - "cmp" - "context" - "encoding/binary" - "encoding/json" - "fmt" - "io" - "log/slog" - "os" - "slices" - "strings" - "sync" - "time" - - "github.com/ollama/ollama/api" - "github.com/ollama/ollama/fs/ggml" - fsgguf "github.com/ollama/ollama/fs/gguf" - "github.com/ollama/ollama/manifest" - "github.com/ollama/ollama/model/parsers" - ollamatemplate "github.com/ollama/ollama/template" - "github.com/ollama/ollama/thinking" - "github.com/ollama/ollama/types/model" -) - -type modelListSummary struct { - Model string - Name string - RemoteModel string - RemoteHost string - Size int64 - Digest string - ModifiedAt time.Time - Details api.ModelDetails - Capabilities []model.Capability -} - -type modelListCacheEntry struct { - Digest string - Summary modelListSummary -} - -type modelListCache struct { - mu sync.RWMutex - - entries map[string]modelListCacheEntry - - once sync.Once - readyOnce sync.Once - ready chan struct{} - hydrateErr error - build func(model.Name, *manifest.Manifest) (modelListSummary, error) -} - -func newModelListCache() *modelListCache { - return &modelListCache{ - entries: make(map[string]modelListCacheEntry), - ready: make(chan struct{}), - build: buildModelListSummary, - } -} - -func (c *modelListCache) Start(ctx context.Context) { - if c == nil { - return - } - - c.once.Do(func() { - slog.Debug("starting model list cache") - go func() { - err := c.hydrate(ctx) - c.markReady(err) - if err != nil { - if ctx != nil && ctx.Err() != nil { - return - } - slog.Warn("model list cache hydration failed", "error", err) - } - }() - }) -} - -func (c *modelListCache) hydrate(ctx context.Context) error { - start := time.Now() - - manifests, err := manifest.Manifests(true) - if err != nil { - return err - } - - var hydrated, failed int - for name, mf := range manifests { - if ctx != nil { - if err := ctx.Err(); err != nil { - return err - } - } - - summary, err := c.build(name, mf) - if err != nil { - failed++ - slog.Warn("failed to hydrate model list cache", "model", name.String(), "error", err) - continue - } - - c.set(name, mf.Digest(), summary) - hydrated++ - } - - slog.Info("model list cache hydration complete", "models", hydrated, "failures", failed, "elapsed", time.Since(start)) - return nil -} - -func (c *modelListCache) markReady(err error) { - c.mu.Lock() - c.hydrateErr = err - c.mu.Unlock() - - c.readyOnce.Do(func() { - close(c.ready) - }) -} - -func (c *modelListCache) Wait(ctx context.Context) error { - if c == nil { - return nil - } - if ctx == nil { - ctx = context.Background() - } - - select { - case <-c.ready: - c.mu.RLock() - err := c.hydrateErr - c.mu.RUnlock() - return err - case <-ctx.Done(): - return ctx.Err() - } -} - -func (c *modelListCache) List(ctx context.Context) ([]api.ListModelResponse, error) { - if err := c.Wait(ctx); err != nil { - return nil, err - } - if err := c.syncManifests(ctx); err != nil { - return nil, err - } - - c.mu.RLock() - models := make([]api.ListModelResponse, 0, len(c.entries)) - for _, entry := range c.entries { - models = append(models, entry.Summary.ListModelResponse()) - } - c.mu.RUnlock() - - sortListModelResponses(models) - return models, nil -} - -func (c *modelListCache) syncManifests(ctx context.Context) error { - manifests, err := manifest.Manifests(true) - if err != nil { - return err - } - - c.mu.RLock() - current := make(map[string]string, len(c.entries)) - for name, entry := range c.entries { - current[name] = entry.Digest - } - c.mu.RUnlock() - - type update struct { - name model.Name - digest string - summary modelListSummary - } - - seen := make(map[string]struct{}, len(manifests)) - stale := make(map[string]struct{}) - var updates []update - for name, mf := range manifests { - if ctx != nil { - if err := ctx.Err(); err != nil { - return err - } - } - - key := name.String() - digest := mf.Digest() - seen[key] = struct{}{} - if current[key] == digest { - continue - } - - summary, err := c.build(name, mf) - if err != nil { - slog.Warn("failed to refresh model list cache", "model", key, "error", err) - if _, ok := current[key]; ok { - stale[key] = struct{}{} - } - continue - } - updates = append(updates, update{name: name, digest: digest, summary: summary}) - } - - c.mu.Lock() - for name := range c.entries { - if _, ok := seen[name]; !ok { - delete(c.entries, name) - continue - } - if _, ok := stale[name]; ok { - delete(c.entries, name) - } - } - for _, update := range updates { - c.entries[update.name.String()] = modelListCacheEntry{ - Digest: update.digest, - Summary: cloneModelListSummary(update.summary), - } - } - c.mu.Unlock() - - return nil -} - -func (c *modelListCache) RefreshModel(name model.Name) error { - if c == nil { - return nil - } - - if !name.IsFullyQualified() { - var err error - name, err = getExistingName(name) - if err != nil { - return err - } - } - - mf, err := manifest.ParseNamedManifest(name) - if err != nil { - c.DeleteModel(name) - return err - } - - summary, err := c.build(name, mf) - if err != nil { - c.DeleteModel(name) - return err - } - - c.set(name, mf.Digest(), summary) - return nil -} - -func (c *modelListCache) DeleteModel(name model.Name) { - if c == nil { - return - } - - c.mu.Lock() - delete(c.entries, name.String()) - c.mu.Unlock() -} - -func (c *modelListCache) Get(name model.Name) (modelListSummary, bool) { - if c == nil { - return modelListSummary{}, false - } - - if !name.IsFullyQualified() { - if existing, err := getExistingName(name); err == nil { - name = existing - } - } - - c.mu.RLock() - entry, ok := c.entries[name.String()] - c.mu.RUnlock() - if !ok { - return modelListSummary{}, false - } - - return cloneModelListSummary(entry.Summary), true -} - -func (c *modelListCache) Len() int { - if c == nil { - return 0 - } - - c.mu.RLock() - defer c.mu.RUnlock() - return len(c.entries) -} - -func (c *modelListCache) set(name model.Name, digest string, summary modelListSummary) { - c.mu.Lock() - c.entries[name.String()] = modelListCacheEntry{ - Digest: digest, - Summary: cloneModelListSummary(summary), - } - c.mu.Unlock() -} - -func buildModelListSummary(name model.Name, mf *manifest.Manifest) (modelListSummary, error) { - cfg, err := readModelListConfig(mf) - if err != nil { - return modelListSummary{}, err - } - - var modified time.Time - if fi := mf.FileInfo(); fi != nil { - modified = fi.ModTime() - } - - summary := modelListSummary{ - Model: name.DisplayShortest(), - Name: name.DisplayShortest(), - RemoteModel: cfg.RemoteModel, - RemoteHost: cfg.RemoteHost, - Size: mf.Size(), - Digest: mf.Digest(), - ModifiedAt: modified, - Details: api.ModelDetails{ - Format: cfg.ModelFormat, - Family: cfg.ModelFamily, - Families: append([]string(nil), cfg.ModelFamilies...), - ParameterSize: cfg.ModelType, - QuantizationLevel: cfg.FileType, - ContextLength: cfg.ContextLen, - EmbeddingLength: cfg.EmbedLen, - }, - } - - modelPath, projectorCount, tmpl, err := readModelListLayers(mf, &summary) - if err != nil { - return modelListSummary{}, err - } - - if cfg.RemoteHost == "" && cfg.RemoteModel == "" && modelPath != "" { - info, err := readModelListGGUF(modelPath) - if err != nil { - slog.Debug("failed to read gguf model metadata", "model", name.String(), "error", err) - } else { - summary.Capabilities = appendModelListCapabilities(summary.Capabilities, info.Capabilities...) - if summary.Details.ContextLength == 0 { - summary.Details.ContextLength = info.ContextLength - } - if summary.Details.EmbeddingLength == 0 { - summary.Details.EmbeddingLength = info.EmbeddingLength - } - if isUnknownQuantization(summary.Details.QuantizationLevel) && !isUnknownQuantization(info.FileType) { - summary.Details.QuantizationLevel = info.FileType - } - } - } - - for _, c := range cfg.Capabilities { - summary.Capabilities = appendModelListCapability(summary.Capabilities, model.Capability(c)) - } - - builtinParser := parsers.ParserForName(cfg.Parser) - if tmpl != nil { - vars, err := tmpl.Vars() - if err != nil { - slog.Warn("model template contains errors", "model", name.String(), "error", err) - } - if slices.Contains(vars, "tools") || (builtinParser != nil && builtinParser.HasToolSupport()) { - summary.Capabilities = appendModelListCapability(summary.Capabilities, model.CapabilityTools) - } - if slices.Contains(vars, "suffix") { - summary.Capabilities = appendModelListCapability(summary.Capabilities, model.CapabilityInsert) - } - - openingTag, closingTag := thinking.InferTags(tmpl.Template) - hasTags := openingTag != "" && closingTag != "" - isGptoss := slices.Contains([]string{"gptoss", "gpt-oss"}, cfg.ModelFamily) - if !slices.Contains(summary.Capabilities, model.CapabilityThinking) && - (hasTags || isGptoss || (builtinParser != nil && builtinParser.HasThinkingSupport())) { - summary.Capabilities = appendModelListCapability(summary.Capabilities, model.CapabilityThinking) - } - } - - if projectorCount > 0 { - summary.Capabilities = appendModelListCapability(summary.Capabilities, model.CapabilityVision) - } - - summary.Capabilities = filterUnsupportedModelListCapabilities(summary.Capabilities, cfg) - - return summary, nil -} - -func filterUnsupportedModelListCapabilities(capabilities []model.Capability, cfg model.ConfigV2) []model.Capability { - if cfg.ModelFormat == "safetensors" && isNemotron3NanoSafetensorsConfig(cfg) { - capabilities = slices.DeleteFunc(capabilities, func(c model.Capability) bool { - return c == model.CapabilityVision || c == model.CapabilityAudio - }) - } - // Mirrors suppressAudioCapability in images.go so /api/tags and /api/show - // agree for safetensors models whose MLX runner serves vision but not audio. - if cfg.ModelFormat == "safetensors" && cfg.Renderer == "glimmer" { - capabilities = slices.DeleteFunc(capabilities, func(c model.Capability) bool { - return c == model.CapabilityAudio - }) - } - - return capabilities -} - -func readModelListConfig(mf *manifest.Manifest) (model.ConfigV2, error) { - var cfg model.ConfigV2 - if mf == nil || mf.Config.Digest == "" { - return cfg, nil - } - - f, err := mf.Config.Open() - if err != nil { - return cfg, err - } - defer f.Close() - - if err := json.NewDecoder(f).Decode(&cfg); err != nil { - return cfg, err - } - - return cfg, nil -} - -func readModelListLayers(mf *manifest.Manifest, summary *modelListSummary) (string, int, *ollamatemplate.Template, error) { - var modelPath string - var projectorCount int - tmpl := ollamatemplate.DefaultTemplate - - for _, layer := range mf.Layers { - switch layer.MediaType { - case "application/vnd.ollama.image.model": - filename, err := manifest.BlobsPath(layer.Digest) - if err != nil { - return "", 0, nil, err - } - modelPath = filename - summary.Details.ParentModel = layer.From - case "application/vnd.ollama.image.projector": - projectorCount++ - case "application/vnd.ollama.image.prompt", - "application/vnd.ollama.image.template": - filename, err := manifest.BlobsPath(layer.Digest) - if err != nil { - return "", 0, nil, err - } - bts, err := os.ReadFile(filename) - if err != nil { - return "", 0, nil, err - } - - tmpl, err = ollamatemplate.Parse(string(bts)) - if err != nil { - return "", 0, nil, err - } - } - } - - return modelPath, projectorCount, tmpl, nil -} - -type modelListGGUF struct { - Capabilities []model.Capability - ContextLength int - EmbeddingLength int - FileType string -} - -const ( - modelListGGUFMagicLE = 0x46554747 - modelListGGUFMagicBE = 0x47475546 -) - -const ( - modelListGGUFTypeUint8 uint32 = iota - modelListGGUFTypeInt8 - modelListGGUFTypeUint16 - modelListGGUFTypeInt16 - modelListGGUFTypeUint32 - modelListGGUFTypeInt32 - modelListGGUFTypeFloat32 - modelListGGUFTypeBool - modelListGGUFTypeString - modelListGGUFTypeArray - modelListGGUFTypeUint64 - modelListGGUFTypeInt64 - modelListGGUFTypeFloat64 -) - -// readModelListGGUF scans only the small GGUF header values launch needs -// and stops before tokenizer arrays. Using gguf.File.KeyValue for missing keys -// can otherwise advance through large arrays just to discover absence. -func readModelListGGUF(path string) (modelListGGUF, error) { - f, err := os.Open(path) - if err != nil { - return modelListGGUF{}, err - } - defer f.Close() - - r := bufio.NewReaderSize(f, 32<<10) - var magic uint32 - if err := binary.Read(r, binary.LittleEndian, &magic); err != nil { - return modelListGGUF{}, err - } - - var byteOrder binary.ByteOrder = binary.LittleEndian - switch magic { - case modelListGGUFMagicLE: - case modelListGGUFMagicBE: - byteOrder = binary.BigEndian - default: - return modelListGGUF{}, fmt.Errorf("invalid file magic") - } - - var version uint32 - if err := binary.Read(r, byteOrder, &version); err != nil { - return modelListGGUF{}, err - } - - var numKV uint64 - switch version { - case 1: - var header struct { - NumTensor uint32 - NumKV uint32 - } - if err := binary.Read(r, byteOrder, &header); err != nil { - return modelListGGUF{}, err - } - numKV = uint64(header.NumKV) - default: - var header struct { - NumTensor uint64 - NumKV uint64 - } - if err := binary.Read(r, byteOrder, &header); err != nil { - return modelListGGUF{}, err - } - numKV = header.NumKV - } - - info := modelListGGUF{} - var architecture string - var hasPoolingType bool - - for range numKV { - key, err := readModelListGGUFString(r, byteOrder, version) - if err != nil { - return modelListGGUF{}, err - } - - var valueType uint32 - if err := binary.Read(r, byteOrder, &valueType); err != nil { - return modelListGGUF{}, err - } - - if key == "general.architecture" { - value, err := readModelListGGUFStringValue(r, byteOrder, version, valueType) - if err != nil { - return modelListGGUF{}, err - } - architecture = value - continue - } - - if key == "general.file_type" { - value, err := readModelListGGUFIntValue(r, byteOrder, version, valueType) - if err != nil { - return modelListGGUF{}, err - } - info.FileType = ggml.FileType(value).String() - continue - } - - if architecture != "" && strings.HasPrefix(key, "tokenizer.") { - break - } - - if architecture != "" && strings.HasPrefix(key, architecture+".") { - switch strings.TrimPrefix(key, architecture+".") { - case "pooling_type": - hasPoolingType = true - case "vision.block_count": - info.Capabilities = appendModelListCapability(info.Capabilities, model.CapabilityVision) - case "audio.block_count": - info.Capabilities = appendModelListCapability(info.Capabilities, model.CapabilityAudio) - case "context_length": - value, err := readModelListGGUFIntValue(r, byteOrder, version, valueType) - if err != nil { - return modelListGGUF{}, err - } - info.ContextLength = value - continue - case "embedding_length": - value, err := readModelListGGUFIntValue(r, byteOrder, version, valueType) - if err != nil { - return modelListGGUF{}, err - } - info.EmbeddingLength = value - continue - } - } - - if err := skipModelListGGUFValue(r, byteOrder, version, valueType); err != nil { - return modelListGGUF{}, err - } - } - - if hasPoolingType { - info.Capabilities = appendModelListCapability(info.Capabilities, model.CapabilityEmbedding) - } else { - info.Capabilities = appendModelListCapability(info.Capabilities, model.CapabilityCompletion) - } - - return info, nil -} - -func readModelListGGUFStringValue(r io.Reader, byteOrder binary.ByteOrder, version uint32, valueType uint32) (string, error) { - if valueType != modelListGGUFTypeString { - if err := skipModelListGGUFValue(r, byteOrder, version, valueType); err != nil { - return "", err - } - return "", fmt.Errorf("unexpected gguf string type %d", valueType) - } - return readModelListGGUFString(r, byteOrder, version) -} - -func readModelListGGUFIntValue(r io.Reader, byteOrder binary.ByteOrder, version uint32, valueType uint32) (int, error) { - switch valueType { - case modelListGGUFTypeUint8: - var value uint8 - if err := binary.Read(r, byteOrder, &value); err != nil { - return 0, err - } - return int(value), nil - case modelListGGUFTypeInt8: - var value int8 - if err := binary.Read(r, byteOrder, &value); err != nil { - return 0, err - } - return int(value), nil - case modelListGGUFTypeUint16: - var value uint16 - if err := binary.Read(r, byteOrder, &value); err != nil { - return 0, err - } - return int(value), nil - case modelListGGUFTypeInt16: - var value int16 - if err := binary.Read(r, byteOrder, &value); err != nil { - return 0, err - } - return int(value), nil - case modelListGGUFTypeUint32: - var value uint32 - if err := binary.Read(r, byteOrder, &value); err != nil { - return 0, err - } - return int(value), nil - case modelListGGUFTypeInt32: - var value int32 - if err := binary.Read(r, byteOrder, &value); err != nil { - return 0, err - } - return int(value), nil - case modelListGGUFTypeUint64: - var value uint64 - if err := binary.Read(r, byteOrder, &value); err != nil { - return 0, err - } - return int(value), nil - case modelListGGUFTypeInt64: - var value int64 - if err := binary.Read(r, byteOrder, &value); err != nil { - return 0, err - } - return int(value), nil - default: - if err := skipModelListGGUFValue(r, byteOrder, version, valueType); err != nil { - return 0, err - } - return 0, fmt.Errorf("unexpected gguf integer type %d", valueType) - } -} - -func skipModelListGGUFValue(r io.Reader, byteOrder binary.ByteOrder, version uint32, valueType uint32) error { - switch valueType { - case modelListGGUFTypeUint8, modelListGGUFTypeInt8, modelListGGUFTypeBool: - return discardModelListGGUFBytes(r, 1) - case modelListGGUFTypeUint16, modelListGGUFTypeInt16: - return discardModelListGGUFBytes(r, 2) - case modelListGGUFTypeUint32, modelListGGUFTypeInt32, modelListGGUFTypeFloat32: - return discardModelListGGUFBytes(r, 4) - case modelListGGUFTypeUint64, modelListGGUFTypeInt64, modelListGGUFTypeFloat64: - return discardModelListGGUFBytes(r, 8) - case modelListGGUFTypeString: - return skipModelListGGUFString(r, byteOrder, version) - case modelListGGUFTypeArray: - var arrayType uint32 - if err := binary.Read(r, byteOrder, &arrayType); err != nil { - return err - } - var count uint64 - if err := binary.Read(r, byteOrder, &count); err != nil { - return err - } - return skipModelListGGUFArray(r, byteOrder, version, arrayType, count) - default: - return fmt.Errorf("unsupported gguf value type %d", valueType) - } -} - -func skipModelListGGUFArray(r io.Reader, byteOrder binary.ByteOrder, version uint32, arrayType uint32, count uint64) error { - if _, err := checkedModelListGGUFLength(count, "array size", fsgguf.MaxArraySize); err != nil { - return err - } - - var size uint64 - switch arrayType { - case modelListGGUFTypeUint8, modelListGGUFTypeInt8, modelListGGUFTypeBool: - size = 1 - case modelListGGUFTypeUint16, modelListGGUFTypeInt16: - size = 2 - case modelListGGUFTypeUint32, modelListGGUFTypeInt32, modelListGGUFTypeFloat32: - size = 4 - case modelListGGUFTypeUint64, modelListGGUFTypeInt64, modelListGGUFTypeFloat64: - size = 8 - case modelListGGUFTypeString: - for range count { - if err := skipModelListGGUFString(r, byteOrder, version); err != nil { - return err - } - } - return nil - default: - return fmt.Errorf("unsupported gguf array type %d", arrayType) - } - - if count > uint64(maxModelListGGUFInt64())/size { - return fmt.Errorf("gguf array byte length %d*%d exceeds maximum %d", count, size, maxModelListGGUFInt64()) - } - return discardModelListGGUFBytes(r, int64(count*size)) -} - -func readModelListGGUFString(r io.Reader, byteOrder binary.ByteOrder, version uint32) (string, error) { - var length uint64 - if err := binary.Read(r, byteOrder, &length); err != nil { - return "", err - } - - n, err := checkedModelListGGUFLength(length, "string", fsgguf.MaxStringLength) - if err != nil { - return "", err - } - - bts := make([]byte, n) - if _, err := io.ReadFull(r, bts); err != nil { - return "", err - } - if version == 1 && len(bts) > 0 && bts[len(bts)-1] == 0 { - bts = bts[:len(bts)-1] - } - return string(bts), nil -} - -func skipModelListGGUFString(r io.Reader, byteOrder binary.ByteOrder, version uint32) error { - var length uint64 - if err := binary.Read(r, byteOrder, &length); err != nil { - return err - } - - n, err := checkedModelListGGUFLength(length, "string", fsgguf.MaxStringLength) - if err != nil { - return err - } - return discardModelListGGUFBytes(r, int64(n)) -} - -func discardModelListGGUFBytes(r io.Reader, n int64) error { - if n <= 0 { - return nil - } - _, err := io.CopyN(io.Discard, r, n) - return err -} - -func checkedModelListGGUFLength(n uint64, kind string, max uint64) (int, error) { - if n > uint64(maxModelListGGUFInt()) { - return 0, fmt.Errorf("gguf %s %d exceeds maximum %d", kind, n, maxModelListGGUFInt()) - } - if n > max { - return 0, fmt.Errorf("gguf %s %d exceeds maximum %d", kind, n, max) - } - return int(n), nil -} - -func maxModelListGGUFInt() int { - return int(^uint(0) >> 1) -} - -func maxModelListGGUFInt64() int64 { - return 1<<63 - 1 -} - -func appendModelListCapabilities(capabilities []model.Capability, values ...model.Capability) []model.Capability { - for _, capability := range values { - capabilities = appendModelListCapability(capabilities, capability) - } - return capabilities -} - -func appendModelListCapability(capabilities []model.Capability, capability model.Capability) []model.Capability { - if capability == "" || slices.Contains(capabilities, capability) { - return capabilities - } - return append(capabilities, capability) -} - -func isUnknownQuantization(quantization string) bool { - return quantization == "" || quantization == "unknown" -} - -func cloneModelListSummary(summary modelListSummary) modelListSummary { - summary.Details.Families = append([]string(nil), summary.Details.Families...) - summary.Capabilities = append([]model.Capability(nil), summary.Capabilities...) - return summary -} - -func (s modelListSummary) ListModelResponse() api.ListModelResponse { - resp := api.ListModelResponse{ - Model: s.Model, - Name: s.Name, - RemoteModel: s.RemoteModel, - RemoteHost: s.RemoteHost, - Size: s.Size, - Digest: s.Digest, - ModifiedAt: s.ModifiedAt, - Details: api.ModelDetails{ - ParentModel: s.Details.ParentModel, - Format: s.Details.Format, - Family: s.Details.Family, - Families: append([]string(nil), s.Details.Families...), - ParameterSize: s.Details.ParameterSize, - QuantizationLevel: s.Details.QuantizationLevel, - ContextLength: s.Details.ContextLength, - EmbeddingLength: s.Details.EmbeddingLength, - }, - } - - resp.Capabilities = append([]model.Capability(nil), s.Capabilities...) - - return resp -} - -func sortListModelResponses(models []api.ListModelResponse) { - slices.SortStableFunc(models, func(i, j api.ListModelResponse) int { - // Preserve the existing /api/tags order: most recently modified first. - return cmp.Compare(j.ModifiedAt.Unix(), i.ModifiedAt.Unix()) - }) -} - -func (s *Server) refreshModelListCache(name model.Name) { - if s == nil || s.modelCaches == nil || s.modelCaches.modelList == nil { - return - } - - if err := s.modelCaches.modelList.RefreshModel(name); err != nil { - slog.Warn("failed to refresh model list cache", "model", name.String(), "error", err) - } -} - -func (s *Server) deleteModelListCache(name model.Name) { - if s == nil || s.modelCaches == nil || s.modelCaches.modelList == nil { - return - } - - s.modelCaches.modelList.DeleteModel(name) -} diff --git a/server/model_list_cache_test.go b/server/model_list_cache_test.go deleted file mode 100644 index f9fb960cd4e..00000000000 --- a/server/model_list_cache_test.go +++ /dev/null @@ -1,371 +0,0 @@ -package server - -import ( - "bytes" - "context" - "encoding/binary" - "errors" - "net/http" - "os" - "slices" - "strings" - "testing" - - "github.com/gin-gonic/gin" - - "github.com/ollama/ollama/api" - fsgguf "github.com/ollama/ollama/fs/gguf" - "github.com/ollama/ollama/manifest" - "github.com/ollama/ollama/types/model" -) - -func TestModelListCacheHydratesSummary(t *testing.T) { - gin.SetMode(gin.TestMode) - setTestHome(t, t.TempDir()) - createListCacheModel(t, "list-cache", map[string]any{ - "test.context_length": uint32(4096), - "test.embedding_length": uint32(384), - }, "{{ .prompt }}{{ if .tools }}{{ .tools }}{{ end }}{{ if .suffix }}{{ .suffix }}{{ end }}") - - cache := newModelListCache() - if err := cache.hydrate(context.Background()); err != nil { - t.Fatalf("hydrate failed: %v", err) - } - - summary, ok := cache.Get(model.ParseName("list-cache")) - if !ok { - t.Fatal("list summary missing") - } - - if summary.Model != "list-cache:latest" || summary.Name != "list-cache:latest" { - t.Fatalf("summary model/name = %q/%q, want list-cache:latest", summary.Model, summary.Name) - } - if summary.Digest == "" { - t.Fatal("summary digest is empty") - } - if summary.Size == 0 { - t.Fatal("summary size is zero") - } - if summary.Details.Family != "test" || summary.Details.Format != "gguf" { - t.Fatalf("summary details = %+v, want gguf/test", summary.Details) - } - if summary.Details.ContextLength != 4096 { - t.Fatalf("context length = %d, want 4096", summary.Details.ContextLength) - } - if summary.Details.EmbeddingLength != 384 { - t.Fatalf("embedding length = %d, want 384", summary.Details.EmbeddingLength) - } - - for _, capability := range []model.Capability{model.CapabilityCompletion, model.CapabilityTools, model.CapabilityInsert} { - if !slices.Contains(summary.Capabilities, capability) { - t.Fatalf("capabilities = %v, want %s", summary.Capabilities, capability) - } - } - - listModel := summary.ListModelResponse() - if !slices.Contains(listModel.Capabilities, model.CapabilityTools) || - listModel.Details.ContextLength != 4096 || - listModel.Details.EmbeddingLength != 384 { - t.Fatalf("list response = %+v, want capabilities/context/embedding", listModel) - } -} - -func TestModelListCacheSuppressesNemotronSafetensorsMedia(t *testing.T) { - caps := []model.Capability{ - model.CapabilityCompletion, - model.CapabilityTools, - model.CapabilityThinking, - model.CapabilityVision, - model.CapabilityAudio, - } - got := filterUnsupportedModelListCapabilities(caps, model.ConfigV2{ - ModelFormat: "safetensors", - Renderer: "nemotron-3-nano", - Parser: "nemotron-3-nano", - }) - - for _, capability := range []model.Capability{ - model.CapabilityCompletion, - model.CapabilityTools, - model.CapabilityThinking, - } { - if !slices.Contains(got, capability) { - t.Fatalf("capabilities = %v, want %s", got, capability) - } - } - for _, capability := range []model.Capability{model.CapabilityVision, model.CapabilityAudio} { - if slices.Contains(got, capability) { - t.Fatalf("capabilities = %v, did not expect %s", got, capability) - } - } -} - -func TestModelListCacheRefreshUpdatesEntry(t *testing.T) { - gin.SetMode(gin.TestMode) - setTestHome(t, t.TempDir()) - createListCacheModel(t, "list-refresh", map[string]any{"test.context_length": uint32(1024)}, "") - - cache := newModelListCache() - if err := cache.hydrate(context.Background()); err != nil { - t.Fatalf("hydrate failed: %v", err) - } - - name := model.ParseName("list-refresh") - first, ok := cache.Get(name) - if !ok { - t.Fatal("list summary missing") - } - - changeShowCacheManifest(t, "list-refresh") - if err := cache.RefreshModel(name); err != nil { - t.Fatalf("refresh failed: %v", err) - } - - refreshed, ok := cache.Get(name) - if !ok { - t.Fatal("refreshed list summary missing") - } - if refreshed.Digest == first.Digest { - t.Fatalf("digest did not change after refresh: %s", refreshed.Digest) - } - if cache.Len() != 1 { - t.Fatalf("cache entries = %d, want 1", cache.Len()) - } -} - -func TestModelListCacheMutationHooks(t *testing.T) { - gin.SetMode(gin.TestMode) - setTestHome(t, t.TempDir()) - - cache := newModelListCache() - s := Server{modelCaches: &modelCaches{modelList: cache}} - - _, digest := createBinFile(t, map[string]any{"test.context_length": uint32(2048)}, nil) - w := createRequest(t, s.CreateHandler, api.CreateRequest{ - Model: "list-hooks", - Files: map[string]string{"model.gguf": digest}, - Stream: &stream, - }) - if w.Code != http.StatusOK { - t.Fatalf("create model status = %d, want 200: %s", w.Code, w.Body.String()) - } - - if _, ok := cache.Get(model.ParseName("list-hooks")); !ok { - t.Fatal("create did not refresh model list cache") - } - - w = createRequest(t, s.CopyHandler, api.CopyRequest{ - Source: "list-hooks", - Destination: "list-hooks-copy", - }) - if w.Code != http.StatusOK { - t.Fatalf("copy model status = %d, want 200: %s", w.Code, w.Body.String()) - } - if _, ok := cache.Get(model.ParseName("list-hooks-copy")); !ok { - t.Fatal("copy did not refresh model list cache") - } - - w = createRequest(t, s.DeleteHandler, api.DeleteRequest{Model: "list-hooks-copy"}) - if w.Code != http.StatusOK { - t.Fatalf("delete model status = %d, want 200: %s", w.Code, w.Body.String()) - } - if _, ok := cache.Get(model.ParseName("list-hooks-copy")); ok { - t.Fatal("delete did not remove model list cache entry") - } -} - -func TestModelListCacheSyncsManifestChanges(t *testing.T) { - gin.SetMode(gin.TestMode) - setTestHome(t, t.TempDir()) - createListCacheModel(t, "list-sync-a", map[string]any{"test.context_length": uint32(1024)}, "") - - cache := newModelListCache() - cache.Start(context.Background()) - if err := cache.Wait(context.Background()); err != nil { - t.Fatalf("wait failed: %v", err) - } - - createListCacheModel(t, "list-sync-b", map[string]any{"test.context_length": uint32(2048)}, "") - models, err := cache.List(context.Background()) - if err != nil { - t.Fatalf("list failed: %v", err) - } - - names := make([]string, 0, len(models)) - for _, m := range models { - names = append(names, m.Name) - } - for _, want := range []string{"list-sync-a:latest", "list-sync-b:latest"} { - if !slices.Contains(names, want) { - t.Fatalf("names = %v, want %s", names, want) - } - } - - var other Server - w := createRequest(t, other.DeleteHandler, api.DeleteRequest{Model: "list-sync-a"}) - if w.Code != http.StatusOK { - t.Fatalf("delete model status = %d, want 200: %s", w.Code, w.Body.String()) - } - - models, err = cache.List(context.Background()) - if err != nil { - t.Fatalf("list after delete failed: %v", err) - } - names = names[:0] - for _, m := range models { - names = append(names, m.Name) - } - if slices.Contains(names, "list-sync-a:latest") || !slices.Contains(names, "list-sync-b:latest") { - t.Fatalf("names after delete = %v, want only list-sync-b", names) - } -} - -func TestModelListCacheSyncDropsStaleEntryOnRefreshFailure(t *testing.T) { - gin.SetMode(gin.TestMode) - setTestHome(t, t.TempDir()) - createListCacheModel(t, "list-stale", map[string]any{"test.context_length": uint32(1024)}, "") - - cache := newModelListCache() - cache.Start(context.Background()) - if err := cache.Wait(context.Background()); err != nil { - t.Fatalf("wait failed: %v", err) - } - - name := model.ParseName("list-stale") - if _, ok := cache.Get(name); !ok { - t.Fatal("list summary missing") - } - - changeShowCacheManifest(t, "list-stale") - cache.build = func(model.Name, *manifest.Manifest) (modelListSummary, error) { - return modelListSummary{}, errors.New("refresh failed") - } - - models, err := cache.List(context.Background()) - if err != nil { - t.Fatalf("list failed: %v", err) - } - if len(models) != 0 { - t.Fatalf("models = %+v, want stale entry removed", models) - } - if _, ok := cache.Get(name); ok { - t.Fatal("stale entry remained in cache after refresh failure") - } -} - -func TestReadModelListGGUFRejectsMalformedMetadata(t *testing.T) { - cases := []struct { - name string - data []byte - want string - }{ - { - name: "oversized key string", - data: modelListGGUFTestFile(func(b *bytes.Buffer) { - writeModelListGGUFHeader(t, b, 1) - writeModelListGGUFUint64(t, b, fsgguf.MaxStringLength+1) - }), - want: "string", - }, - { - name: "oversized skipped string", - data: modelListGGUFTestFile(func(b *bytes.Buffer) { - writeModelListGGUFHeader(t, b, 1) - writeModelListGGUFString(t, b, "unused") - writeModelListGGUFUint32(t, b, modelListGGUFTypeString) - writeModelListGGUFUint64(t, b, fsgguf.MaxStringLength+1) - }), - want: "string", - }, - { - name: "oversized skipped array", - data: modelListGGUFTestFile(func(b *bytes.Buffer) { - writeModelListGGUFHeader(t, b, 1) - writeModelListGGUFString(t, b, "unused") - writeModelListGGUFUint32(t, b, modelListGGUFTypeArray) - writeModelListGGUFUint32(t, b, modelListGGUFTypeUint8) - writeModelListGGUFUint64(t, b, fsgguf.MaxArraySize+1) - }), - want: "array size", - }, - } - - for _, tt := range cases { - t.Run(tt.name, func(t *testing.T) { - defer func() { - if r := recover(); r != nil { - t.Fatalf("readModelListGGUF panicked: %v", r) - } - }() - - path := t.TempDir() + "/model.gguf" - if err := os.WriteFile(path, tt.data, 0o600); err != nil { - t.Fatal(err) - } - - _, err := readModelListGGUF(path) - if err == nil { - t.Fatal("expected error") - } - if !strings.Contains(err.Error(), tt.want) { - t.Fatalf("error = %v, want substring %q", err, tt.want) - } - }) - } -} - -func createListCacheModel(t *testing.T, name string, kv map[string]any, tmpl string) { - t.Helper() - _, digest := createBinFile(t, kv, nil) - - req := api.CreateRequest{ - Model: name, - Files: map[string]string{"model.gguf": digest}, - Stream: &stream, - } - if tmpl != "" { - req.Template = tmpl - } - - var s Server - w := createRequest(t, s.CreateHandler, req) - if w.Code != http.StatusOK { - t.Fatalf("create model status = %d, want 200: %s", w.Code, w.Body.String()) - } -} - -func modelListGGUFTestFile(fn func(*bytes.Buffer)) []byte { - var b bytes.Buffer - fn(&b) - return b.Bytes() -} - -func writeModelListGGUFHeader(t *testing.T, b *bytes.Buffer, numKV uint64) { - t.Helper() - writeModelListGGUFUint32(t, b, modelListGGUFMagicLE) - writeModelListGGUFUint32(t, b, 3) - writeModelListGGUFUint64(t, b, 0) - writeModelListGGUFUint64(t, b, numKV) -} - -func writeModelListGGUFString(t *testing.T, b *bytes.Buffer, s string) { - t.Helper() - writeModelListGGUFUint64(t, b, uint64(len(s))) - if _, err := b.WriteString(s); err != nil { - t.Fatal(err) - } -} - -func writeModelListGGUFUint32(t *testing.T, b *bytes.Buffer, v uint32) { - t.Helper() - if err := binary.Write(b, binary.LittleEndian, v); err != nil { - t.Fatal(err) - } -} - -func writeModelListGGUFUint64(t *testing.T, b *bytes.Buffer, v uint64) { - t.Helper() - if err := binary.Write(b, binary.LittleEndian, v); err != nil { - t.Fatal(err) - } -} diff --git a/server/model_list_test.go b/server/model_list_test.go new file mode 100644 index 00000000000..44fd83036d8 --- /dev/null +++ b/server/model_list_test.go @@ -0,0 +1,168 @@ +package server + +import ( + "context" + "os" + "slices" + "testing" + + "github.com/gin-gonic/gin" + + "github.com/ollama/ollama/api" + "github.com/ollama/ollama/manifest" + "github.com/ollama/ollama/types/model" +) + +func listedModel(t *testing.T, name string) api.ListModelResponse { + t.Helper() + models, err := listModels(context.Background()) + if err != nil { + t.Fatalf("listModels failed: %v", err) + } + for _, m := range models { + if m.Name == name { + return m + } + } + t.Fatalf("%s not listed; got %v", name, models) + return api.ListModelResponse{} +} + +func TestListModelsDescribesModel(t *testing.T) { + gin.SetMode(gin.TestMode) + setTestHome(t, t.TempDir()) + createListedModelFromKV(t, "list-describe", map[string]any{ + "test.context_length": uint32(4096), + "test.embedding_length": uint32(384), + }, "{{ .prompt }}{{ if .tools }}{{ .tools }}{{ end }}{{ if .suffix }}{{ .suffix }}{{ end }}") + + got := listedModel(t, "list-describe:latest") + + if got.Model != "list-describe:latest" { + t.Errorf("model = %q", got.Model) + } + if got.Digest == "" || got.Size == 0 { + t.Errorf("digest = %q size = %d, want both set", got.Digest, got.Size) + } + if got.Details.Family != "test" || got.Details.Format != "gguf" { + t.Errorf("details = %+v, want gguf/test", got.Details) + } + if got.Details.ContextLength != 4096 { + t.Errorf("context length = %d, want 4096", got.Details.ContextLength) + } + if got.Details.EmbeddingLength != 384 { + t.Errorf("embedding length = %d, want 384", got.Details.EmbeddingLength) + } + // The test GGUF has no general.file_type; the list must not guess F32 + // (FileType 0) from a missing key. + if !isUnknownQuantization(got.Details.QuantizationLevel) { + t.Errorf("quantization = %q, want unknown (no general.file_type in GGUF)", got.Details.QuantizationLevel) + } + for _, capability := range []model.Capability{model.CapabilityCompletion, model.CapabilityTools, model.CapabilityInsert} { + if !slices.Contains(got.Capabilities, capability) { + t.Errorf("capabilities = %v, want %s", got.Capabilities, capability) + } + } +} + +// Listing reads the manifests every time, so a model created or deleted after +// the last call shows up without anything needing to be told about it. +func TestListModelsFollowsManifestChanges(t *testing.T) { + gin.SetMode(gin.TestMode) + setTestHome(t, t.TempDir()) + createListedModelFromKV(t, "list-follow-a", map[string]any{"test.context_length": uint32(1024)}, "") + + listedModel(t, "list-follow-a:latest") + + createListedModelFromKV(t, "list-follow-b", map[string]any{"test.context_length": uint32(2048)}, "") + listedModel(t, "list-follow-a:latest") + listedModel(t, "list-follow-b:latest") + + deleteModelNamed(t, "list-follow-a") + + models, err := listModels(context.Background()) + if err != nil { + t.Fatal(err) + } + names := make([]string, 0, len(models)) + for _, m := range models { + names = append(names, m.Name) + } + if slices.Contains(names, "list-follow-a:latest") || !slices.Contains(names, "list-follow-b:latest") { + t.Fatalf("names after delete = %v, want only list-follow-b", names) + } +} + +func TestCapabilitiesSuppressNemotronSafetensorsMedia(t *testing.T) { + caps := []model.Capability{ + model.CapabilityCompletion, + model.CapabilityTools, + model.CapabilityThinking, + model.CapabilityVision, + model.CapabilityAudio, + } + got := (&Model{Config: model.ConfigV2{ + ModelFormat: "safetensors", + Renderer: "nemotron-3-nano", + Parser: "nemotron-3-nano", + }}).filterUnsupportedCapabilities(caps, "") + + for _, capability := range []model.Capability{ + model.CapabilityCompletion, + model.CapabilityTools, + model.CapabilityThinking, + } { + if !slices.Contains(got, capability) { + t.Errorf("capabilities = %v, want %s", got, capability) + } + } + for _, capability := range []model.Capability{model.CapabilityVision, model.CapabilityAudio} { + if slices.Contains(got, capability) { + t.Errorf("capabilities = %v, did not expect %s", got, capability) + } + } +} + +// A model whose layers no longer parse must still be listed: it is the one the +// user most needs to see, in order to remove it. +func TestListModelsKeepsUnloadableModel(t *testing.T) { + gin.SetMode(gin.TestMode) + setTestHome(t, t.TempDir()) + createListedModelFromKV(t, "broken", map[string]any{"test.context_length": uint32(1024)}, "{{ .Prompt }}") + + mf, err := manifest.ParseNamedManifest(model.ParseName("broken")) + if err != nil { + t.Fatal(err) + } + var corrupted bool + for _, layer := range mf.Layers { + if layer.MediaType != "application/vnd.ollama.image.template" { + continue + } + path, err := manifest.BlobsPath(layer.Digest) + if err != nil { + t.Fatal(err) + } + if err := os.WriteFile(path, []byte("{{ if }"), 0o644); err != nil { + t.Fatal(err) + } + corrupted = true + } + if !corrupted { + t.Fatal("no template layer to corrupt") + } + if _, err := GetModel("broken"); err == nil { + t.Fatal("model still loads, so the listing is not being asked the question") + } + + got := listedModel(t, "broken:latest") + if got.Details.Family != "test" { + t.Errorf("details = %+v, want the manifest config's family", got.Details) + } +} + +func createListedModelFromKV(t *testing.T, name string, kv map[string]any, tmpl string) { + t.Helper() + _, digest := createBinFile(t, kv, nil) + createModelFromBlob(t, name, digest, tmpl) +} diff --git a/server/routes.go b/server/routes.go index 9b73139a1a6..177a7441f70 100644 --- a/server/routes.go +++ b/server/routes.go @@ -286,7 +286,7 @@ func (s *Server) GenerateHandler(c *gin.Context) { return } - m, err := s.getModel(name.String()) + m, err := GetModel(name.String()) if err != nil { switch { case errors.Is(err, fs.ErrNotExist): @@ -850,7 +850,7 @@ func (s *Server) EmbedHandler(c *gin.Context) { return } - m, err := s.getModel(name.String()) + m, err := GetModel(name.String()) if err != nil { handleScheduleError(c, req.Model, err) return @@ -1068,7 +1068,7 @@ func (s *Server) EmbeddingsHandler(c *gin.Context) { name := modelRef.Name - m, err := s.getModel(name.String()) + m, err := GetModel(name.String()) if err != nil { handleScheduleError(c, req.Model, err) return @@ -1153,8 +1153,6 @@ func (s *Server) PullHandler(c *gin.Context) { ch <- gin.H{"error": err.Error()} return } - - s.refreshModelListCache(name) }() if req.Stream != nil && !*req.Stream { @@ -1294,9 +1292,9 @@ func (s *Server) DeleteHandler(c *gin.Context) { return } - s.deleteModelListCache(n) - - if err := m.RemoveLayers(); err != nil { + removed, err := m.RemoveLayers() + removeGGUFMetadata(removed...) + if err != nil { c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()}) return } @@ -1635,12 +1633,7 @@ func selectedModelTemplate(m *Model, kv ggml.KV) string { } func (s *Server) ListHandler(c *gin.Context) { - if s.modelCaches == nil || s.modelCaches.modelList == nil { - c.JSON(http.StatusInternalServerError, gin.H{"error": "model list cache unavailable"}) - return - } - - models, err := s.modelCaches.modelList.List(c.Request.Context()) + models, err := listModels(c.Request.Context()) if err != nil { c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()}) return @@ -1685,8 +1678,6 @@ func (s *Server) CopyHandler(c *gin.Context) { c.JSON(http.StatusNotFound, gin.H{"error": fmt.Sprintf("model %q not found", r.Source)}) } else if err != nil { c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()}) - } else { - s.refreshModelListCache(dst) } } @@ -2492,7 +2483,7 @@ func (s *Server) ChatHandler(c *gin.Context) { return } - m, err := s.getModel(name.String()) + m, err := GetModel(name.String()) if err != nil { switch { case os.IsNotExist(err): diff --git a/server/routes_list_test.go b/server/routes_list_test.go index eb409c8b073..9417c629288 100644 --- a/server/routes_list_test.go +++ b/server/routes_list_test.go @@ -1,7 +1,6 @@ package server import ( - "context" "encoding/json" "io" "net/http" @@ -36,11 +35,7 @@ func TestList(t *testing.T) { "myhost/mynamespace/lips:code", } - s := Server{modelCaches: &modelCaches{modelList: newModelListCache()}} - s.modelCaches.modelList.Start(context.Background()) - if err := s.modelCaches.modelList.Wait(context.Background()); err != nil { - t.Fatal(err) - } + var s Server for _, n := range expectNames { _, digest := createBinFile(t, nil, nil) @@ -88,12 +83,7 @@ func TestOpenAIListMatchesTagsModels(t *testing.T) { gin.SetMode(gin.TestMode) setTestHome(t, t.TempDir()) - cache := newModelListCache() - s := Server{modelCaches: &modelCaches{modelList: cache}} - cache.Start(context.Background()) - if err := cache.Wait(context.Background()); err != nil { - t.Fatal(err) - } + var s Server createModel := func(name string) { t.Helper() @@ -123,9 +113,6 @@ func TestOpenAIListMatchesTagsModels(t *testing.T) { if err := os.Chtimes(path, modified, modified); err != nil { t.Fatalf("set manifest time for %s: %v", name, err) } - if err := cache.RefreshModel(parsed); err != nil { - t.Fatalf("refresh %s: %v", name, err) - } } older := time.Unix(1000, 0).UTC() diff --git a/server/routes_test.go b/server/routes_test.go index 8164454e9ef..44639f6f067 100644 --- a/server/routes_test.go +++ b/server/routes_test.go @@ -95,11 +95,7 @@ func TestRoutes(t *testing.T) { Expected func(t *testing.T, resp *http.Response) } - s := &Server{modelCaches: &modelCaches{modelList: newModelListCache()}} - s.modelCaches.modelList.Start(context.Background()) - if err := s.modelCaches.modelList.Wait(context.Background()); err != nil { - t.Fatal(err) - } + s := &Server{} createTestModel := func(t *testing.T, name string) { t.Helper() @@ -135,7 +131,6 @@ func TestRoutes(t *testing.T) { if err := createModel(r, modelName, baseLayers, config, fn); err != nil { t.Fatal(err) } - s.refreshModelListCache(modelName) } testCases := []testCase{ From 8d66f083557d9287b78090de6741b5d7f6bd6e79 Mon Sep 17 00:00:00 2001 From: Jesse Gross Date: Wed, 9 Sep 2026 12:45:17 -0700 Subject: [PATCH 04/24] mlxrunner: capture whole-state at close for nodes split out of an edge KV snapshots must cover a node's edge exactly. Recurrent and sliding-window state is only useful at a node's end, and a node may have none: a request resuming there lands on the previous checkpoint and begin schedules a capture at the match. The header claimed every node carries its snapshots from creation, which a node split out of an existing edge at close cannot. That hid a gap: when a response is a prefix of a stored one, close lands on the split-off head with the caches resting at its end, and pageOut skipped the capture because the node already had a KV snapshot. Restate the header as the rules that hold, and make pageOut capture whatever layers a node is missing. The scheduling comment also said eviction preserves user nodes; it only resists compaction. --- x/mlxrunner/prefix_cache.go | 31 ++++++++++++++++++++----------- 1 file changed, 20 insertions(+), 11 deletions(-) diff --git a/x/mlxrunner/prefix_cache.go b/x/mlxrunner/prefix_cache.go index 1adbb2376cc..b149ad896f7 100644 --- a/x/mlxrunner/prefix_cache.go +++ b/x/mlxrunner/prefix_cache.go @@ -3,14 +3,22 @@ // optional per-layer snapshots that can be paged in/out of the live MLX cache // arrays. // -// Key properties: +// Invariants: // - Only one path through the trie is "active" (backed by live MLX arrays) // at a time. Switching paths pages in the new path from its snapshots. -// - Every node carries its snapshots from creation: prefill captures for -// prompt segments, a page-out at close for generated ones. Sliceable -// (KV) layers always span exactly the node's edge; whole-state layers -// (recurrent, rotating) keep entries only at node ends. +// - Sliceable (KV) layers: every node holds a snapshot covering exactly its +// edge, so the layer's history is complete along any path from the root. +// - Whole-state (recurrent, rotating) layers: what a node holds is the +// state at its end offset. A node may hold none. +// - Whole-state is captured only while the live caches sit at that offset +// (prefill captures, the page-out at close) and is never rebuilt later. A +// node split out of an existing edge afterward therefore holds none. +// - A request resumes at the deepest node at or below its match that holds +// whole-state. begin schedules a capture at the match, so any node a +// request resumes at holds whole-state afterward. // - All cache layers must stay at the same token offset. +// - Draft caches are settled whenever the trie captures, pages out, or +// rewinds: no entry is still waiting on the next token. // - Sibling edges must not share a common token prefix (compressed trie // invariant). // - begin() always re-evaluates at least one token so the pipeline can seed @@ -281,7 +289,7 @@ pageIn: // state that prefix alone determines (offset - draftLookahead), which is where // a prompt sharing exactly that prefix restores. The offsets are merged with // any snapshots begin already scheduled (e.g. a branch point), with coinciding -// offsets upgraded to user so eviction preserves them. +// offsets upgraded to user so compaction keeps them. // // Offsets at or before the current cache position, or past the end of the // prompt, are dropped: callers only request offsets ahead of the prefill base, @@ -471,20 +479,21 @@ func (c *prefixCache) compactPath() { c.activePath = c.activePath[:n-1] } -// pageOut captures a fresh node's state from the live caches, which rest -// exactly at its end. Nodes reached by traversal already carry snapshots. +// pageOut captures the snapshots a node is missing from the live caches, which +// rest exactly at its end. func (c *prefixCache) pageOut(node *trieNode) { - if node.hasSnapshots() { + if hasAllSnapshots(node, c.caches) { return } snaps := make([]cache.Snapshot, len(c.caches)) + copy(snaps, node.snapshots) for i, kv := range c.caches { - if kv == nil { + if kv == nil || snaps[i] != nil { continue } snaps[i] = kv.Snapshot(node.startOffset()) } - node.setSnapshots(snaps, &c.pagedOutBytes) + node.swapSnapshots(snaps, &c.pagedOutBytes) logutil.Trace(fmt.Sprintf("page out: [%d, %d)", node.startOffset(), node.endOffset)) c.enforceEvictionPolicy() } From 45a02807e2dd9946b618876ee47d31b0f2e8e963 Mon Sep 17 00:00:00 2001 From: Jesse Gross Date: Wed, 9 Sep 2026 13:16:16 -0700 Subject: [PATCH 05/24] mlxrunner: keep cache boundaries out of non-causal media items The tokens of a non-causal media item attend to each other in both directions, so the item has to be evaluated in one forward. Prefill honors that when it picks chunk boundaries, but the prefix cache did not: a snapshot could be taken partway through an item, and a request that resumed there would evaluate the rest of the item alone and compute different attention for it. Snapshots scheduled inside a non-causal item now land at its end, and a match that ends inside one resumes at its start. --- x/mlxrunner/prefix_cache.go | 26 +++++++++++++++++++- x/mlxrunner/prefix_cache_test.go | 41 ++++++++++++++++++++++++++++++++ 2 files changed, 66 insertions(+), 1 deletion(-) diff --git a/x/mlxrunner/prefix_cache.go b/x/mlxrunner/prefix_cache.go index b149ad896f7..79050de6f65 100644 --- a/x/mlxrunner/prefix_cache.go +++ b/x/mlxrunner/prefix_cache.go @@ -19,6 +19,8 @@ // - All cache layers must stay at the same token offset. // - Draft caches are settled whenever the trie captures, pages out, or // rewinds: no entry is still waiting on the next token. +// - A non-causal media item's tokens are evaluated in one batch: no node +// boundary or resume point lies strictly inside them. // - Sibling edges must not share a common token prefix (compressed trie // invariant). // - begin() always re-evaluates at least one token so the pipeline can seed @@ -64,6 +66,7 @@ type cacheSession struct { cache *prefixCache inputs []int32 effInputs []uint32 // inputs' key alphabet, media folds applied + items []mediaItem outputs []int32 caches []cache.Cache @@ -104,6 +107,10 @@ func (c *prefixCache) begin(inputs []int32, items []mediaItem) *cacheSession { if matched == len(inputs) && matched > 0 { matchPath, matched = findBestMatch(c.root, keys[:matched-1]) } + // A match ending inside a non-causal media item resumes before it. + if item := insideAtomicItem(items, matched); item != nil { + matchPath, matched = findBestMatch(c.root, keys[:item.pos]) + } // Switch to the matched path, paging in/out as needed. c.switchToPath(matchPath, matched) @@ -116,6 +123,7 @@ func (c *prefixCache) begin(inputs []int32, items []mediaItem) *cacheSession { cache: c, inputs: inputs, effInputs: effInputs, + items: items, caches: c.caches, remaining: remaining, } @@ -135,6 +143,18 @@ func (c *prefixCache) begin(inputs []int32, items []mediaItem) *cacheSession { return session } +// insideAtomicItem returns the non-causal media item that offset lies strictly +// inside, or nil. +func insideAtomicItem(items []mediaItem, offset int) *mediaItem { + for i := range items { + item := &items[i] + if item.atomic() && item.pos < offset && offset < item.pos+item.length { + return item + } + } + return nil +} + // effectiveKeyTokens returns the per-position key alphabet: the token ID // outside media expansions, the item's fold value across each expansion's // whole range. @@ -287,7 +307,8 @@ pageIn: // prefill records interior states without the caller breaking the batch. A // passed offset names a token prefix; the capture lands at the deepest // state that prefix alone determines (offset - draftLookahead), which is where -// a prompt sharing exactly that prefix restores. The offsets are merged with +// a prompt sharing exactly that prefix restores. An offset inside a non-causal +// media item's tokens moves past them. The offsets are merged with // any snapshots begin already scheduled (e.g. a branch point), with coinciding // offsets upgraded to user so compaction keeps them. // @@ -299,6 +320,9 @@ func (s *cacheSession) schedulePrefillSnapshots(offsets []int) { base := c.minCacheOffset() for _, offset := range offsets { offset -= c.draftLookahead + if item := insideAtomicItem(s.items, offset); item != nil { + offset = item.pos + item.length + } if offset <= base || offset > len(s.inputs) { continue } diff --git a/x/mlxrunner/prefix_cache_test.go b/x/mlxrunner/prefix_cache_test.go index cb7165b665f..e5617534bf9 100644 --- a/x/mlxrunner/prefix_cache_test.go +++ b/x/mlxrunner/prefix_cache_test.go @@ -7,6 +7,7 @@ import ( "github.com/ollama/ollama/x/mlxrunner/cache" "github.com/ollama/ollama/x/mlxrunner/mlx" + "github.com/ollama/ollama/x/mlxrunner/model/base" ) // snapshotTracker records every fakeSnapshot created and every Close() call @@ -1011,6 +1012,46 @@ func TestSnapshotBeyondPrefillSkipped(t *testing.T) { }) } +// TestAtomicMediaBoundaries verifies that a non-causal media item is never +// split by the cache: a capture scheduled inside its tokens lands at their +// end, and a prompt whose match ends inside them resumes before them. +func TestAtomicMediaBoundaries(t *testing.T) { + forEachEnv(t, func(t *testing.T, env *testEnv) { + pc := env.pc + inputs := []int32{1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12} + const itemPos, itemLen = 3, 6 + items := []mediaItem{{pos: itemPos, length: itemLen, fold: 1 << 31, item: &base.PreparedItem{}}} + + session := pc.begin(inputs, items) + session.schedulePrefillSnapshots([]int{itemPos + itemLen/2}) + feedAll(pc.caches, inputs[pc.minCacheOffset():len(inputs)-1]) + session.attachPrefillSnapshots() + session.close() + + walkNodes(pc.root, func(n *trieNode) bool { + if itemPos < n.endOffset && n.endOffset < itemPos+itemLen { + t.Errorf("trie node ends at %d, inside the item's tokens [%d,%d)", n.endOffset, itemPos, itemPos+itemLen) + } + return true + }) + if !nodeExistsAtOffset(pc.root, itemPos+itemLen) { + t.Errorf("capture inside the item did not move to its end %d", itemPos+itemLen) + } + + // A prompt ending inside the item matches the stored path through its + // last token, which would put the resume point inside the item. + short := inputs[:itemPos+itemLen-2] + shortItems := []mediaItem{{pos: itemPos, length: len(short) - itemPos, fold: 1 << 31, item: &base.PreparedItem{}}} + session = pc.begin(short, shortItems) + if resumed := len(short) - len(session.remaining); resumed > itemPos { + t.Errorf("resumed at %d, inside the item's tokens starting at %d", resumed, itemPos) + } + session.close() + + checkTrieInvariants(t, pc.root) + }) +} + // TestPrefillSnapshotsKeptOnCancel mirrors a prefill canceled after the caches // captured interior snapshots but before the success-path attach ran. Closing // the session attaches the crossed captures so a retry can resume from them, From b859a94509d19ad9018fdbadc83b0000b3e6c0ec Mon Sep 17 00:00:00 2001 From: Jesse Gross Date: Thu, 10 Sep 2026 11:27:20 -0700 Subject: [PATCH 06/24] mlxrunner: keep the reused head of a cached edge safe from eviction When a request resumes partway through a cached edge, the node holding that edge was dropped from the active path, because the path has to end at the live offset for close and the prefill captures to extend the trie from its last node. Off the path, the node was an ordinary leaf with a stale last-used time, so eviction removed it first. The captures taken during that request start at the resume offset, but attach rebuilds the missing node from the path's last node, so the new node's edge begins earlier than its KV snapshot. A later request resuming there was refused by the KV cache and re-prefilled from scratch. If eviction first merged the node into its parent, the two snapshots were concatenated as if adjacent, and the restore reported a hit while the buffer held tokens from other positions. Split the node at the live offset instead. The head stays on the path and gets the last-used update. Only the unused tail can be evicted, and losing it costs nothing. The split only happens when every layer can rewind into the edge, so it never involves a recurrent layer, and the head gets the same KV-only snapshots a close-time split already produces. When the request follows the edge, compaction merges the halves back. --- x/mlxrunner/prefix_cache.go | 14 +++++++++++--- 1 file changed, 11 insertions(+), 3 deletions(-) diff --git a/x/mlxrunner/prefix_cache.go b/x/mlxrunner/prefix_cache.go index 79050de6f65..4ced29fb192 100644 --- a/x/mlxrunner/prefix_cache.go +++ b/x/mlxrunner/prefix_cache.go @@ -283,11 +283,19 @@ pageIn: } } } + + // If the live offset falls inside the last node, split it so the reused + // head stays on the active path and only the unused tail can be evicted. for i := len(c.activePath) - 1; i >= 0; i-- { - if c.activePath[i].endOffset <= minOff { - c.activePath = c.activePath[:i+1] - break + node := c.activePath[i] + if i > 0 && node.startOffset() >= minOff { + continue + } + if node.endOffset > minOff { + node = splitNode(node, minOff-node.startOffset(), c.caches, &c.pagedOutBytes) } + c.activePath = append(c.activePath[:i], node) + break } // Update last-used time on only the final used node. For recurrent From 6137793ac495305c4a7d82a9a33e3f563e084c03 Mon Sep 17 00:00:00 2001 From: Jesse Gross Date: Wed, 9 Sep 2026 13:52:48 -0700 Subject: [PATCH 07/24] mlxrunner: evict the active conversation's own checkpoints under the budget Eviction skipped every node on the active path, so a conversation's own turn checkpoints were never reclaimed no matter how far over budget the trie was. On models with sliding-window or recurrent layers each turn's checkpoint is a full copy of that state, 800 MiB per turn on gemma4:31b-mlx, and a long chat grows without bound. The scheduler then counts that memory as in use and evicts the model to load anything else. Only the frontier and branch points are protected now. Any other node, active or not, is evicted least recently used first. On the active path that merges a turn into the next one: the merged node keeps the newer whole-state, and the KV snapshots there are lazy views of the live buffer, so nothing is copied. Rewinding to an evicted turn resumes at the newest surviving checkpoint before it. qwen3.8:27b-mlx on an M5 Max, the same short question every turn with 24 tokens generated per reply, 8 GiB budget, 17.2 GiB of weights: turn | before: paged out nodes reported | after: paged out nodes reported 11 | 4.61 GiB 33 21.6 GiB | 4.61 GiB 33 21.6 GiB 21 | 7.91 GiB 56 24.9 GiB | 7.92 GiB 56 24.9 GiB 31 | 8.46 GiB 60 25.5 GiB | 7.94 GiB 56 25.0 GiB 41 | 9.90 GiB 70 26.9 GiB | 7.96 GiB 56 25.0 GiB 50 | 11.19 GiB 79 28.2 GiB | 7.98 GiB 56 25.0 GiB Fixes #17783 --- x/mlxrunner/prefix_cache.go | 14 ++++++++------ x/mlxrunner/prefix_cache_test.go | 26 ++++++++++++++++++++++---- 2 files changed, 30 insertions(+), 10 deletions(-) diff --git a/x/mlxrunner/prefix_cache.go b/x/mlxrunner/prefix_cache.go index 4ced29fb192..0f825dcd36d 100644 --- a/x/mlxrunner/prefix_cache.go +++ b/x/mlxrunner/prefix_cache.go @@ -606,15 +606,13 @@ func (c *prefixCache) enforceEvictionPolicy() { return } - activeSet := make(map[*trieNode]bool, len(c.activePath)) - for _, n := range c.activePath { - activeSet[n] = true - } - for c.pagedOutBytes > maxPagedOutBytes { + // Evicting the frontier's parent merges the frontier into it, so + // resolve the frontier again after every eviction. + frontier := c.activePath[len(c.activePath)-1] var best *trieNode walkNodes(c.root, func(n *trieNode) bool { - if n == c.root || activeSet[n] || len(n.children) > 1 { + if n == c.root || n == frontier || len(n.children) > 1 { return true } // Evict: oldest, then deepest, then largest. @@ -644,7 +642,11 @@ func (c *prefixCache) evictNode(node *trieNode) { // Interior node with one child: merge with child. before := c.pagedOutBytes tokens := len(node.tokens) + child := node.children[0] mergeWithChild(node, c.caches, &c.pagedOutBytes) + if i := slices.Index(c.activePath, child); i >= 0 { + c.activePath = slices.Delete(c.activePath, i, i+1) + } slog.Debug("evicting interior node", "offset", node.startOffset(), "tokens", tokens, "freed", mlx.PrettyBytes(int(before-c.pagedOutBytes))) } else { panic("evictNode called on multi-child branch point") diff --git a/x/mlxrunner/prefix_cache_test.go b/x/mlxrunner/prefix_cache_test.go index e5617534bf9..b870fd7ea50 100644 --- a/x/mlxrunner/prefix_cache_test.go +++ b/x/mlxrunner/prefix_cache_test.go @@ -881,7 +881,7 @@ func TestEvictionPreservesActiveConversations(t *testing.T) { t.Fatalf("pagedOutBytes = %d, want <= %d", pc.pagedOutBytes, maxPagedOutBytes) } - // Active path should be untouched. + // The branch point and the frontier survive. if len(pc.activePath) < 2 { t.Fatalf("activePath should have >= 2 nodes, got %d", len(pc.activePath)) } @@ -954,8 +954,26 @@ func TestUserSnapshotResistsAutoMerge(t *testing.T) { t.Fatalf("user node children = %d, want 2", len(userNode.children)) } - // Inflate snapshot sizes and evict. The non-active branch should be - // evicted, leaving the user node with one child. + // Inflate snapshot sizes so that evicting the non-active branch alone + // brings the trie under budget, leaving the user node with one child. + var kept, evicted int + walkNodes(pc.root, func(n *trieNode) bool { + for _, s := range n.snapshots { + if s == nil { + continue + } + if n.parent == userNode && !slices.Contains(pc.activePath, n) { + evicted++ + } else { + kept++ + } + } + return true + }) + if evicted == 0 { + t.Fatal("no snapshots on the non-active branch") + } + size := int(maxPagedOutBytes) / kept walkNodes(pc.root, func(n *trieNode) bool { if !n.hasSnapshots() { return true @@ -963,7 +981,7 @@ func TestUserSnapshotResistsAutoMerge(t *testing.T) { snaps := make([]cache.Snapshot, len(n.snapshots)) for i, s := range n.snapshots { if s != nil { - snaps[i] = &fakeSnapshot{byteSize: 5 * 1024 * 1024 * 1024} + snaps[i] = &fakeSnapshot{byteSize: size} } } n.setSnapshots(snaps, &pc.pagedOutBytes) From 1548f78c733bc159de6ca07d705f55c6a65085a5 Mon Sep 17 00:00:00 2001 From: Jesse Gross Date: Thu, 10 Sep 2026 11:49:42 -0700 Subject: [PATCH 08/24] mlxrunner: bound MLX loads by system free memory while other models are loaded On Apple silicon the scheduler's free-memory figure for the GPU is the Metal working set minus what Ollama's own runners report. It does not see memory held by other applications, so a second MLX model can pass the fit check on a machine that is already short of memory, and the load pushes the system into swap and compression. While other models are loaded, the MLX fit check now also bounds the available memory by the system's free memory on shared-memory GPUs, the same rule llama-server loads already apply. A miss evicts an idle model and retries instead of starting the load. First loads are unchanged: with nothing else loaded, the model loads against the working-set figure alone, as both engines do today. The check also does not cover memory that grows after load, such as KV caches and prefix-cache snapshots. --- x/mlxrunner/client.go | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) diff --git a/x/mlxrunner/client.go b/x/mlxrunner/client.go index 7f242fbd424..728ba654db2 100644 --- a/x/mlxrunner/client.go +++ b/x/mlxrunner/client.go @@ -301,11 +301,14 @@ func (c *Client) HasExited() bool { } // Load checks whether the model fits in GPU memory and starts the subprocess. -func (c *Client) Load(ctx context.Context, _ ml.SystemInfo, gpus []ml.DeviceInfo, requireFull bool) ([]ml.DeviceID, error) { +func (c *Client) Load(ctx context.Context, systemInfo ml.SystemInfo, gpus []ml.DeviceInfo, requireFull bool) ([]ml.DeviceID, error) { if len(gpus) > 0 { modelSize := c.memory.Load() // We currently only use the first GPU with MLX available := gpus[0].FreeMemory + if requireFull && gpus[0].Integrated && systemInfo.FreeMemory > 0 && systemInfo.FreeMemory < available { + available = systemInfo.FreeMemory + } overhead := gpus[0].MinimumMemory() + envconfig.GpuOverhead() if available > overhead { available -= overhead From f09d55d0e247e2c2c0cb059a091cf38569eab5dc Mon Sep 17 00:00:00 2001 From: Jesse Gross Date: Wed, 9 Sep 2026 11:36:53 -0700 Subject: [PATCH 09/24] mlxrunner: wait for a killed runner to exit before the scheduler loads the next model The scheduler starts the next load as soon as Close returns. The MLX client sent SIGINT, gave the process five seconds, then sent SIGKILL and returned without waiting, so a runner that could not take the signal was still exiting, with its memory still held, when the next load began. The runner has no signal handler, so SIGINT was already a kill. Load also started the process and recorded it without the client's mutex, so a Close racing with a load at server shutdown could find nothing to stop and leave the runner it missed running. Close now kills the process and waits for it to be reaped, as the llama-server client does. Load starts and records the process under the mutex and refuses to start once Close has run. --- x/mlxrunner/client.go | 22 ++++++++++++---------- 1 file changed, 12 insertions(+), 10 deletions(-) diff --git a/x/mlxrunner/client.go b/x/mlxrunner/client.go index 728ba654db2..5bdfb56e82c 100644 --- a/x/mlxrunner/client.go +++ b/x/mlxrunner/client.go @@ -42,6 +42,7 @@ type Client struct { status *llm.StatusWriter mu sync.Mutex cmd *exec.Cmd + closed bool } // NewClient prepares a new MLX runner client for LLM models. @@ -141,15 +142,11 @@ func (c *Client) Close() error { c.mu.Lock() defer c.mu.Unlock() + c.closed = true if c.cmd != nil && c.cmd.Process != nil { slog.Info("stopping mlx runner subprocess", "pid", c.cmd.Process.Pid) - c.cmd.Process.Signal(os.Interrupt) - - select { - case <-c.done: - case <-time.After(5 * time.Second): - c.cmd.Process.Kill() - } + c.cmd.Process.Kill() + <-c.done c.cmd = nil } return nil @@ -405,19 +402,24 @@ func (c *Client) Load(ctx context.Context, systemInfo ml.SystemInfo, gpus []ml.D } } - c.cmd = cmd - status := llm.NewStatusWriter(os.Stderr) - c.status = status // os/exec serializes Write calls when shared, which keeps the status writer // from seeing concurrent stdout/stderr fragments. cmd.Stdout = status cmd.Stderr = status + c.mu.Lock() + defer c.mu.Unlock() + if c.closed { + return nil, errors.New("mlx runner client is closed") + } + + c.status = status slog.Info("starting mlx runner subprocess", "model", c.modelName, "port", c.port) if err := cmd.Start(); err != nil { return nil, fmt.Errorf("failed to start mlx runner: %w", err) } + c.cmd = cmd // Reap subprocess when it exits go func() { From 4512d2b76dd90c1387e4219eccee160ed47a1192 Mon Sep 17 00:00:00 2001 From: frob Date: Fri, 11 Sep 2026 00:46:58 +0200 Subject: [PATCH 10/24] llm: raise token repeat limit to 100 and return error instead of silently closing (#18374) --- llm/llama_server.go | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/llm/llama_server.go b/llm/llama_server.go index fc3edfbbe65..0f40197995e 100644 --- a/llm/llama_server.go +++ b/llm/llama_server.go @@ -1700,9 +1700,9 @@ func (s *llamaServerRunner) Completion(ctx context.Context, req CompletionReques lastToken = strings.TrimSpace(lsResp.Content) tokenRepeat = 0 } - if tokenRepeat > 30 { + if tokenRepeat > 100 { slog.Debug("prediction aborted, token repeat limit reached") - return ctx.Err() + return fmt.Errorf("prediction aborted, token repeat limit reached") } if lsResp.Content != "" && !lsResp.Stop { From 13037ecb14f5159b3be0259710d0e7af981b8a32 Mon Sep 17 00:00:00 2001 From: Jesse Gross Date: Thu, 3 Sep 2026 20:05:49 -0700 Subject: [PATCH 11/24] mlx: scope array lifetimes instead of pinning and sweeping The bindings freed arrays by sweeping everything not pinned, so freeing anything required knowing what every other caller still held, and code that never swept accumulated until memory ran out. The prefix cache's eviction of a long stored path did exactly that: each merge copied the KV snapshots and nothing freed the consumed copies until the request ended, which drove a second long request past physical memory. Every array now belongs to a scope. A function scope, entered with Scoped or one of the ScopedEval forms, frees what was created in it when the function returns; results leave only by being returned. A held scope is closed by its holder and frees what was attached to it. A graph is built in a function scope and evaluated after it, so the eval frees each intermediate as it consumes it. Pin, Unpin, Sweep, and the array list's mutex are gone. On an M5 Max with qwen3.8:27b-mlx, the second 84k-token request after a stored one peaks at 35 GB instead of 57 GB; the cold path is unchanged. The copies themselves are untouched, so restoring an owned path can still exceed memory. --- x/create/draft.go | 2 +- x/create/mlxthread.go | 5 +- x/create/pipeline.go | 2 +- x/create/quantize.go | 176 ++++++++++++------------ x/mlxrunner/cache/cache.go | 4 +- x/mlxrunner/cache/kvcache.go | 88 ++++++------ x/mlxrunner/cache/lazy_test.go | 101 +++++++------- x/mlxrunner/cache/recurrent.go | 22 +-- x/mlxrunner/cache/rotating.go | 53 ++++---- x/mlxrunner/dflash.go | 109 ++++++++------- x/mlxrunner/grammar.go | 10 +- x/mlxrunner/media.go | 32 ++--- x/mlxrunner/media_test.go | 4 +- x/mlxrunner/mlx/act_test.go | 19 +-- x/mlxrunner/mlx/array.go | 91 ++----------- x/mlxrunner/mlx/compile.go | 25 +--- x/mlxrunner/mlx/compile_test.go | 42 +++--- x/mlxrunner/mlx/scope.go | 203 ++++++++++++++++++++++++++++ x/mlxrunner/mlx/scope_test.go | 151 +++++++++++++++++++++ x/mlxrunner/mlx/thread_test.go | 11 +- x/mlxrunner/model/base/base.go | 19 --- x/mlxrunner/mtp.go | 197 +++++++++++++++------------ x/mlxrunner/mtp_test.go | 15 --- x/mlxrunner/pipeline.go | 208 +++++++++++++++-------------- x/mlxrunner/prefix_cache.go | 18 +-- x/mlxrunner/runner.go | 173 +++++++++++++----------- x/mlxrunner/sample/logprob_test.go | 14 +- x/mlxrunner/sample/sample.go | 145 +++++++++++--------- x/mlxrunner/sample/sample_test.go | 11 -- x/mlxrunner/server.go | 5 +- x/mlxrunner/speculate.go | 116 ++++++++-------- x/models/qwen3_5/vision.go | 2 +- x/models/qwen4_exp/engram_cache.go | 26 ++-- 33 files changed, 1184 insertions(+), 915 deletions(-) create mode 100644 x/mlxrunner/mlx/scope.go create mode 100644 x/mlxrunner/mlx/scope_test.go diff --git a/x/create/draft.go b/x/create/draft.go index 759b2fd2115..dfcaeca91de 100644 --- a/x/create/draft.go +++ b/x/create/draft.go @@ -22,7 +22,7 @@ func CreateDraftLayers(modelDir, tensorPrefix, configPrefix, quantize string, st if configPrefix == "" { return nil, fmt.Errorf("draft config prefix must not be empty") } - defer sweepMLX() + defer releaseMLXCache() inv, err := ReadInventory(modelDir) if err != nil { diff --git a/x/create/mlxthread.go b/x/create/mlxthread.go index 590a15dfba1..faae5e8a2a1 100644 --- a/x/create/mlxthread.go +++ b/x/create/mlxthread.go @@ -59,14 +59,13 @@ func runOnMLXThread(f func() error) error { return <-done } -// sweepMLX releases the MLX buffer cache. It is a no-op if no MLX work has run. -func sweepMLX() { +// releaseMLXCache releases the MLX buffer cache. It is a no-op if no MLX work has run. +func releaseMLXCache() { if !mlxThreadStarted.Load() { return } _ = runOnMLXThread(func() error { mlx.ClearCache() - mlx.Sweep() return nil }) } diff --git a/x/create/pipeline.go b/x/create/pipeline.go index 22ef3f5c82e..927ce581542 100644 --- a/x/create/pipeline.go +++ b/x/create/pipeline.go @@ -14,7 +14,7 @@ import ( // server-side entry point — the caller supplies blob storage (store) and // manifest assembly (writeManifest). func Create(modelName, modelDir, quantize string, store BlobStore, writeManifest ManifestWriter, fn func(status string)) error { - defer sweepMLX() + defer releaseMLXCache() inv, err := ReadInventory(modelDir) if err != nil { diff --git a/x/create/quantize.go b/x/create/quantize.go index b91a356bd50..473f38490af 100644 --- a/x/create/quantize.go +++ b/x/create/quantize.go @@ -43,11 +43,8 @@ func quantizeBlob(items []quantizeItem) ([]byte, error) { func quantizeBlobLocked(items []quantizeItem) ([]byte, error) { allArrays := make(map[string]*mlx.Array) - var pinned []*mlx.Array - defer func() { - mlx.Unpin(pinned...) - mlx.Sweep() - }() + held := mlx.NewScope() + defer held.Close() tmpDir, err := os.MkdirTemp("", "ollama-quantize-*") if err != nil { @@ -81,35 +78,20 @@ func quantizeBlobLocked(items []quantizeItem) ([]byte, error) { } for _, it := range items { - if err := func() error { - defer mlx.Sweep() - tmpPath, toEval, st, err := loadAndQuantizeArray(it.reader, it.name, it.quantize, it.decodeFP8, allArrays, tmpDir) - if tmpPath != "" { - defer os.Remove(tmpPath) - } - if err != nil { - return err - } - if st != nil { - defer st.Free() - } - mlx.Eval(toEval...) - final := arraysForItem(allArrays, it) - mlx.Pin(final...) - pinned = append(pinned, final...) - - if mixed && it.quantize != "" { - if gs, _, _ := quant.Params(it.quantize); gs > 0 { - if metadata == nil { - metadata = make(map[string]string) - } - metadata[it.name+".quant_type"] = it.quantize - metadata[it.name+".group_size"] = strconv.Itoa(gs) + if err := quantizeItemArrays(it, allArrays, tmpDir, held); err != nil { + return nil, err + } + // The item's intermediates are free; hand their buffers back before + // the next item, which may never reuse those sizes. + mlx.ClearCache() + if mixed && it.quantize != "" { + if gs, _, _ := quant.Params(it.quantize); gs > 0 { + if metadata == nil { + metadata = make(map[string]string) } + metadata[it.name+".quant_type"] = it.quantize + metadata[it.name+".group_size"] = strconv.Itoa(gs) } - return nil - }(); err != nil { - return nil, err } } @@ -120,18 +102,20 @@ func quantizeBlobLocked(items []quantizeItem) ([]byte, error) { return os.ReadFile(outPath) } -func arraysForItem(all map[string]*mlx.Array, it quantizeItem) []*mlx.Array { - keys := []string{it.name} - if it.quantize != "" { - keys = append(keys, it.name+".scale", it.name+".bias") +// quantizeItemArrays loads and quantizes one item into arrays and holds its +// finished arrays in held. +func quantizeItemArrays(it quantizeItem, arrays map[string]*mlx.Array, tmpDir string, held *mlx.Scope) error { + tmpPath, toEval, st, err := loadAndQuantizeArray(it.reader, it.name, it.quantize, it.decodeFP8, arrays, tmpDir) + if tmpPath != "" { + defer os.Remove(tmpPath) } - out := make([]*mlx.Array, 0, len(keys)) - for _, k := range keys { - if a := all[k]; a != nil { - out = append(out, a) - } + if err != nil { + return err } - return out + defer st.Free() + mlx.Eval(toEval...) + held.Attach(toEval...) + return nil } // loadAndQuantizeArray writes a safetensors reader to a temp file, loads it @@ -164,65 +148,69 @@ func loadAndQuantizeArray(r io.Reader, name, quantize string, decodeFP8 bool, ar return tmpPath, nil, nil, fmt.Errorf("failed to load safetensors for %s: %w", name, err) } - arr := st.Get(name) - if arr == nil { - st.Free() - return tmpPath, nil, nil, fmt.Errorf("tensor %q not found in safetensors", name) - } - - // Decode an FP8 source tensor (using its block scale) before quantizing, - // so a decode-only request (quantize == "") still yields usable float data. - if decodeFP8 { - scaleKey := name + ".scale_inv" - scaleInv := st.Get(scaleKey) - if scaleInv == nil { - scaleKey = name + ".scale" - scaleInv = st.Get(scaleKey) + toEval = mlx.ScopedArrays(func() []*mlx.Array { + arr := st.Get(name) + if arr == nil { + err = fmt.Errorf("tensor %q not found in safetensors", name) + return nil } - if scaleInv == nil { - st.Free() - return tmpPath, nil, nil, fmt.Errorf("missing companion tensor %q or %q for fp8 source tensor %q", name+".scale_inv", name+".scale", name) + + // Decode an FP8 source tensor (using its block scale) before quantizing, + // so a decode-only request (quantize == "") still yields usable float data. + if decodeFP8 { + scaleKey := name + ".scale_inv" + scaleInv := st.Get(scaleKey) + if scaleInv == nil { + scaleKey = name + ".scale" + scaleInv = st.Get(scaleKey) + } + if scaleInv == nil { + err = fmt.Errorf("missing companion tensor %q or %q for fp8 source tensor %q", name+".scale_inv", name+".scale", name) + return nil + } + arr, err = decodeSourceFP8Tensor(arr, scaleInv) + if err != nil { + err = fmt.Errorf("failed to decode fp8 tensor %s: %w", name, err) + return nil + } } - arr, err = decodeSourceFP8Tensor(arr, scaleInv) - if err != nil { - st.Free() - return tmpPath, nil, nil, fmt.Errorf("failed to decode fp8 tensor %s: %w", name, err) + + if quantize == "" { + arr = mlx.Contiguous(arr, false) + arrays[name] = arr + return []*mlx.Array{arr} } - mlx.Eval(arr) - } - if quantize == "" { - arr = mlx.Contiguous(arr, false) - arrays[name] = arr - return tmpPath, []*mlx.Array{arr}, st, nil - } + if arr.DType() != mlx.DTypeBFloat16 && arr.DType() != mlx.DTypeFloat32 && arr.DType() != mlx.DTypeFloat16 { + arr = arr.AsType(mlx.DTypeBFloat16) + } - if arr.DType() != mlx.DTypeBFloat16 && arr.DType() != mlx.DTypeFloat32 && arr.DType() != mlx.DTypeFloat16 { - arr = arr.AsType(mlx.DTypeBFloat16) - mlx.Eval(arr) - } + groupSize, bits, mode := quant.Params(quantize) + qweight, scales, qbiases := mlx.Quantize(arr, groupSize, bits, mode) + if len(qweight.Dims()) == 0 || qweight.Dims()[0] == 0 { + err = fmt.Errorf("mlx.Quantize produced empty weight for %s (quantize=%s, groupSize=%d, bits=%d, mode=%s)", name, quantize, groupSize, bits, mode) + return nil + } + if len(scales.Dims()) == 0 || scales.Dims()[0] == 0 { + err = fmt.Errorf("mlx.Quantize produced empty scales for %s (quantize=%s, groupSize=%d, bits=%d, mode=%s)", name, quantize, groupSize, bits, mode) + return nil + } - groupSize, bits, mode := quant.Params(quantize) - qweight, scales, qbiases := mlx.Quantize(arr, groupSize, bits, mode) - mlx.Eval(qweight, scales) - if len(qweight.Dims()) == 0 || qweight.Dims()[0] == 0 { - st.Free() - return tmpPath, nil, nil, fmt.Errorf("mlx.Quantize produced empty weight for %s (quantize=%s, groupSize=%d, bits=%d, mode=%s)", name, quantize, groupSize, bits, mode) - } - if len(scales.Dims()) == 0 || scales.Dims()[0] == 0 { + qweight = mlx.Contiguous(qweight, false) + scales = mlx.Contiguous(scales, false) + arrays[name] = qweight + arrays[name+".scale"] = scales + out := []*mlx.Array{qweight, scales} + if qbiases != nil { + qbiases = mlx.Contiguous(qbiases, false) + arrays[name+".bias"] = qbiases + out = append(out, qbiases) + } + return out + }) + if err != nil { st.Free() - return tmpPath, nil, nil, fmt.Errorf("mlx.Quantize produced empty scales for %s (quantize=%s, groupSize=%d, bits=%d, mode=%s)", name, quantize, groupSize, bits, mode) - } - - qweight = mlx.Contiguous(qweight, false) - scales = mlx.Contiguous(scales, false) - arrays[name] = qweight - arrays[name+".scale"] = scales - toEval = append(toEval, qweight, scales) - if qbiases != nil { - qbiases = mlx.Contiguous(qbiases, false) - arrays[name+".bias"] = qbiases - toEval = append(toEval, qbiases) + return tmpPath, nil, nil, err } return tmpPath, toEval, st, nil } diff --git a/x/mlxrunner/cache/cache.go b/x/mlxrunner/cache/cache.go index 828fbeef859..be404a92f65 100644 --- a/x/mlxrunner/cache/cache.go +++ b/x/mlxrunner/cache/cache.go @@ -15,7 +15,7 @@ type Cache interface { Offset() int // Snapshot copies cache state from fromOffset to current offset into - // pinned VRAM arrays. The active cache is unchanged. + // owned VRAM arrays. The active cache is unchanged. Snapshot(fromOffset int) Snapshot // PrepareSnapshots schedules the cache to capture a snapshot as its @@ -67,7 +67,7 @@ type Snapshot interface { // never lazy may treat this as a no-op. SetMaterializeHook(func(delta int)) - // Close unpins the snapshot's arrays so they can be freed by Sweep. + // Close frees the snapshot's arrays. Close() } diff --git a/x/mlxrunner/cache/kvcache.go b/x/mlxrunner/cache/kvcache.go index 4bd91865ba2..68afac07cf3 100644 --- a/x/mlxrunner/cache/kvcache.go +++ b/x/mlxrunner/cache/kvcache.go @@ -23,6 +23,7 @@ type Attention interface { type KVCache struct { keys, values *mlx.Array + scope *mlx.Scope offset int step int @@ -35,7 +36,7 @@ type KVCache struct { } func NewKVCache() *KVCache { - return &KVCache{step: 256} + return &KVCache{step: 256, scope: mlx.NewScope()} } // Assumes B = 1; heterogeneous batches are not supported. @@ -78,7 +79,7 @@ func (c *KVCache) appendKV(keys, values *mlx.Array) (*mlx.Array, *mlx.Array) { c.values.Set(c.values.Concatenate(2, newValues)) } else { c.keys, c.values = newKeys, newValues - mlx.Pin(c.keys, c.values) + c.scope.Attach(c.keys, c.values) } } @@ -133,6 +134,7 @@ func (c *KVCache) captureLazySnapshots(start, end int) { // and cache is nil. type kvSnapshot struct { keys, values *mlx.Array + scope *mlx.Scope // holds keys and values once copied out fromOffset, toOffset int cache *KVCache // issuer while lazy; nil once copied out @@ -154,7 +156,7 @@ func (s *kvSnapshot) Size() int { func (s *kvSnapshot) SetMaterializeHook(fn func(delta int)) { s.onMaterialize = fn } func (s *kvSnapshot) Close() { - mlx.Unpin(s.keys, s.values) + s.scope.Close() if s.cache != nil { s.cache.dropLazySnapshot(s) s.cache = nil @@ -170,14 +172,15 @@ func (s *kvSnapshot) copyOut() { return } c := s.cache - kSlice := c.keys.Slice(mlx.Slice(), mlx.Slice(), mlx.Slice(s.fromOffset, s.toOffset), mlx.Slice()) - vSlice := c.values.Slice(mlx.Slice(), mlx.Slice(), mlx.Slice(s.fromOffset, s.toOffset), mlx.Slice()) - kCopy := mlx.Contiguous(kSlice, false) - vCopy := mlx.Contiguous(vSlice, false) - mlx.Pin(kCopy, vCopy) - mlx.AsyncEval(kCopy, vCopy) - - s.keys, s.values = kCopy, vCopy + copies := mlx.ScopedAsyncEval(func() []*mlx.Array { + kSlice := c.keys.Slice(mlx.Slice(), mlx.Slice(), mlx.Slice(s.fromOffset, s.toOffset), mlx.Slice()) + vSlice := c.values.Slice(mlx.Slice(), mlx.Slice(), mlx.Slice(s.fromOffset, s.toOffset), mlx.Slice()) + return []*mlx.Array{mlx.Contiguous(kSlice, false), mlx.Contiguous(vSlice, false)} + }) + s.scope = mlx.NewScope() + s.scope.Attach(copies...) + + s.keys, s.values = copies[0], copies[1] c.dropLazySnapshot(s) s.cache = nil @@ -249,7 +252,7 @@ func (c *KVCache) Restore(snapshot Snapshot, target int) bool { // Rewind to snapshot start, then feed snapshot. c.offset = snap.fromOffset - c.appendKV(snap.keys, snap.values) + mlx.Scoped(func() { c.appendKV(snap.keys, snap.values) }) // Clamp to target if needed (target may be less than full snapshot). if target < c.offset { @@ -287,20 +290,21 @@ func (c *KVCache) Merge(parent, child Snapshot) Snapshot { p.copyOut() ch.copyOut() - mk := p.keys.Concatenate(2, ch.keys) - mv := p.values.Concatenate(2, ch.values) - mlx.Pin(mk, mv) - mlx.AsyncEval(mk, mv) - - p.Close() - ch.Close() - - return &kvSnapshot{ - keys: mk, - values: mv, + merged := &kvSnapshot{ + scope: mlx.NewScope(), fromOffset: p.fromOffset, toOffset: ch.toOffset, } + joined := mlx.ScopedAsyncEval(func() []*mlx.Array { + joined := []*mlx.Array{p.keys.Concatenate(2, ch.keys), p.values.Concatenate(2, ch.values)} + p.Close() + ch.Close() + return joined + }) + merged.scope.Attach(joined...) + merged.keys, merged.values = joined[0], joined[1] + + return merged } func (c *KVCache) Split(snapshot Snapshot, at int) (Snapshot, Snapshot) { @@ -327,27 +331,23 @@ func (c *KVCache) Split(snapshot Snapshot, at int) (Snapshot, Snapshot) { return p, ch } - pk := mlx.Contiguous(snap.keys.Slice(mlx.Slice(), mlx.Slice(), mlx.Slice(0, splitIdx), mlx.Slice()), false) - pv := mlx.Contiguous(snap.values.Slice(mlx.Slice(), mlx.Slice(), mlx.Slice(0, splitIdx), mlx.Slice()), false) - ck := mlx.Contiguous(snap.keys.Slice(mlx.Slice(), mlx.Slice(), mlx.Slice(splitIdx, seqLen), mlx.Slice()), false) - cv := mlx.Contiguous(snap.values.Slice(mlx.Slice(), mlx.Slice(), mlx.Slice(splitIdx, seqLen), mlx.Slice()), false) - mlx.Pin(pk, pv, ck, cv) - mlx.AsyncEval(pk, pv, ck, cv) - - snap.Close() + p := &kvSnapshot{scope: mlx.NewScope(), fromOffset: snap.fromOffset, toOffset: at} + ch := &kvSnapshot{scope: mlx.NewScope(), fromOffset: at, toOffset: snap.toOffset} + halves := mlx.ScopedAsyncEval(func() []*mlx.Array { + halves := []*mlx.Array{ + mlx.Contiguous(snap.keys.Slice(mlx.Slice(), mlx.Slice(), mlx.Slice(0, splitIdx), mlx.Slice()), false), + mlx.Contiguous(snap.values.Slice(mlx.Slice(), mlx.Slice(), mlx.Slice(0, splitIdx), mlx.Slice()), false), + mlx.Contiguous(snap.keys.Slice(mlx.Slice(), mlx.Slice(), mlx.Slice(splitIdx, seqLen), mlx.Slice()), false), + mlx.Contiguous(snap.values.Slice(mlx.Slice(), mlx.Slice(), mlx.Slice(splitIdx, seqLen), mlx.Slice()), false), + } + snap.Close() + return halves + }) + p.scope.Attach(halves[0], halves[1]) + ch.scope.Attach(halves[2], halves[3]) + p.keys, p.values = halves[0], halves[1] + ch.keys, ch.values = halves[2], halves[3] - p := &kvSnapshot{ - keys: pk, - values: pv, - fromOffset: snap.fromOffset, - toOffset: at, - } - ch := &kvSnapshot{ - keys: ck, - values: cv, - fromOffset: at, - toOffset: snap.toOffset, - } return p, ch } @@ -358,7 +358,7 @@ func (c *KVCache) Free() { for _, s := range slices.Clone(c.lazySnapshots) { s.copyOut() } - mlx.Unpin(c.keys, c.values) + c.scope.Close() c.keys, c.values = nil, nil c.offset = 0 c.snapshots = pendingSnapshots{} diff --git a/x/mlxrunner/cache/lazy_test.go b/x/mlxrunner/cache/lazy_test.go index 4739e7e42d2..01232a610bf 100644 --- a/x/mlxrunner/cache/lazy_test.go +++ b/x/mlxrunner/cache/lazy_test.go @@ -31,10 +31,10 @@ func firstKeyAt(arr *mlx.Array, p, D int) float32 { return arr.Floats()[p*D] } -// settledActiveMemory drains unpinned arrays and the allocator cache, then -// reports active (allocated, in-use) bytes. +// settledActiveMemory drains the allocator cache, then reports active +// (allocated, in-use) bytes. The work between two readings runs in a scope +// so its intermediates are gone by the second one. func settledActiveMemory() int { - mlx.Sweep() mlx.ClearCache() return mlx.ActiveMemory() } @@ -62,26 +62,28 @@ func TestKVSpeculationCaptureAllocatesNothing(t *testing.T) { baseline := settledActiveMemory() - snaps := c.TakeSnapshots() - // Every captured snapshot is a lazy snapshot (no owned buffer). - for i, s := range snaps { - if s == nil { - continue + mlx.Scoped(func() { + snaps := c.TakeSnapshots() + // Every captured snapshot is a lazy snapshot (no owned buffer). + for i, s := range snaps { + if s == nil { + continue + } + if ks := s.(*kvSnapshot); ks.keys != nil { + t.Fatalf("snaps[%d] owns a buffer at capture; want a lazy snapshot", i) + } } - if ks := s.(*kvSnapshot); ks.keys != nil { - t.Fatalf("snaps[%d] owns a buffer at capture; want a lazy snapshot", i) - } - } - // MTP commit: rewind to a partial accept, then discard all snapshots. - if !c.Restore(nil, before+draft/2) { - t.Fatal("live rewind failed") - } - for _, s := range snaps { - if s != nil { - s.Close() + // MTP commit: rewind to a partial accept, then discard all snapshots. + if !c.Restore(nil, before+draft/2) { + t.Fatal("live rewind failed") } - } + for _, s := range snaps { + if s != nil { + s.Close() + } + } + }) after := settledActiveMemory() // Lazy snapshots allocate nothing; allow a tiny slack for allocator noise but @@ -219,25 +221,28 @@ func TestKVLazySnapshotSplitMergeNoCopy(t *testing.T) { base := settledActiveMemory() - // Lazy snapshot [2,10), split at 5. - snap := c.Snapshot(2) - p, ch := c.Split(snap, 5) - ps, cs := p.(*kvSnapshot), ch.(*kvSnapshot) - if ps.keys != nil || cs.keys != nil { - t.Fatal("Split of a lazy snapshot should yield lazy snapshots (no copy)") - } - if ps.fromOffset != 2 || ps.toOffset != 5 || cs.fromOffset != 5 || cs.toOffset != 10 { - t.Fatalf("split ranges = [%d,%d)/[%d,%d), want [2,5)/[5,10)", ps.fromOffset, ps.toOffset, cs.fromOffset, cs.toOffset) - } + var merged *kvSnapshot + mlx.Scoped(func() { + // Lazy snapshot [2,10), split at 5. + snap := c.Snapshot(2) + p, ch := c.Split(snap, 5) + ps, cs := p.(*kvSnapshot), ch.(*kvSnapshot) + if ps.keys != nil || cs.keys != nil { + t.Fatal("Split of a lazy snapshot should yield lazy snapshots (no copy)") + } + if ps.fromOffset != 2 || ps.toOffset != 5 || cs.fromOffset != 5 || cs.toOffset != 10 { + t.Fatalf("split ranges = [%d,%d)/[%d,%d), want [2,5)/[5,10)", ps.fromOffset, ps.toOffset, cs.fromOffset, cs.toOffset) + } - // Merge them back into [2,10). - merged := c.Merge(p, ch).(*kvSnapshot) - if merged.keys != nil { - t.Fatal("Merge of adjacent lazy snapshots should yield a lazy snapshot (no Concatenate)") - } - if merged.fromOffset != 2 || merged.toOffset != 10 { - t.Fatalf("merged range = [%d,%d), want [2,10)", merged.fromOffset, merged.toOffset) - } + // Merge them back into [2,10). + merged = c.Merge(p, ch).(*kvSnapshot) + if merged.keys != nil { + t.Fatal("Merge of adjacent lazy snapshots should yield a lazy snapshot (no Concatenate)") + } + if merged.fromOffset != 2 || merged.toOffset != 10 { + t.Fatalf("merged range = [%d,%d), want [2,10)", merged.fromOffset, merged.toOffset) + } + }) if after := settledActiveMemory(); after > base { t.Fatalf("Split/Merge of lazy snapshots allocated %d bytes; want 0", after-base) @@ -313,15 +318,17 @@ func TestKVRestoreLiveLazySnapshotIsOffsetMove(t *testing.T) { // Restore the snapshot back to 10. Its slots [5,10) were never overwritten, // so it is still lazy and the data is already in the buffer — a pure offset // move, no allocation. - if !c.Restore(snap, 10) { - t.Fatal("restore failed") - } - if snap.keys != nil { - t.Fatal("snapshot was copied out; expected the offset-move fast path") - } - if c.Offset() != 10 { - t.Fatalf("offset after restore = %d, want 10", c.Offset()) - } + mlx.Scoped(func() { + if !c.Restore(snap, 10) { + t.Fatal("restore failed") + } + if snap.keys != nil { + t.Fatal("snapshot was copied out; expected the offset-move fast path") + } + if c.Offset() != 10 { + t.Fatalf("offset after restore = %d, want 10", c.Offset()) + } + }) if after := settledActiveMemory(); after > base { t.Fatalf("restore of a live lazy snapshot allocated %d bytes; want 0 (offset move)", after-base) } diff --git a/x/mlxrunner/cache/recurrent.go b/x/mlxrunner/cache/recurrent.go index c394161f9fb..80c45daad25 100644 --- a/x/mlxrunner/cache/recurrent.go +++ b/x/mlxrunner/cache/recurrent.go @@ -20,6 +20,7 @@ import ( type RecurrentCache struct { convState *mlx.Array deltaState *mlx.Array + scope *mlx.Scope offset int convTail int @@ -77,13 +78,14 @@ func (c *RecurrentCache) captureBoundary(reached int, conv, delta *mlx.Array) { func (c *RecurrentCache) setState(old, v *mlx.Array) *mlx.Array { v = v.Clone() - mlx.Pin(v) - mlx.Unpin(old) + c.scope.Attach(v) + c.scope.Discard(old) return v } func NewRecurrentCache(convTail, convDim, numVHeads, headVDim, headKDim int32) *RecurrentCache { return &RecurrentCache{ + scope: mlx.NewScope(), convTail: int(convTail), convDim: int(convDim), numVHeads: int(numVHeads), @@ -174,26 +176,28 @@ func (c *RecurrentCache) State() []*mlx.Array { // does not depend on any parent state. type recurrentSnapshot struct { convState, deltaState *mlx.Array + scope *mlx.Scope offset int } func (s *recurrentSnapshot) Size() int { return s.convState.NumBytes() + s.deltaState.NumBytes() } -func (s *recurrentSnapshot) Close() { mlx.Unpin(s.convState, s.deltaState) } +func (s *recurrentSnapshot) Close() { s.scope.Close() } // SetMaterializeHook is a no-op: recurrent snapshots own their compact copy from // construction. func (s *recurrentSnapshot) SetMaterializeHook(func(int)) {} -// newRecurrentSnapshot clones and pins conv/delta into an owned snapshot at +// newRecurrentSnapshot clones conv/delta into an owned snapshot at // offset. It does not schedule the eval — capture-path snapshots ride the // cache's State into the caller's batched eval. func newRecurrentSnapshot(conv, delta *mlx.Array, offset int) *recurrentSnapshot { snap := &recurrentSnapshot{ convState: conv.Clone(), deltaState: delta.Clone(), + scope: mlx.NewScope(), offset: offset, } - mlx.Pin(snap.convState, snap.deltaState) + snap.scope.Attach(snap.convState, snap.deltaState) return snap } @@ -226,8 +230,10 @@ func (c *RecurrentCache) Restore(snapshot Snapshot, target int) bool { return false } - c.convState = c.setState(c.convState, snap.convState) - c.deltaState = c.setState(c.deltaState, snap.deltaState) + mlx.Scoped(func() { + c.convState = c.setState(c.convState, snap.convState) + c.deltaState = c.setState(c.deltaState, snap.deltaState) + }) c.offset = snap.offset return true @@ -248,7 +254,7 @@ func (c *RecurrentCache) Split(snapshot Snapshot, at int) (Snapshot, Snapshot) { } func (c *RecurrentCache) Free() { - mlx.Unpin(c.convState, c.deltaState) + c.scope.Close() c.convState, c.deltaState = nil, nil c.offset = 0 c.snapshots = pendingSnapshots{} diff --git a/x/mlxrunner/cache/rotating.go b/x/mlxrunner/cache/rotating.go index 68171ff77c2..8aa3f9c2da7 100644 --- a/x/mlxrunner/cache/rotating.go +++ b/x/mlxrunner/cache/rotating.go @@ -13,6 +13,7 @@ import ( // RotatingKVCache implements sliding window attention with bounded memory. type RotatingKVCache struct { keys, values *mlx.Array + scope *mlx.Scope offset int step int maxSize int @@ -29,7 +30,7 @@ type RotatingKVCache struct { } func NewRotatingKVCache(maxSize int) *RotatingKVCache { - return &RotatingKVCache{maxSize: maxSize, step: 256} + return &RotatingKVCache{maxSize: maxSize, step: 256, scope: mlx.NewScope()} } // Assumes B = 1; heterogeneous batches are not supported. @@ -65,7 +66,7 @@ func (c *RotatingKVCache) concat(keys, values *mlx.Array) (newK *mlx.Array, newV if c.keys == nil { c.keys, c.values = keys.Clone(), values.Clone() - mlx.Pin(c.keys, c.values) + c.scope.Attach(c.keys, c.values) } else { if c.idx < c.keys.Dim(2) { if c.offset <= c.maxSize { @@ -123,7 +124,7 @@ func (c *RotatingKVCache) update(keys, values *mlx.Array) (*mlx.Array, *mlx.Arra c.values.Set(c.values.Concatenate(2, newValues)) } else { c.keys, c.values = newKeys, newValues - mlx.Pin(c.keys, c.values) + c.scope.Attach(c.keys, c.values) } c.idx = prev } @@ -194,18 +195,18 @@ func (c *RotatingKVCache) State() []*mlx.Array { } } -// replaceBuffer swaps in newK/newV as the cache's keys/values, unpinning the old -// buffer and pinning the new one. +// replaceBuffer swaps in newK/newV as the cache's keys/values, releasing the +// old buffer and holding the new one. func (c *RotatingKVCache) replaceBuffer(newK, newV *mlx.Array) { - mlx.Unpin(c.keys, c.values) + c.scope.Discard(c.keys, c.values) c.keys, c.values = newK, newV - mlx.Pin(c.keys, c.values) + c.scope.Attach(c.keys, c.values) } func (c *RotatingKVCache) Free() { // Freeing drops the buffer lazy snapshots index into; copy them out first. c.copyOutLazySnapshots() - mlx.Unpin(c.keys, c.values) + c.scope.Close() c.keys, c.values = nil, nil c.offset = 0 c.idx = 0 @@ -283,6 +284,7 @@ func (c *RotatingKVCache) lazyRotatingSnapshot(o int) Snapshot { // reorders or drops those slots. type rotatingSnapshot struct { keys, values *mlx.Array // owned window once copied out; nil while lazy + scope *mlx.Scope // holds keys and values once copied out fromOffset, toOffset int // absolute offset range the window covers idx int // buffer write position a restore installs @@ -307,7 +309,7 @@ func (s *rotatingSnapshot) Size() int { func (s *rotatingSnapshot) SetMaterializeHook(fn func(delta int)) { s.onMaterialize = fn } func (s *rotatingSnapshot) Close() { - mlx.Unpin(s.keys, s.values) + s.scope.Close() if s.cache != nil { s.cache.dropLazySnapshot(s) s.cache = nil @@ -322,14 +324,15 @@ func (s *rotatingSnapshot) copyOut() { return } c := s.cache - kSlice := c.keys.Slice(mlx.Slice(), mlx.Slice(), mlx.Slice(s.sliceStart, s.sliceEnd), mlx.Slice()) - vSlice := c.values.Slice(mlx.Slice(), mlx.Slice(), mlx.Slice(s.sliceStart, s.sliceEnd), mlx.Slice()) - k := mlx.Contiguous(kSlice, false) - v := mlx.Contiguous(vSlice, false) - mlx.Pin(k, v) - mlx.AsyncEval(k, v) - - s.keys, s.values = k, v + copies := mlx.ScopedAsyncEval(func() []*mlx.Array { + kSlice := c.keys.Slice(mlx.Slice(), mlx.Slice(), mlx.Slice(s.sliceStart, s.sliceEnd), mlx.Slice()) + vSlice := c.values.Slice(mlx.Slice(), mlx.Slice(), mlx.Slice(s.sliceStart, s.sliceEnd), mlx.Slice()) + return []*mlx.Array{mlx.Contiguous(kSlice, false), mlx.Contiguous(vSlice, false)} + }) + s.scope = mlx.NewScope() + s.scope.Attach(copies...) + + s.keys, s.values = copies[0], copies[1] c.dropLazySnapshot(s) s.cache = nil @@ -363,11 +366,13 @@ func (c *RotatingKVCache) Snapshot(fromOffset int) Snapshot { state := c.State() k := state[0].Clone() v := state[1].Clone() - mlx.Pin(k, v) + scope := mlx.NewScope() + scope.Attach(k, v) return &rotatingSnapshot{ keys: k, values: v, + scope: scope, fromOffset: fromOffset, toOffset: c.offset, idx: c.idx, @@ -414,10 +419,12 @@ func (c *RotatingKVCache) Restore(snapshot Snapshot, target int) bool { c.dropLazySnapshot(snap) c.copyOutLazySnapshots() liveLen := snap.sliceEnd - snap.sliceStart - c.replaceBuffer( - c.keys.Slice(mlx.Slice(), mlx.Slice(), mlx.Slice(snap.sliceStart, snap.sliceEnd), mlx.Slice()), - c.values.Slice(mlx.Slice(), mlx.Slice(), mlx.Slice(snap.sliceStart, snap.sliceEnd), mlx.Slice()), - ) + mlx.Scoped(func() { + c.replaceBuffer( + c.keys.Slice(mlx.Slice(), mlx.Slice(), mlx.Slice(snap.sliceStart, snap.sliceEnd), mlx.Slice()), + c.values.Slice(mlx.Slice(), mlx.Slice(), mlx.Slice(snap.sliceStart, snap.sliceEnd), mlx.Slice()), + ) + }) snap.sliceStart, snap.sliceEnd = 0, liveLen c.lazySnapshots = append(c.lazySnapshots, snap) c.offset = snap.toOffset @@ -435,7 +442,7 @@ func (c *RotatingKVCache) Restore(snapshot Snapshot, target int) bool { snap.copyOut() c.copyOutLazySnapshots() - c.replaceBuffer(snap.keys.Clone(), snap.values.Clone()) + mlx.Scoped(func() { c.replaceBuffer(snap.keys.Clone(), snap.values.Clone()) }) c.offset = snap.toOffset c.idx = snap.idx diff --git a/x/mlxrunner/dflash.go b/x/mlxrunner/dflash.go index 7dad4ad1f38..628c0bf43fc 100644 --- a/x/mlxrunner/dflash.go +++ b/x/mlxrunner/dflash.go @@ -8,7 +8,7 @@ import ( "github.com/ollama/ollama/x/mlxrunner/model/base" ) -// dflashPendingFlushTokens bounds the pinned feature rows between flushes. +// dflashPendingFlushTokens bounds the held feature rows between flushes. const dflashPendingFlushTokens = 256 // dflashDrafter drafts with a block-diffusion draft model (DFlash): one @@ -48,6 +48,7 @@ type dflashDraftSession struct { ctxOffset int pendingFeatures []*mlx.Array pendingCount int + pending *mlx.Scope // holds the rows until the flush // blockOutstanding tracks the proposal's scheduled rollback point, which // commitBlock has to drain even when it needs no rewind. @@ -66,7 +67,10 @@ func (d *dflashDraftSession) committed(tokens, features *mlx.Array, position int } if start < n { f := features.Slice(mlx.Slice(), mlx.Slice(start, n), mlx.Slice()) - mlx.Pin(f) + if d.pending == nil { + d.pending = mlx.NewScope() + } + d.pending.Attach(f) d.pendingFeatures = append(d.pendingFeatures, f) d.pendingCount += n - start if d.pendingCount >= dflashPendingFlushTokens { @@ -91,7 +95,8 @@ func (d *dflashDraftSession) takePending() *mlx.Array { return nil } features := mlx.Concatenate(d.pendingFeatures, 1) - mlx.Unpin(d.pendingFeatures...) + d.pending.Close() + d.pending = nil d.pendingFeatures = nil d.ctxOffset += d.pendingCount d.pendingCount = 0 @@ -116,24 +121,25 @@ func (d *dflashDraftSession) flush() { spec := d.drafter.spec d.commitBlock() - offset := d.ctxOffset - features := d.takePending() - if features == nil { + if len(d.pendingFeatures) == 0 { return } - spec.draft.Forward(&batch.Batch{ - SeqOffsets: []int32{int32(offset)}, - Hidden: features, - Layout: d.layout, - }, spec.targets, spec.draftKV) - - // Force the cache writes: a session that never drafts would otherwise - // leave the flush chain unevaluated, pinning every feature until close. - state := make([]*mlx.Array, 0, 2*len(spec.draftKV)) - for _, c := range spec.draftKV { - state = append(state, c.State()...) - } - mlx.AsyncEval(state...) + offset := d.ctxOffset + // Evaluating the cache state forces the writes: a session that never + // drafts would otherwise leave the flush chain unevaluated, holding + // every feature until close. + mlx.ScopedAsyncEval(func() []*mlx.Array { + spec.draft.Forward(&batch.Batch{ + SeqOffsets: []int32{int32(offset)}, + Hidden: d.takePending(), + Layout: d.layout, + }, spec.targets, spec.draftKV) + state := make([]*mlx.Array, 0, 2*len(spec.draftKV)) + for _, c := range spec.draftKV { + state = append(state, c.State()...) + } + return state + }) } // propose drafts a block after the not-yet-validated current token, one @@ -152,34 +158,39 @@ func (d *dflashDraftSession) propose(current *mlx.Array, maxTokens int) *draftCa // Send only the anchor and the rows being sampled, not the full trained // block. Exact for causal layers, and measured as free for bidirectional // ones. - masks := make([]int32, n) - for i := range masks { - masks[i] = d.drafter.maskToken - } - block := current.ExpandDims(-1).Concatenate(1, mlx.FromValues(masks, 1, len(masks))) - - offset := d.ctxOffset - features := d.takePending() - - scheduleSpeculation(spec.draftKV, d.ctxOffset, 1) - d.blockOutstanding = true - - hidden, _ := spec.draft.Forward(&batch.Batch{ - InputIDs: block, - SeqOffsets: []int32{int32(offset)}, - SeqQueryLens: []int32{int32(n + 1)}, - Hidden: features, - Layout: d.layout, - }, spec.targets, spec.draftKV) - - // Row i predicts the token at its own position, so the anchor row is - // unused. Rows 1..n are sampled from one batched distribution; penalties - // see only the committed history, not the other rows of the block. - logits := spec.draft.Unembed(hidden.Slice(mlx.Slice(), mlx.Slice(1, n+1), mlx.Slice())) - dist := r.Sampler.Distribution(pipelineSlot, logits, nil) - tokens := r.Sampler.SampleDistribution(pipelineSlot, dist) - return &draftCandidates{ - tokens: tokens.ExpandDims(0), - dist: dist, - } + var candidates *draftCandidates + mlx.ScopedArrays(func() []*mlx.Array { + masks := make([]int32, n) + for i := range masks { + masks[i] = d.drafter.maskToken + } + block := current.ExpandDims(-1).Concatenate(1, mlx.FromValues(masks, 1, len(masks))) + + offset := d.ctxOffset + features := d.takePending() + + scheduleSpeculation(spec.draftKV, d.ctxOffset, 1) + d.blockOutstanding = true + + hidden, _ := spec.draft.Forward(&batch.Batch{ + InputIDs: block, + SeqOffsets: []int32{int32(offset)}, + SeqQueryLens: []int32{int32(n + 1)}, + Hidden: features, + Layout: d.layout, + }, spec.targets, spec.draftKV) + + // Row i predicts the token at its own position, so the anchor row is + // unused. Rows 1..n are sampled from one batched distribution; penalties + // see only the committed history, not the other rows of the block. + logits := spec.draft.Unembed(hidden.Slice(mlx.Slice(), mlx.Slice(1, n+1), mlx.Slice())) + dist := r.Sampler.Distribution(pipelineSlot, logits, nil) + tokens := r.Sampler.SampleDistribution(pipelineSlot, dist) + candidates = &draftCandidates{ + tokens: tokens.ExpandDims(0), + dist: dist, + } + return candidates.Arrays() + }) + return candidates } diff --git a/x/mlxrunner/grammar.go b/x/mlxrunner/grammar.go index 511d0cde7e8..e4c49f90189 100644 --- a/x/mlxrunner/grammar.go +++ b/x/mlxrunner/grammar.go @@ -32,7 +32,7 @@ const ( ) // grammarEngine is the runner's structured-output subsystem: xgrammar bound -// to the model's vocabulary, plus the pinned lookup table for expanding +// to the model's vocabulary, plus the held lookup table for expanding // packed token masks on the device. type grammarEngine struct { // compileMu is the single compile slot: one native compile at a time @@ -46,6 +46,7 @@ type grammarEngine struct { maskTable *mlx.Array byteShifts *mlx.Array + scope *mlx.Scope } func newGrammarEngine(logitsWidth int, tokenizer *tokenizer.Tokenizer) *grammarEngine { @@ -91,7 +92,7 @@ func validateGrammarVocab(logitsWidth, tokenizerSize int) error { } // initMask builds the byte-to-mask lookup table for expanding packed token -// masks on the device, pinned for the runner's lifetime: row v holds, for +// masks on the device, held for the runner's lifetime: row v holds, for // each of the byte value v's eight bits low to high, 0 where the bit is set // (token allowed) and -inf where it is clear. func (e *grammarEngine) initMask(vocabSize int) { @@ -106,7 +107,8 @@ func (e *grammarEngine) initMask(vocabSize int) { } e.maskTable = mlx.FromValues(vals, 256, 8) e.byteShifts = mlx.FromValues([]int32{0, 8, 16, 24}, 4) - mlx.Pin(e.maskTable, e.byteShifts) + e.scope = mlx.NewScope() + e.scope.Attach(e.maskTable, e.byteShifts) } func (e *grammarEngine) close() { @@ -116,7 +118,7 @@ func (e *grammarEngine) close() { e.compiler.Close() e.compiler = nil } - mlx.Unpin(e.maskTable, e.byteShifts) + e.scope.Close() e.maskTable, e.byteShifts = nil, nil } diff --git a/x/mlxrunner/media.go b/x/mlxrunner/media.go index c29a1d7f388..03d0ff8cc28 100644 --- a/x/mlxrunner/media.go +++ b/x/mlxrunner/media.go @@ -54,6 +54,7 @@ type requestMedia struct { // toggled in place so every batch shares the same slice. manifest []batch.MediaItem features []*mlx.Array // parallel to items; nil until encoded + scope *mlx.Scope // layout is the request's one-row Batch.Layout, shared by every batch // like the manifest; nil when the model returned no layout. @@ -70,6 +71,7 @@ func (r *Runner) openMedia(request Request) *requestMedia { inputLen: len(request.Tokens), manifest: make([]batch.MediaItem, len(request.MediaItems)), features: make([]*mlx.Array, len(request.MediaItems)), + scope: mlx.NewScope(), } if request.Layout != nil { m.layout = []any{request.Layout} @@ -114,7 +116,7 @@ func (m *requestMedia) extendChunk(pos, n int) int { } // batchMedia returns the manifest for chunk [pos, pos+n), encoding and -// pinning each item's features on first overlap; nothing evaluates here — +// holding each item's features on first overlap; nothing evaluates here — // the consuming forward pulls the encoder. func (m *requestMedia) batchMedia(pos, n int) []batch.MediaItem { if m == nil { @@ -125,9 +127,11 @@ func (m *requestMedia) batchMedia(pos, n int) []batch.MediaItem { continue } if m.features[i] == nil { - data := mlx.FromValues(item.item.MediaData, item.item.Dims...) - m.features[i] = m.model.EncodeMedia(item.item, data) - mlx.Pin(m.features[i]) + m.features[i] = mlx.ScopedArrays(func() []*mlx.Array { + data := mlx.FromValues(item.item.MediaData, item.item.Dims...) + return []*mlx.Array{m.model.EncodeMedia(item.item, data)} + })[0] + m.scope.Attach(m.features[i]) // The upload copied the pixels; free them here — release never // passes the end of an expansion reaching the prompt's last token. item.item.MediaData = nil @@ -137,9 +141,9 @@ func (m *requestMedia) batchMedia(pos, n int) []batch.MediaItem { return m.manifest } -// release frees what items fully evaluated or restored at position pos no -// longer need: the pinned features and the preprocessed pixel buffer. -func (m *requestMedia) release(pos int) { +// free frees what items fully evaluated or restored at position pos no +// longer need: the held features and the preprocessed pixel buffer. +func (m *requestMedia) free(pos int) { if m == nil { return } @@ -147,7 +151,7 @@ func (m *requestMedia) release(pos int) { if item.pos+item.length <= pos { item.item.MediaData = nil if m.features[i] != nil { - mlx.Unpin(m.features[i]) + m.scope.Discard(m.features[i]) m.features[i] = nil m.manifest[i].Features = nil } @@ -155,18 +159,16 @@ func (m *requestMedia) release(pos int) { } } -// close unpins whatever remains when the pipeline exits. +// close frees whatever remains when the pipeline exits. func (m *requestMedia) close() { if m == nil { return } - for i, f := range m.features { - if f != nil { - mlx.Unpin(f) - m.features[i] = nil - m.manifest[i].Features = nil - } + for i := range m.features { + m.features[i] = nil + m.manifest[i].Features = nil } + m.scope.Close() } // expandMedia tokenizes the [img-N]-tagged prompt into segments, expands diff --git a/x/mlxrunner/media_test.go b/x/mlxrunner/media_test.go index b1052a5dd1a..ab150f8c590 100644 --- a/x/mlxrunner/media_test.go +++ b/x/mlxrunner/media_test.go @@ -105,11 +105,11 @@ func TestBatchMediaLifecycle(t *testing.T) { t.Fatalf("second overlap re-encoded (calls=%d)", calls) } - m.release(4) + m.free(4) if m.manifest[0].Features == nil { t.Fatal("release dropped features before the expansion was evaluated") } - m.release(6) + m.free(6) if m.manifest[0].Features != nil { t.Fatal("release kept features past the expansion end") } diff --git a/x/mlxrunner/mlx/act_test.go b/x/mlxrunner/mlx/act_test.go index cb330ced367..c2c8ff5dcb4 100644 --- a/x/mlxrunner/mlx/act_test.go +++ b/x/mlxrunner/mlx/act_test.go @@ -23,7 +23,6 @@ func TestGELUCompiledMatchesEager(t *testing.T) { withMLXThread(t, func(t *mlxthreadtest.T) { EnableCompile() input := FromValues(values, len(values)).AsType(tt.dtype) - Pin(input) want := gelu(input) got := GELU(input) @@ -38,7 +37,6 @@ func TestGELUCompiledMatchesEager(t *testing.T) { t.Fatalf("%s GELU[%d] = %v, want %v (delta %v)", tt.name, i, gotValues[i], wantValues[i], delta) } } - Unpin(input) }) }) } @@ -58,22 +56,13 @@ func benchmarkGELU(b *testing.B, fn func(*Array) *Array) { EnableCompile() input := AddScalar(Zeros(DTypeBFloat16, 1, 4096, 8192), 1) Eval(input) - Pin(input) - defer func() { - Unpin(input) - Sweep() - ClearCache() - }() + defer ClearCache() - warmup := fn(input) - Eval(warmup) - Sweep() + Scoped(func() { Eval(fn(input)) }) b.ResetTimer() for range b.N { - output := fn(input) - Eval(output) - Sweep() + Scoped(func() { Eval(fn(input)) }) } return nil }); err != nil { @@ -85,8 +74,6 @@ func TestReLUSquared(t *testing.T) { var got []float32 withMLXThread(t, func(t *mlxthreadtest.T) { x := FromValues([]float32{-2, -0, 0.5, 2}, 4) - Pin(x) - defer Unpin(x) y := ReLUSquared(x) Eval(y) diff --git a/x/mlxrunner/mlx/array.go b/x/mlxrunner/mlx/array.go index bc3f07b6fbd..675be0c86e8 100644 --- a/x/mlxrunner/mlx/array.go +++ b/x/mlxrunner/mlx/array.go @@ -8,40 +8,22 @@ import ( "fmt" "log/slog" "reflect" - "sort" "strings" - "sync" - "sync/atomic" "unsafe" - - "github.com/ollama/ollama/logutil" ) +// An Array's lifetime is governed by the scope it belongs to; see scope.go. type Array struct { - ctx C.mlx_array - name string - pinned atomic.Int32 + ctx C.mlx_array + name string + scope *Scope } -var ( - arrays []*Array - arraysMu sync.Mutex -) - // constructor utilities func New(name string) *Array { t := &Array{name: name} - - if tracing { - traceScratch = append(traceScratch, t) - } else { - arraysMu.Lock() - defer arraysMu.Unlock() - - arrays = append(arrays, t) - } - + currentScope.take(t) return t } @@ -135,52 +117,17 @@ func (t *Array) Clone() *Array { return tt } -// lifecycle utilities - -// Pin marks arrays as in-use so they are retained during Sweep. -func Pin(s ...*Array) { - for _, t := range s { - if t != nil { - t.pinned.Add(1) - } - } -} - -// Unpin marks arrays as no longer in-use, allowing Sweep to free them. -func Unpin(s ...*Array) { - for _, t := range s { - if t != nil { - if t.pinned.Add(-1) < 0 { - panic(fmt.Sprintf("mlx.Unpin: negative pin count on array %q", t.name)) - } - } - } -} - -// Sweep releases all unpinned arrays, primarily intermediate tensors. MLX will truly -// free them when there are no other references, including dependencies in the graph. -func Sweep() { - arraysMu.Lock() - defer arraysMu.Unlock() - n := 0 - for _, t := range arrays { - if t.pinned.Load() > 0 && t.Valid() { - arrays[n] = t - n++ - } else if t.Valid() { - mlxCheck(C.mlx_array_free(t.ctx)) - t.ctx.ctx = nil - } - } - arrays = arrays[:n] -} - // misc. utilities func (t *Array) Valid() bool { return t.ctx.ctx != nil } +func (t *Array) free() { + mlxCheck(C.mlx_array_free(t.ctx)) + t.ctx.ctx = nil +} + func (t *Array) String() string { str := mlxCheck(C.mlx_string_new()) mlxCheck(C.mlx_array_tostring(&str, t.ctx)) @@ -191,7 +138,6 @@ func (t *Array) String() string { func (t *Array) LogValue() slog.Value { attrs := []slog.Attr{ slog.String("name", t.name), - slog.Int("pinned", int(t.pinned.Load())), } if t.Valid() { attrs = append(attrs, @@ -288,20 +234,3 @@ func (t *Array) Save(name string) error { } return nil } - -// LogArrays logs all live arrays, sorted by size -func LogArrays() { - arraysMu.Lock() - defer arraysMu.Unlock() - sort.Slice(arrays, func(i, j int) bool { - return arrays[i].NumBytes() > arrays[j].NumBytes() - }) - - var total int - for _, t := range arrays { - nb := t.NumBytes() - total += nb - logutil.Trace(fmt.Sprintf("tensor %-60s %5s %5s pinned=%d %v", t.name, t.DType(), PrettyBytes(nb), t.pinned.Load(), t.Dims())) - } - logutil.Trace(fmt.Sprintf("tensors total: %d, size: %s, active: %s", len(arrays), PrettyBytes(total), PrettyBytes(ActiveMemory()))) -} diff --git a/x/mlxrunner/mlx/compile.go b/x/mlxrunner/mlx/compile.go index a6f91fc79f5..21cad2c95eb 100644 --- a/x/mlxrunner/mlx/compile.go +++ b/x/mlxrunner/mlx/compile.go @@ -120,10 +120,6 @@ func Compile3(name string, fn func(*Array, *Array, *Array) *Array, opts ...Compi // single-threaded at this level a plain Go bool suffices. var tracing bool -// traceScratch collects arrays created during a compile trace so they can be -// freed as a group when the callback returns. -var traceScratch []*Array - //export closureCallback func closureCallback(res *C.mlx_vector_array, input C.mlx_vector_array, payload unsafe.Pointer) (rc C.int) { defer func() { @@ -136,26 +132,19 @@ func closureCallback(res *C.mlx_vector_array, input C.mlx_vector_array, payload handle := *(*cgo.Handle)(payload) fn := handle.Value().(CompileFunc) - // When tracing, we track all of the intermediates that are created and free them separately at the end of - // the process. This will give the effect of a single op - inputs are owned by the original caller (via - // the MLX layer) and outputs are transferred back to MLX to create a new Go side tensor. + // The trace runs in its own scope so its intermediates are freed as a + // group when the callback returns. This gives the effect of a single op: + // inputs are owned by the original caller (via the MLX layer) and outputs + // are transferred back to MLX to create a new Go side tensor. if tracing { panic("mlx: nested compile trace") } tracing = true - traceScratch = nil + s := enterScope() + s.noEscape = true defer func() { - for _, a := range traceScratch { - if a.pinned.Load() > 0 { - panic("mlx: traced array was pinned during compilation") - } - if a.Valid() { - mlxCheck(C.mlx_array_free(a.ctx)) - a.ctx.ctx = nil - } - } tracing = false - traceScratch = nil + exitScope(s) }() n := int(mlxCheck(C.mlx_vector_array_size(input))) diff --git a/x/mlxrunner/mlx/compile_test.go b/x/mlxrunner/mlx/compile_test.go index cd264c69f9c..956b7a7c61b 100644 --- a/x/mlxrunner/mlx/compile_test.go +++ b/x/mlxrunner/mlx/compile_test.go @@ -33,29 +33,27 @@ func testCompileFusion(t *mlxthreadtest.T) { a := FromValues(data, n) b := FromValues(data, n) - Pin(a, b) - defer Unpin(a, b) // Compiled: ops fused into a single kernel. EnableCompile() fn := Compile2("diamond", body, Shapeless()) - warm := fn(a, b) - Eval(warm) - Sweep() + Scoped(func() { Eval(fn(a, b)) }) ClearCache() ResetPeakMemory() - y := fn(a, b) - Eval(y) - compiledPeak := PeakMemory() - Sweep() + var compiledPeak int + Scoped(func() { + Eval(fn(a, b)) + compiledPeak = PeakMemory() + }) // Uncompiled: ops evaluated individually, intermediates materialized. ClearCache() ResetPeakMemory() - z := body(a, b) - Eval(z) - uncompiledPeak := PeakMemory() - Sweep() + var uncompiledPeak int + Scoped(func() { + Eval(body(a, b)) + uncompiledPeak = PeakMemory() + }) if compiledPeak == 0 && uncompiledPeak == 0 { t.Skip("peak memory tracking not available") @@ -88,8 +86,6 @@ func testCompileNested(t *mlxthreadtest.T) { gate := FromValues([]float32{0, 1, 2}, 3) up := FromValues([]float32{1, 1, 1}, 3) - Pin(gate, up) - defer Unpin(gate, up) y := outer(gate, up) Eval(y) @@ -116,8 +112,6 @@ func testCompileCallbackPanicRecovers(t *mlxthreadtest.T) { }) x := FromValues([]float32{1}, 1) - Pin(x) - defer Unpin(x) defer func() { r := recover() @@ -139,26 +133,22 @@ func TestCompileNoTrackingGrowth(t *testing.T) { func testCompileNoTrackingGrowth(t *mlxthreadtest.T) { // Repeated invocations of a compiled kernel should not grow the - // tracked-arrays list; the callback's traceScratch collects - // intermediates during tracing and frees them when the callback returns. + // tracked-arrays list; the callback's scope collects intermediates + // during tracing and frees them when the callback returns. fn := Compile2("mul_add", func(a, b *Array) *Array { return a.Multiply(b).Add(b) }) a := FromValues([]float32{1, 2}, 2) b := FromValues([]float32{3, 4}, 2) - Pin(a, b) - defer Unpin(a, b) - Sweep() - before := len(arrays) + before := len(currentScope.arrays) for range 100 { - _ = fn(a, b) - Sweep() + Scoped(func() { _ = fn(a, b) }) } - after := len(arrays) + after := len(currentScope.arrays) if after > before+2 { t.Fatalf("tracked arrays grew from %d to %d across 100 calls (includes initial trace)", before, after) } diff --git a/x/mlxrunner/mlx/scope.go b/x/mlxrunner/mlx/scope.go new file mode 100644 index 00000000000..92156246a5a --- /dev/null +++ b/x/mlxrunner/mlx/scope.go @@ -0,0 +1,203 @@ +package mlx + +import ( + "fmt" +) + +// Array lifetimes +// +// Every handle belongs to a scope, a set of arrays freed together. A +// function scope is entered with Scoped, ScopedArrays, ScopedEval, or +// ScopedAsyncEval and ends when the function returns; a held scope is +// created with NewScope and ends when its holder closes it. A function +// scope frees what was created in it or detached into it; a held scope +// frees what was attached to it. An array moves between scopes in three +// ways only: by being returned from a function scope to the caller's scope, +// by Attach into a held scope, or by Detach from a held scope back into the +// current one for a caller that still reads them. +// +// MLX frees a buffer once no handle and no queued graph refers to it, so a +// graph is built in one function scope and evaluated after that scope ends, +// and the eval frees each intermediate as it consumes it. The function that +// finishes a graph opens the scope, returns the arrays that leave it, and +// discards or closes inside it whatever the graph consumed. An operation +// that only adds to its caller's graph, a model forward, a cache update, a +// sampling distribution, builds into the open scope and never opens its +// own; one that finishes a graph of its own, a cache copy, a sample, opens +// one like any other finisher. Whoever needs the values evaluates them once +// the scope has ended: the finisher itself with ScopedEval or +// ScopedAsyncEval, or its caller with Eval or AsyncEval. A holder holds +// what outlives the function that produced it, and only a holder gives its +// arrays up. + +type Scope struct { + arrays []*Array + parent *Scope + // noEscape refuses to let an array move out of the scope. + noEscape bool +} + +// Function scopes + +// Scoped runs fn in a function scope. Arrays created or released inside it +// are freed when fn returns. +func Scoped(fn func()) { + s := enterScope() + defer exitScope(s) + fn() +} + +// ScopedArrays runs fn in a function scope and moves the arrays it returns to +// the caller's scope. +func ScopedArrays(fn func() []*Array) []*Array { + s := enterScope() + defer exitScope(s) + ts := fn() + escape(ts...) + return ts +} + +// ScopedEval runs fn in a function scope, moves the arrays it returns to the +// caller's scope, ends the scope, and then evaluates them. +func ScopedEval(fn func() []*Array) []*Array { + ts := ScopedArrays(fn) + Eval(ts...) + return ts +} + +// ScopedAsyncEval is ScopedEval with an asynchronous evaluation. +func ScopedAsyncEval(fn func() []*Array) []*Array { + ts := ScopedArrays(fn) + AsyncEval(ts...) + return ts +} + +// Held scopes + +func NewScope() *Scope { + return &Scope{} +} + +// Attach takes arrays from the function scope that built them or from the +// root. An array some held scope already holds, this one included, is that +// holder's to discard or detach first: a second Attach means two owners. +func (s *Scope) Attach(arrays ...*Array) { + for _, t := range arrays { + if t == nil { + continue + } + if t.scope != rootScope && t.scope.parent == nil { + panic(fmt.Sprintf("mlx: array %q is already held", t.name)) + } + s.take(t) + } +} + +// Detach moves arrays back to the current scope, for a caller that still +// reads them. +func (s *Scope) Detach(arrays ...*Array) { + for _, t := range arrays { + if t == nil { + continue + } + if t.scope != s { + panic(fmt.Sprintf("mlx: array %q is not held by this scope", t.name)) + } + currentScope.take(t) + } +} + +func (s *Scope) Discard(arrays ...*Array) { + for _, t := range arrays { + if t == nil { + continue + } + if t.scope != s { + panic(fmt.Sprintf("mlx: array %q is not held by this scope", t.name)) + } + s.remove(t) + t.free() + } +} + +// Close frees whatever the scope still holds. A nil *Scope holds nothing. +func (s *Scope) Close() { + if s == nil { + return + } + s.end() +} + +// Internals + +var ( + rootScope = &Scope{} + currentScope = rootScope +) + +func enterScope() *Scope { + s := &Scope{parent: currentScope} + currentScope = s + return s +} + +// exitScope ends s. +func exitScope(s *Scope) { + if currentScope != s { + panic("mlx: scope exited out of order") + } + currentScope = s.parent + s.end() +} + +// escape moves arrays of the current scope to the caller's scope so they +// outlive the current one. Arrays held elsewhere are left where they are. +func escape(arrays ...*Array) { + if currentScope.parent == nil { + return + } + for _, t := range arrays { + if t != nil && t.scope == currentScope { + currentScope.parent.take(t) + } + } +} + +// end frees the arrays in s. +func (s *Scope) end() { + for _, t := range s.arrays { + t.free() + } + s.arrays = nil +} + +// take moves t into s, out of the scope it was in. +func (s *Scope) take(t *Array) { + if from := t.scope; from != nil { + if from == s { + return + } + if !t.Valid() { + panic(fmt.Sprintf("mlx: array %q used after its scope ended", t.name)) + } + if from.noEscape { + panic(fmt.Sprintf("mlx: array %q escaped a scope that allows no escape", t.name)) + } + from.remove(t) + } + t.scope = s + s.arrays = append(s.arrays, t) +} + +// remove drops t from s's list. An array usually leaves the scope that just +// built it, so the search runs from the end; order in the list is free. +func (s *Scope) remove(t *Array) { + for i := len(s.arrays) - 1; i >= 0; i-- { + if s.arrays[i] == t { + last := len(s.arrays) - 1 + s.arrays[i] = s.arrays[last] + s.arrays = s.arrays[:last] + return + } + } +} diff --git a/x/mlxrunner/mlx/scope_test.go b/x/mlxrunner/mlx/scope_test.go new file mode 100644 index 00000000000..05994d78653 --- /dev/null +++ b/x/mlxrunner/mlx/scope_test.go @@ -0,0 +1,151 @@ +package mlx + +import ( + "testing" + + "github.com/ollama/ollama/x/internal/mlxthreadtest" +) + +// A function scope frees what was created in it. What fn returns moves to the +// caller's scope instead, and a returned array the scope does not own stays +// where it is. +func TestScopeFreesWhatIsNotReturned(t *testing.T) { + withMLXThread(t, func(t *mlxthreadtest.T) { + held := NewScope() + defer held.Close() + var kept, returned, dropped *Array + Scoped(func() { + kept = FromValue(1) + held.Attach(kept) + out := ScopedArrays(func() []*Array { + dropped = FromValue(2) + return []*Array{FromValue(3), nil, kept} + }) + returned = out[0] + if !returned.Valid() { + t.Fatal("returned array was freed with the scope that created it") + } + if dropped.Valid() { + t.Fatal("array not returned survived its scope") + } + }) + if returned.Valid() { + t.Fatal("returned array survived the scope it was returned into") + } + if !kept.Valid() { + t.Fatal("returning a held array moved it out of its scope") + } + }) +} + +// ScopedEval ends the build scope before it evaluates, so the intermediates +// are gone by then and the returned arrays come back evaluated. +func TestScopedEvalEvaluatesAfterBuild(t *testing.T) { + withMLXThread(t, func(t *mlxthreadtest.T) { + Scoped(func() { + var tmp *Array + out := ScopedEval(func() []*Array { + tmp = FromValue(2) + return []*Array{FromValue(1).Add(tmp)} + }) + if tmp.Valid() { + t.Fatal("intermediate survived the build scope") + } + if !out[0].Valid() || out[0].Int() != 3 { + t.Fatal("returned array was not evaluated after the build scope") + } + }) + }) +} + +// A held scope keeps arrays past the function scope that created them. +// Discard frees one now, Detach hands one to the current scope, and Close +// frees what remains. A scope refuses to discard or detach an array it does +// not hold, and to attach one that is already held, by itself or another +// scope. +func TestHeldScope(t *testing.T) { + withMLXThread(t, func(t *mlxthreadtest.T) { + held, other := NewScope(), NewScope() + defer other.Close() + var kept, discarded, detached *Array + Scoped(func() { + kept, discarded, detached = FromValue(1), FromValue(2), FromValue(3) + held.Attach(kept, discarded, detached) + held.Discard(discarded) + if discarded.Valid() { + t.Fatal("discarded array survived") + } + Scoped(func() { held.Detach(detached) }) + if detached.Valid() { + t.Fatal("detached array survived the scope it was detached into") + } + if !panics(func() { other.Discard(kept) }) { + t.Fatal("no panic discarding an array the scope does not hold") + } + if !panics(func() { other.Detach(kept) }) { + t.Fatal("no panic detaching an array the scope does not hold") + } + if !panics(func() { other.Attach(kept) }) { + t.Fatal("no panic holding an array another scope holds") + } + if !panics(func() { held.Attach(kept) }) { + t.Fatal("no panic holding an array twice") + } + }) + if !kept.Valid() { + t.Fatal("held array was freed with the scope that created it") + } + held.Close() + if kept.Valid() { + t.Fatal("held array survived its scope's close") + } + }) +} + +// A scope ends when its function panics, so whatever recovers is back in +// the scope it started from. +func TestScopeEndsOnPanic(t *testing.T) { + withMLXThread(t, func(t *mlxthreadtest.T) { + start := currentScope + var a *Array + func() { + defer func() { _ = recover() }() + Scoped(func() { + a = FromValue(1) + panic("build failed") + }) + }() + if a.Valid() { + t.Fatal("array survived the scope that panicked") + } + if currentScope != start { + t.Fatal("registry not back in the caller's scope after a panic") + } + }) +} + +// Nothing built in a compile trace may outlive it: holding a trace array +// fails the compiled call. +func TestCompileTraceRefusesEscapes(t *testing.T) { + withMLXThread(t, func(t *mlxthreadtest.T) { + held := NewScope() + defer held.Close() + double := Compile("scope_test_escape", func(in ...*Array) []*Array { + out := in[0].Add(in[0]) + held.Attach(out) + return []*Array{out} + }) + defer func() { + if recover() == nil { + t.Fatal("no panic for an array held out of a compile trace") + } + }() + Scoped(func() { double(FromValue(1)) }) + }) +} + +func panics(fn func()) (panicked bool) { + defer func() { panicked = recover() != nil }() + fn() + return false +} diff --git a/x/mlxrunner/mlx/thread_test.go b/x/mlxrunner/mlx/thread_test.go index cd3a5160d0a..0cebb04a960 100644 --- a/x/mlxrunner/mlx/thread_test.go +++ b/x/mlxrunner/mlx/thread_test.go @@ -55,11 +55,12 @@ func TestThreadedMLXOperations(t *testing.T) { for range iterations { if err := thread.Do(context.Background(), func() error { - a := FromValues([]float32{1, 2, 3, 4}, 2, 2) - b := Matmul(a, a) - AsyncEval(b) - Eval(b) - Sweep() + Scoped(func() { + a := FromValues([]float32{1, 2, 3, 4}, 2, 2) + b := Matmul(a, a) + AsyncEval(b) + Eval(b) + }) ClearCache() return nil }); err != nil { diff --git a/x/mlxrunner/model/base/base.go b/x/mlxrunner/model/base/base.go index 243ac688ec4..9cef8253cd1 100644 --- a/x/mlxrunner/model/base/base.go +++ b/x/mlxrunner/model/base/base.go @@ -176,22 +176,3 @@ func NewDraft(root *model.Root, target Model) (DraftModel, error) { return fn(root, target) } - -// Weights returns a function that loads model weights, then pins all -// arrays reachable from the model struct and sweeps everything else. -func Weights(m Model) func(map[string]*mlx.Array) error { - return func(tensors map[string]*mlx.Array) error { - if err := m.LoadWeights(tensors); err != nil { - return err - } - - collected := mlx.Collect(m) - for _, arr := range collected { - mlx.Pin(arr) - } - mlx.Sweep() - mlx.Eval(collected...) - - return nil - } -} diff --git a/x/mlxrunner/mtp.go b/x/mlxrunner/mtp.go index d8d8aa8193e..aed607f31dd 100644 --- a/x/mlxrunner/mtp.go +++ b/x/mlxrunner/mtp.go @@ -10,7 +10,7 @@ import ( ) // mtpPendingFlushTokens caps how many committed look-ahead tokens wait in the -// pending buffer before a batched flush, bounding the pinned hidden states +// pending buffer before a batched flush, bounding the held hidden states // regardless of what else triggers a flush. const mtpPendingFlushTokens = 256 @@ -37,7 +37,7 @@ func (d *mtpDrafter) draftLimit() int { return 0 } // open returns the drafting session for one request, its pairing frontier // synced to the draft caches' restored offset. func (d *mtpDrafter) open(layout []any) draftSession { - s := &mtpDraftSession{drafter: d, layout: layout} + s := &mtpDraftSession{drafter: d, layout: layout, scope: mlx.NewScope()} if kv := d.spec.draftKV; len(kv) > 0 { // A restored prefix arrives with the draft caches already written; // pairing resumes from their absolute offset. @@ -54,20 +54,22 @@ func (d *mtpDrafter) open(layout []any) draftSession { type mtpDraftSession struct { drafter *mtpDrafter layout []any + scope *mlx.Scope // frontier is the slot after the last reported token; frontierHidden is - // the pinned target hidden at frontier-1, fused into the next pair. + // the held target hidden at frontier-1, fused into the next pair. frontier int frontierHidden *mlx.Array // committedDraftOffset is the slot after the last pair written to the - // draft caches; later pairs wait pinned in the pending lists until + // draft caches; later pairs wait held in the pending lists until // flushed. pendingCount is the look-ahead tokens those lists hold, summed // across the buffered runs. committedDraftOffset int pendingTokens []*mlx.Array pendingHiddens []*mlx.Array pendingCount int + pending *mlx.Scope // holds the lists' arrays until the flush // heldHidden is the frontier row's pre-unembed hidden and heldAuxHidden // its fusion hidden, carried from the last flush so the first proposal @@ -76,7 +78,7 @@ type mtpDraftSession struct { heldAuxHidden *mlx.Array // pendingMedia holds manifest rows the deferred flush may still embed, - // pinned since prefill releases them after the target's chunk; + // on their own handles since prefill releases them after the target's chunk; // lastDelivered marks each row's newest delivered end. pendingMedia map[int]batch.MediaItem lastDelivered map[int]int @@ -113,7 +115,7 @@ func (d *mtpDraftSession) committed(tokens, hiddens *mlx.Array, position int, me d.setFrontierHidden(lastHiddenRow(hiddens)) } -// captureMedia pins the run's feature-bearing rows for the deferred +// captureMedia holds the run's feature-bearing rows for the deferred // flush, which embeds them after prefill has released the features. A row // spanning chunks arrives once per chunk. func (d *mtpDraftSession) captureMedia(media []batch.MediaItem, end int) { @@ -129,7 +131,8 @@ func (d *mtpDraftSession) captureMedia(media []batch.MediaItem, end int) { d.pendingMedia = make(map[int]batch.MediaItem) d.lastDelivered = make(map[int]int) } - mlx.Pin(item.Features) + item.Features = item.Features.Clone() + d.scope.Attach(item.Features) d.pendingMedia[item.Pos] = item } d.lastDelivered[item.Pos] = end @@ -150,7 +153,8 @@ func (d *mtpDraftSession) flushMedia(embedEnd int) []batch.MediaItem { slices.SortFunc(manifest, func(a, b batch.MediaItem) int { return a.Pos - b.Pos }) for pos, last := range d.lastDelivered { if last <= embedEnd { - mlx.Unpin(d.pendingMedia[pos].Features) + // The flush's forward still reads the row; it dies with the build. + d.scope.Detach(d.pendingMedia[pos].Features) delete(d.pendingMedia, pos) delete(d.lastDelivered, pos) } @@ -160,7 +164,7 @@ func (d *mtpDraftSession) flushMedia(embedEnd int) []batch.MediaItem { func (d *mtpDraftSession) closeMedia() { for pos, item := range d.pendingMedia { - mlx.Unpin(item.Features) + d.scope.Discard(item.Features) delete(d.pendingMedia, pos) delete(d.lastDelivered, pos) } @@ -174,7 +178,7 @@ func (d *mtpDraftSession) settle(next *mlx.Array) { return } if d.frontierHidden != nil && d.frontier-1 == d.committedDraftOffset+d.pendingCount { - d.queueCacheWrites(next.ExpandDims(-1), d.frontierHidden) + d.queueCacheWrites(next.ExpandDims(-1), d.frontierHidden.Clone()) } d.flush() } @@ -184,14 +188,18 @@ func (d *mtpDraftSession) close() { d.closeMedia() d.setFrontierHidden(nil) d.setHeld(nil, nil) + d.scope.Close() } // queueCacheWrites buffers completed draft-cache writes — look-ahead tokens // fused with their target hiddens — flushing once the buffer reaches the token -// cap so the pinned hiddens stay bounded. flush coalesces the buffered writes +// cap so the held hiddens stay bounded. flush coalesces the buffered writes // into one head forward, so a contiguous run lands in a single draft-cache extend. func (d *mtpDraftSession) queueCacheWrites(tokens, hiddens *mlx.Array) { - mlx.Pin(tokens, hiddens) + if d.pending == nil { + d.pending = mlx.NewScope() + } + d.pending.Attach(tokens, hiddens) d.pendingTokens = append(d.pendingTokens, tokens) d.pendingHiddens = append(d.pendingHiddens, hiddens) d.pendingCount += tokens.Dim(1) @@ -216,45 +224,51 @@ func (d *mtpDraftSession) flush() { } } - ids := mlx.Concatenate(d.pendingTokens, 1) - hiddens := mlx.Concatenate(d.pendingHiddens, 1) - // The pair at slot S embeds the look-ahead token S+1, so this flush - // embeds prompt tokens up to committedDraftOffset+len+1. - hidden, auxHidden := spec.draft.Forward(&batch.Batch{ - InputIDs: ids, - SeqOffsets: []int32{int32(d.committedDraftOffset)}, - SeqQueryLens: []int32{int32(ids.Dim(1))}, - Hidden: hiddens, - Media: d.flushMedia(d.committedDraftOffset + ids.Dim(1) + 1), - Layout: d.layout, - }, spec.targets, spec.draftKV) - d.setHeld(lastHiddenRow(hidden), lastHiddenRow(auxHidden)) - d.committedDraftOffset += ids.Dim(1) - - // Force the draft writes: a session that never drafts would otherwise - // leave the flush chain unevaluated, pinning every hidden until close. - state := make([]*mlx.Array, 0, 2*len(spec.draftKV)) - for _, c := range spec.draftKV { - state = append(state, c.State()...) - } - mlx.AsyncEval(state...) + // Evaluating the state forces the draft writes: a session that never + // drafts would otherwise leave the flush chain unevaluated, holding + // every hidden until close. + n := d.pendingCount + out := mlx.ScopedAsyncEval(func() []*mlx.Array { + ids, hiddens := d.takePending() + // The pair at slot S embeds the look-ahead token S+1, so this flush + // embeds prompt tokens up to committedDraftOffset+len+1. + hidden, auxHidden := spec.draft.Forward(&batch.Batch{ + InputIDs: ids, + SeqOffsets: []int32{int32(d.committedDraftOffset)}, + SeqQueryLens: []int32{int32(n)}, + Hidden: hiddens, + Media: d.flushMedia(d.committedDraftOffset + n + 1), + Layout: d.layout, + }, spec.targets, spec.draftKV) + out := []*mlx.Array{lastHiddenRow(hidden), lastHiddenRow(auxHidden)} + for _, c := range spec.draftKV { + out = append(out, c.State()...) + } + return out + }) + d.setHeld(out[0], out[1]) + d.committedDraftOffset += n +} - mlx.Unpin(d.pendingTokens...) - mlx.Unpin(d.pendingHiddens...) - d.pendingTokens, d.pendingHiddens = nil, nil - d.pendingCount = 0 +// takePending returns the buffered pairs as one batch, advancing past them. +func (d *mtpDraftSession) takePending() (ids, hiddens *mlx.Array) { + ids = mlx.Concatenate(d.pendingTokens, 1) + hiddens = mlx.Concatenate(d.pendingHiddens, 1) + d.pending.Close() + d.pending, d.pendingTokens, d.pendingHiddens, d.pendingCount = nil, nil, nil, 0 + return ids, hiddens } func (d *mtpDraftSession) setFrontierHidden(h *mlx.Array) { - mlx.Pin(h) - mlx.Unpin(d.frontierHidden) + d.scope.Attach(h) + d.scope.Discard(d.frontierHidden) d.frontierHidden = h } -// setHeld replaces the held flush outputs, pinned until the next flush or close. +// setHeld replaces the held flush outputs, kept until the next flush or close. func (d *mtpDraftSession) setHeld(hidden, auxHidden *mlx.Array) { - mlx.Pin(hidden, auxHidden) - mlx.Unpin(d.heldHidden, d.heldAuxHidden) + d.scope.Attach(hidden, auxHidden) + d.scope.Discard(d.heldHidden, d.heldAuxHidden) d.heldHidden, d.heldAuxHidden = hidden, auxHidden } @@ -276,55 +290,60 @@ func (d *mtpDraftSession) propose(current *mlx.Array, maxTokens int) *draftCandi } } - lastToken := current.ExpandDims(-1) - lastHidden := d.frontierHidden - draftDists := make([]sampler.Distribution, 0, maxTokens) - var prefix *mlx.Array - - for i := range maxTokens { - var hidden, auxHidden *mlx.Array - if i == 0 && len(spec.draftKV) > 0 { - // The settle flush already produced the frontier row; reuse it - // instead of re-running the head. - hidden, auxHidden = d.heldHidden, d.heldAuxHidden - } else { - // A head with draft caches writes each draft token to the next - // draft-cache slot, advancing one per step from the last committed - // slot (the held i==0 step stands in for that slot). A cacheless - // head stays at the last committed slot every step, re-attending - // the committed prefix read-only ("single-position"). - pos := d.frontier - 1 - if len(spec.draftKV) > 0 { - pos = d.frontier - 1 + i + var candidates *draftCandidates + mlx.ScopedArrays(func() []*mlx.Array { + lastToken := current.ExpandDims(-1) + lastHidden := d.frontierHidden + draftDists := make([]sampler.Distribution, 0, maxTokens) + var prefix *mlx.Array + + for i := range maxTokens { + var hidden, auxHidden *mlx.Array + if i == 0 && len(spec.draftKV) > 0 { + // The settle flush already produced the frontier row; reuse it + // instead of re-running the head. + hidden, auxHidden = d.heldHidden, d.heldAuxHidden + } else { + // A head with draft caches writes each draft token to the next + // draft-cache slot, advancing one per step from the last committed + // slot (the held i==0 step stands in for that slot). A cacheless + // head stays at the last committed slot every step, re-attending + // the committed prefix read-only ("single-position"). + pos := d.frontier - 1 + if len(spec.draftKV) > 0 { + pos = d.frontier - 1 + i + } + hidden, auxHidden = spec.draft.Forward(&batch.Batch{ + InputIDs: lastToken, + SeqOffsets: []int32{int32(pos)}, + SeqQueryLens: []int32{1}, + Hidden: lastHidden, + Layout: d.layout, + }, spec.targets, spec.draftKV) + } + // Unembed only the row being sampled, never the batch. + stepLogits := spec.draft.Unembed(hidden).Squeeze(1) + lastHidden = auxHidden + // The chain's earlier drafts ride along as the row's history, so + // penalties shape proposals the same way they shape validation. + dist := r.Sampler.Distribution(pipelineSlot, stepLogits, prefix) + nextToken := r.Sampler.SampleDistribution(pipelineSlot, dist) + + lastToken = nextToken.ExpandDims(-1) + draftDists = append(draftDists, dist) + if prefix == nil { + prefix = lastToken + } else { + prefix = prefix.Concatenate(1, lastToken) } - hidden, auxHidden = spec.draft.Forward(&batch.Batch{ - InputIDs: lastToken, - SeqOffsets: []int32{int32(pos)}, - SeqQueryLens: []int32{1}, - Hidden: lastHidden, - Layout: d.layout, - }, spec.targets, spec.draftKV) } - // Unembed only the row being sampled, never the batch. - stepLogits := spec.draft.Unembed(hidden).Squeeze(1) - lastHidden = auxHidden - // The chain's earlier drafts ride along as the row's history, so - // penalties shape proposals the same way they shape validation. - dist := r.Sampler.Distribution(pipelineSlot, stepLogits, prefix) - nextToken := r.Sampler.SampleDistribution(pipelineSlot, dist) - - lastToken = nextToken.ExpandDims(-1) - draftDists = append(draftDists, dist) - if prefix == nil { - prefix = lastToken - } else { - prefix = prefix.Concatenate(1, lastToken) + candidates = &draftCandidates{ + tokens: prefix, + dist: sampler.ConcatenateDistributions(draftDists), } - } - return &draftCandidates{ - tokens: prefix, - dist: sampler.ConcatenateDistributions(draftDists), - } + return candidates.Arrays() + }) + return candidates } func lastHiddenRow(hidden *mlx.Array) *mlx.Array { diff --git a/x/mlxrunner/mtp_test.go b/x/mlxrunner/mtp_test.go index 881362d8ec0..43142fa25b5 100644 --- a/x/mlxrunner/mtp_test.go +++ b/x/mlxrunner/mtp_test.go @@ -274,8 +274,6 @@ func TestAcceptMTPDraftsGreedyAcceptAll(t *testing.T) { spec := testSpeculationSession(r, caches) current := sampler.Result{Token: mlx.FromValues([]int32{1}, 1)} - unpin := pinAcceptInputs(current, candidates) - defer unpin() results, accepted, observed, err := spec.accept(&position, current, candidates, nil) if err != nil { t.Fatalf("accept: %v", err) @@ -311,8 +309,6 @@ func TestAcceptMTPDraftsGreedyMismatch(t *testing.T) { spec := testSpeculationSession(r, caches) current := sampler.Result{Token: mlx.FromValues([]int32{1}, 1)} - unpin := pinAcceptInputs(current, candidates) - defer unpin() results, accepted, observed, err := spec.accept(&position, current, candidates, nil) if err != nil { t.Fatalf("accept: %v", err) @@ -350,8 +346,6 @@ func TestAcceptMTPDraftsGreedyEOS(t *testing.T) { spec := testSpeculationSession(r, caches) current := sampler.Result{Token: mlx.FromValues([]int32{1}, 1)} - unpin := pinAcceptInputs(current, candidates) - defer unpin() results, accepted, observed, err := spec.accept(&position, current, candidates, nil) if err != nil { t.Fatalf("accept: %v", err) @@ -1377,12 +1371,3 @@ func scriptedCandidates(r *Runner, tokens []int32) *draftCandidates { d.committed(mlx.FromValues([]int32{0}, 1, 1), mlx.Zeros(mlx.DTypeFloat32, 1, 1, mtpTestVocab), 0, nil) return d.propose(mlx.FromValues([]int32{0}, 1), len(tokens)) } - -// pinAcceptInputs pins the arrays accept's caller must keep alive across -// accept's internal sweep — current and the candidate tokens — as the decoder -// and next do in the live engine. It returns the matching unpin. -func pinAcceptInputs(current sampler.Result, candidates *draftCandidates) func() { - arrays := append(current.Arrays(), candidates.tokens) - mlx.Pin(arrays...) - return func() { mlx.Unpin(arrays...) } -} diff --git a/x/mlxrunner/pipeline.go b/x/mlxrunner/pipeline.go index c7be35ee153..188f3610b9c 100644 --- a/x/mlxrunner/pipeline.go +++ b/x/mlxrunner/pipeline.go @@ -90,21 +90,19 @@ func (r *Runner) Prepare(request *Request) (err error) { // The runner serializes requests today so we just use a fixed slot ID. const pipelineSlot = 0 -func (r *Runner) TextGenerationPipeline(ctx context.Context, request Request) error { +func (r *Runner) TextGenerationPipeline(ctx context.Context, request Request) (err error) { mlx.ResetPeakMemory() + mlx.Scoped(func() { err = r.generate(ctx, request) }) + mlx.ClearCache() - defer func() { - r.Sampler.Remove(pipelineSlot) - mlx.Sweep() - mlx.ClearCache() - - if slog.Default().Enabled(context.TODO(), logutil.LevelTrace) { - mlx.LogArrays() - r.cache.dumpTree() - } - slog.Info("peak memory", "size", mlx.PrettyBytes(mlx.PeakMemory())) - }() + if slog.Default().Enabled(context.TODO(), logutil.LevelTrace) { + r.cache.dumpTree() + } + slog.Info("memory", "peak", mlx.PrettyBytes(mlx.PeakMemory()), "held", mlx.PrettyBytes(mlx.ActiveMemory())) + return err +} +func (r *Runner) generate(ctx context.Context, request Request) error { inputs := request.Tokens session := r.cache.begin(inputs, request.MediaItems) @@ -126,6 +124,7 @@ func (r *Runner) TextGenerationPipeline(ctx context.Context, request Request) er // Register the sampler after prefill completes. r.Sampler.Add(pipelineSlot, request.SamplerOpts, inputs) + defer r.Sampler.Remove(pipelineSlot) grammar, err := request.Grammar.resolve(ctx) if err != nil { @@ -166,7 +165,7 @@ func (r *Runner) prefill(ctx context.Context, session *cacheSession, spec *specu snapshotOffsets = append(snapshotOffsets, end) } - materializeCaches := func() { + cacheState := func() []*mlx.Array { state := make([]*mlx.Array, 0, 2*len(caches)) for _, c := range caches { if c == nil { @@ -174,10 +173,7 @@ func (r *Runner) prefill(ctx context.Context, session *cacheSession, spec *specu } state = append(state, c.State()...) } - if len(state) == 0 { - return - } - mlx.Eval(state...) + return state } session.schedulePrefillSnapshots(snapshotOffsets) @@ -185,7 +181,7 @@ func (r *Runner) prefill(ctx context.Context, session *cacheSession, spec *specu total, processed := len(tokens), 0 position := len(inputs) - len(tokens) // Free restored items' buffers now: on a full cache hit the loop never runs. - media.release(position) + media.free(position) for total-processed > 1 { if err := ctx.Err(); err != nil { // Settle the drafter with the next prompt token so the caches @@ -198,26 +194,27 @@ func (r *Runner) prefill(ctx context.Context, session *cacheSession, spec *specu n := min(prefillChunk, total-processed-1) n = media.extendChunk(position, n) - chunkIDs := mlx.FromValues(tokens[processed:processed+n], 1, n) - manifest := media.batchMedia(position, n) - _, auxHidden := r.Model.Forward(&batch.Batch{ - InputIDs: chunkIDs, - SeqOffsets: []int32{int32(position)}, - SeqQueryLens: []int32{int32(n)}, - Media: manifest, - Layout: media.rowLayout(), - }, caches) - // Report to the drafter only after the chunk's eval: a draft flush - // evaluates, and an eval before the sweep cannot free any buffer the - // chunk's live handles retain — on media chunks, the whole vision tower. - mlx.Pin(chunkIDs, auxHidden) - mlx.Sweep() - materializeCaches() - spec.committed(chunkIDs, auxHidden, position, manifest) - mlx.Unpin(chunkIDs, auxHidden) - // Released after committed so the drafter can capture rows its - // deferred flush still embeds. - media.release(position + n) + mlx.Scoped(func() { + chunkIDs := mlx.FromValues(tokens[processed:processed+n], 1, n) + chunkMedia := media.batchMedia(position, n) + auxHidden := mlx.ScopedArrays(func() []*mlx.Array { + _, auxHidden := r.Model.Forward(&batch.Batch{ + InputIDs: chunkIDs, + SeqOffsets: []int32{int32(position)}, + SeqQueryLens: []int32{int32(n)}, + Media: chunkMedia, + Layout: media.rowLayout(), + }, caches) + return []*mlx.Array{auxHidden} + })[0] + mlx.Eval(cacheState()...) + // Report to the drafter only after the chunk's eval: a draft + // flush evaluates. + spec.committed(chunkIDs, auxHidden, position, chunkMedia) + // Freed after committed so the drafter can capture rows its + // deferred flush still embeds. + media.free(position + n) + }) processed += n position += n slog.Info("Prompt processing progress", "processed", processed, "total", total) @@ -290,53 +287,61 @@ func (r *Runner) decode(ctx context.Context, request Request, session *cacheSess return err } - results, err := d.next(request.Options.NumPredict - generated) - if err != nil { - return err - } - - // Record the whole run before streaming any of it: a cancelled - // stream returns early and must not leave the caches ahead of - // session.outputs. - done := false - stream := len(results) - for i, res := range results { - id := res.Token.Int() - session.outputs = append(session.outputs, id) - if done { - continue - } - if r.Tokenizer.IsEOS(id) { - final.DoneReason = 0 - done = true - stream = i - continue + var done bool + var err error + mlx.Scoped(func() { + var results []sampler.Result + results, err = d.next(request.Options.NumPredict - generated) + if err != nil { + return } - generated++ - if generated >= request.Options.NumPredict { - done = true - stream = i + 1 - } - } - for _, res := range results[:stream] { - resp, ok := detok.detokenize(res) - if !ok { - continue + // Record the whole run before streaming any of it: a cancelled + // stream returns early and must not leave the caches ahead of + // session.outputs. + stream := len(results) + for i, res := range results { + id := res.Token.Int() + session.outputs = append(session.outputs, id) + if done { + continue + } + if r.Tokenizer.IsEOS(id) { + final.DoneReason = 0 + done = true + stream = i + continue + } + generated++ + if generated >= request.Options.NumPredict { + done = true + stream = i + 1 + } } - // Two-pass structured output cancels the first pass before its final response. - if request.IncludeIntermediateMetrics { - resp.PromptEvalCount = len(request.Tokens) - resp.PromptEvalCachedCount = final.PromptEvalCachedCount - resp.PromptEvalDuration = promptEval - resp.EvalCount = generated - resp.EvalDuration = time.Since(now) - } - select { - case <-ctx.Done(): - return ctx.Err() - case request.Responses <- resp: + + for _, res := range results[:stream] { + resp, ok := detok.detokenize(res) + if !ok { + continue + } + // Two-pass structured output cancels the first pass before its final response. + if request.IncludeIntermediateMetrics { + resp.PromptEvalCount = len(request.Tokens) + resp.PromptEvalCachedCount = final.PromptEvalCachedCount + resp.PromptEvalDuration = promptEval + resp.EvalCount = generated + resp.EvalDuration = time.Since(now) + } + select { + case <-ctx.Done(): + err = ctx.Err() + return + case request.Responses <- resp: + } } + }) + if err != nil { + return err } if done { @@ -375,6 +380,7 @@ type pipelinedDecoder struct { grammars []*grammar // row i's grammar; nil rows are unconstrained position int pending sampler.Result // in flight: sampled, not yet forwarded + scope *mlx.Scope // holds pending across steps // Steps run ahead asynchronously: when one faults, its token is already // forwarded and still has to be returned, so err waits for the next call. err error @@ -383,18 +389,15 @@ type pipelinedDecoder struct { func (r *Runner) pipelinedDecoder(spec *speculationSession, caches []cache.Cache, seed *mlx.Array, position int, layout []any, g *grammar) *pipelinedDecoder { t := &pipelinedDecoder{ r: r, spec: spec, caches: caches, layout: layout, position: position, - grammars: []*grammar{g}, + grammars: []*grammar{g}, scope: mlx.NewScope(), } logits := t.forward(seed) - mlx.Pin(logits) - defer mlx.Unpin(logits) if r.grammarEngine.hasGrammar(t.grammars) { // Dispatch the forward before the host builds the first masks. The // first sample commits nothing, so there is nothing to accept. A mask // fault here is a step fault like any other: the seed is already // forwarded, so the error waits for the first call. - mlx.Sweep() mlx.AsyncEval(logits) var errs []error logits, errs = r.grammarEngine.mask(t.grammars, logits, nil) @@ -421,12 +424,9 @@ func (t *pipelinedDecoder) next(int) ([]sampler.Result, error) { } out := t.pending logits := t.forward(out.Token.ExpandDims(-1)) - mlx.Pin(logits) - defer mlx.Unpin(logits) if t.r.grammarEngine.hasGrammar(t.grammars) { // Dispatch the forward before the host's grammar work. - mlx.Sweep() mlx.AsyncEval(logits) err := t.failRows(t.r.grammarEngine.accept(t.grammars, out.Token.Ints())) @@ -438,23 +438,25 @@ func (t *pipelinedDecoder) next(int) ([]sampler.Result, error) { t.pending = t.sample(logits) - mlx.Unpin(out.Arrays()...) + t.scope.Detach(out.Arrays()...) return []sampler.Result{out}, nil } // forward runs the model one step over token, shaped [B, L], and returns the // final position's [B, 1, V] logits, still lazy. func (t *pipelinedDecoder) forward(token *mlx.Array) *mlx.Array { - hidden, auxHidden := t.r.Model.Forward(&batch.Batch{ - InputIDs: token, - SeqOffsets: []int32{int32(t.position)}, - SeqQueryLens: []int32{int32(token.Dim(1))}, - Layout: t.layout, - }, t.caches) - t.spec.committed(token, auxHidden, t.position, nil) - t.position += token.Dim(1) - logits := t.r.Model.Unembed(hidden) - return logits.Slice(mlx.Slice(), mlx.Slice(logits.Dim(1)-1), mlx.Slice()) + return mlx.ScopedArrays(func() []*mlx.Array { + hidden, auxHidden := t.r.Model.Forward(&batch.Batch{ + InputIDs: token, + SeqOffsets: []int32{int32(t.position)}, + SeqQueryLens: []int32{int32(token.Dim(1))}, + Layout: t.layout, + }, t.caches) + t.spec.committed(token, auxHidden, t.position, nil) + t.position += token.Dim(1) + logits := t.r.Model.Unembed(hidden) + return []*mlx.Array{logits.Slice(mlx.Slice(), mlx.Slice(logits.Dim(1)-1), mlx.Slice())} + })[0] } // sample dispatches the batched sample over the decoder's rows. On an @@ -462,9 +464,8 @@ func (t *pipelinedDecoder) forward(token *mlx.Array) *mlx.Array { // is in flight before the previous tokens are synchronized. func (t *pipelinedDecoder) sample(logits *mlx.Array) sampler.Result { next := t.r.Sampler.Sample([]int{pipelineSlot}, logits.Squeeze(1)) - mlx.Pin(next.Arrays()...) - mlx.Sweep() mlx.AsyncEval(next.Arrays()...) + t.scope.Attach(next.Arrays()...) return next } @@ -477,6 +478,9 @@ func (t *pipelinedDecoder) drain() ([]sampler.Result, int, error) { // The sample leaves without its forward, so its accept runs here. err = t.failRows(t.r.grammarEngine.accept(t.grammars, t.pending.Token.Ints())) } + if err == nil { + t.scope.Detach(t.pending.Arrays()...) + } return []sampler.Result{t.pending}, t.position, err } @@ -484,7 +488,7 @@ func (t *pipelinedDecoder) close() { // The in-flight sample's forward was never dispatched; its report settles // the drafter level with the caches' resting offset. t.spec.settle(t.pending.Token) - mlx.Unpin(t.pending.Arrays()...) + t.scope.Close() } // detokenizer serializes sampled tokens into response chunks, holding bytes diff --git a/x/mlxrunner/prefix_cache.go b/x/mlxrunner/prefix_cache.go index 0f825dcd36d..9a75e5761d8 100644 --- a/x/mlxrunner/prefix_cache.go +++ b/x/mlxrunner/prefix_cache.go @@ -656,16 +656,18 @@ func (c *prefixCache) evictNode(node *trieNode) { func (c *prefixCache) dumpTree() { // Summary stats var cacheBytes int - for _, kv := range c.caches { - if kv == nil { - continue - } - for _, a := range kv.State() { - if a != nil { - cacheBytes += a.NumBytes() + mlx.Scoped(func() { + for _, kv := range c.caches { + if kv == nil { + continue + } + for _, a := range kv.State() { + if a != nil { + cacheBytes += a.NumBytes() + } } } - } + }) // Build active path set for marking. active := make(map[*trieNode]bool, len(c.activePath)) diff --git a/x/mlxrunner/runner.go b/x/mlxrunner/runner.go index 2db3d0f8bf6..626353c9c14 100644 --- a/x/mlxrunner/runner.go +++ b/x/mlxrunner/runner.go @@ -42,6 +42,7 @@ type Request struct { type Runner struct { Model base.Model + weights *mlx.Scope Tokenizer *tokenizer.Tokenizer Requests chan Request Sampler *sample.Sampler @@ -57,82 +58,95 @@ type Runner struct { } func (r *Runner) Load(modelName string) error { - root, err := model.Open(modelName) + weights, err := r.loadModel(modelName) if err != nil { return err } - defer root.Close() + mlx.Eval(weights...) + r.weights = mlx.NewScope() + r.weights.Attach(weights...) + configureWiredMemory() + return nil +} - m, err := base.New(root) - if err != nil { - return err - } +func (r *Runner) loadModel(modelName string) (weights []*mlx.Array, err error) { + weights = mlx.ScopedArrays(func() []*mlx.Array { + root, e := model.Open(modelName) + if e != nil { + err = e + return nil + } + defer root.Close() - // Load all tensor blobs from manifest - tensors, err := loadTensorsFromManifest(root) - if err != nil { - return err - } + m, e := base.New(root) + if e != nil { + err = e + return nil + } - // On Metal, materialize the loaded tensors with CPU reads before any - // weight graph exists, so the weight eval never commits a command buffer - // that waits on file data. CUDA loads read at dispatch and need no pre-pass. - if mlx.MetalIsAvailable() { - mlx.Eval(slices.Collect(maps.Values(tensors))...) - } + // Load all tensor blobs from manifest + tensors, e := loadTensorsFromManifest(root) + if e != nil { + err = e + return nil + } - // Assign weights to model (model-specific logic). Target and draft weights - // must be loaded before sweeping so tensors from a combined manifest are - // not discarded before the draft model can retain them. - if err := m.LoadWeights(tensors); err != nil { - return err - } + // On Metal, materialize the loaded tensors with CPU reads before any + // weight graph exists, so the weight eval never commits a command buffer + // that waits on file data. CUDA loads read at dispatch and need no pre-pass. + if mlx.MetalIsAvailable() { + mlx.Eval(slices.Collect(maps.Values(tensors))...) + } - var draftModel base.DraftModel - draft, err := base.NewDraft(root, m) - if err != nil { - return err - } - if draft != nil { - if err := draft.LoadWeights(tensors); err != nil { - return err + // Assign weights to model (model-specific logic). Target and draft weights + // must be loaded before the load scope ends so tensors from a combined + // manifest are not discarded before the draft model can retain them. + if err = m.LoadWeights(tensors); err != nil { + return nil } - draftModel = draft - } else if sd, ok := m.(base.SelfDraft); ok { - // Inline draft head: already loaded with the target; nil if none shipped. - draftModel = sd.SelfDraft() - } - collected := mlx.Collect(m) - if draft != nil { - draftArrays := mlx.Collect(draft) - collected = append(collected, draftArrays...) - if root.Draft != nil { - slog.Info("Loaded draft model", "tensor_prefix", root.Draft.TensorPrefix, "config", root.Draft.Config, "arrays", len(draftArrays)) - } else { - slog.Info("Loaded draft model", "arrays", len(draftArrays)) + var draftModel base.DraftModel + draft, e := base.NewDraft(root, m) + if e != nil { + err = e + return nil + } + if draft != nil { + if err = draft.LoadWeights(tensors); err != nil { + return nil + } + draftModel = draft + } else if sd, ok := m.(base.SelfDraft); ok { + // Inline draft head: already loaded with the target; nil if none shipped. + draftModel = sd.SelfDraft() } - } - for _, arr := range collected { - mlx.Pin(arr) - } - mlx.Sweep() - mlx.Eval(collected...) - configureWiredMemory() - r.Model = m - r.Tokenizer = m.Tokenizer() - r.contextLength = m.MaxContextLength() - caches := m.NewCaches() - draftCaches := newDraftCaches(draftModel) - r.cache = newPrefixCache(slices.Concat(caches, draftCaches)) - r.Sampler = sample.New(r.contextLength) - r.spec = newSpeculation(r, draftModel, caches, draftCaches) - r.grammarEngine = newGrammarEngine(logitsWidth(m), r.Tokenizer) + w := mlx.Collect(m) + if draft != nil { + draftArrays := mlx.Collect(draft) + w = append(w, draftArrays...) + if root.Draft != nil { + slog.Info("Loaded draft model", "tensor_prefix", root.Draft.TensorPrefix, "config", root.Draft.Config, "arrays", len(draftArrays)) + } else { + slog.Info("Loaded draft model", "arrays", len(draftArrays)) + } + } - mlx.EnableCompile() + r.Model = m + r.Tokenizer = m.Tokenizer() + r.contextLength = m.MaxContextLength() + caches := m.NewCaches() + draftCaches := newDraftCaches(draftModel) + r.cache = newPrefixCache(slices.Concat(caches, draftCaches)) + r.Sampler = sample.New(r.contextLength) + r.spec = newSpeculation(r, draftModel, caches, draftCaches) + r.grammarEngine = newGrammarEngine(logitsWidth(m), r.Tokenizer) - return nil + mlx.EnableCompile() + + return w + }) + return weights, err } func (r *Runner) Close() { @@ -140,6 +154,8 @@ func (r *Runner) Close() { r.grammarEngine.close() r.grammarEngine = nil } + r.weights.Close() + r.weights = nil } // newDraftCaches returns nil when the model ships no draft. @@ -152,24 +168,23 @@ func newDraftCaches(draft base.DraftModel) []cache.Cache { // logitsWidth reads a model's logits width off a one-token forward's static // shape — the same Forward and Unembed path decode logits take. Nothing is -// evaluated, and the probe's caches and graph are released before returning, -// which sweeps every unpinned array: call this only at load, after the -// model's weights are pinned. -func logitsWidth(m base.Model) int { - caches := m.NewCaches() - hidden, _ := m.Forward(&batch.Batch{ - InputIDs: mlx.FromValues([]int32{0}, 1, 1), - SeqOffsets: []int32{0}, - SeqQueryLens: []int32{1}, - }, caches) - logits := m.Unembed(hidden) - width := logits.Dim(logits.NumDims() - 1) - for _, c := range caches { - if c != nil { - c.Free() +// evaluated. +func logitsWidth(m base.Model) (width int) { + mlx.Scoped(func() { + caches := m.NewCaches() + hidden, _ := m.Forward(&batch.Batch{ + InputIDs: mlx.FromValues([]int32{0}, 1, 1), + SeqOffsets: []int32{0}, + SeqQueryLens: []int32{1}, + }, caches) + logits := m.Unembed(hidden) + width = logits.Dim(logits.NumDims() - 1) + for _, c := range caches { + if c != nil { + c.Free() + } } - } - mlx.Sweep() + }) return width } diff --git a/x/mlxrunner/sample/logprob_test.go b/x/mlxrunner/sample/logprob_test.go index 2afb6b0f238..97abf343be2 100644 --- a/x/mlxrunner/sample/logprob_test.go +++ b/x/mlxrunner/sample/logprob_test.go @@ -26,16 +26,15 @@ func runSampleLogprobs(t *mlxtest.T, logits []float32, topK int) (int32, float64 s := New(128) t.Cleanup(func() { s.Free() - mlx.Sweep() }) s.Add(0, Options{Logprobs: true, TopLogprobs: topK}, nil) tensor := mlx.FromValues(logits, 1, len(logits)) - res := s.Sample([]int{0}, tensor) - - mlx.Pin(res.Arrays()...) - t.Cleanup(func() { mlx.Unpin(res.Arrays()...) }) - mlx.Sweep() + var res Result + mlx.ScopedArrays(func() []*mlx.Array { + res = s.Sample([]int{0}, tensor) + return res.Arrays() + }) mlx.Eval(res.Arrays()...) selected := res.Token.Int() @@ -249,15 +248,12 @@ func TestBatchedLogprobsPerRow(t *testing.T) { s := New(128) t.Cleanup(func() { s.Free() - mlx.Sweep() }) s.Add(1, Options{Logprobs: true}, nil) s.Add(2, Options{Logprobs: true}, nil) logits := mlx.FromValues(append(append([]float32{}, rowA...), rowB...), 2, 3) res := s.Sample([]int{1, 2}, logits) - mlx.Pin(res.Arrays()...) - t.Cleanup(func() { mlx.Unpin(res.Arrays()...) }) mlx.Eval(res.Arrays()...) got := res.Logprob.Floats() diff --git a/x/mlxrunner/sample/sample.go b/x/mlxrunner/sample/sample.go index 3b4c2d3a095..4e85778b475 100644 --- a/x/mlxrunner/sample/sample.go +++ b/x/mlxrunner/sample/sample.go @@ -38,8 +38,8 @@ type Result struct { } // Arrays returns the tensor fields as a slice so callers can drive the mlx -// lifecycle verbs (Pin, Unpin, Eval, AsyncEval) over the whole group. Unset -// fields stay nil; the mlx helpers skip them. +// lifecycle verbs (Eval, AsyncEval, a held scope's Attach) over the whole +// group. Unset fields stay nil; the mlx helpers skip them. func (r Result) Arrays() []*mlx.Array { return []*mlx.Array{r.Token, r.Logprob, r.TopTokens, r.TopLogprobs} } @@ -174,6 +174,7 @@ type Sampler struct { // belongs to slots[i]; W is max(RepeatLastN) across penalty slots. // Allocated on the first penalty slot, rebuilt only in Add/Remove. history *mlx.Array + scope *mlx.Scope // allSameOpts: every registered slot shares Options. When true the // canonical shared value is s.slots[0].opts. @@ -209,6 +210,7 @@ func New(numCtx int) *Sampler { byID: make(map[int]*slotState), allSameOpts: true, numCtx: numCtx, + scope: mlx.NewScope(), } } @@ -270,29 +272,32 @@ func (s *Sampler) Add(seqID int, opts Options, priorTokens []int32) { // Grow the pool to hold this slot's row. The pool is lazy — the first // penalty slot allocates it — and thereafter every registered slot // gets a row (rows for non-penalty slots are zero and never read). - // Invariant: s.history is pinned whenever non-nil. if s.history != nil || opts.usesHistory() { - targetWidth := max(opts.RepeatLastN, s.historyWidth()) - newRow := makeHistoryRow(priorTokens, opts.RepeatLastN, targetWidth) - - var pool *mlx.Array - switch { - case s.history == nil && len(s.slots) == 0: - pool = newRow - case s.history == nil: - // First penalty slot with non-penalty slots already registered; - // seed zero rows so s.slots and pool row indices stay aligned. - zeros := mlx.Zeros(mlx.DTypeInt32, len(s.slots), targetWidth) - pool = zeros.Concatenate(0, newRow) - case targetWidth > s.historyWidth(): - pad := mlx.Zeros(mlx.DTypeInt32, s.history.Dim(0), targetWidth-s.historyWidth()) - pool = s.history.Concatenate(1, pad).Concatenate(0, newRow) - default: - pool = s.history.Concatenate(0, newRow) - } + pool := mlx.ScopedArrays(func() []*mlx.Array { + targetWidth := max(opts.RepeatLastN, s.historyWidth()) + newRow := makeHistoryRow(priorTokens, opts.RepeatLastN, targetWidth) + + var pool *mlx.Array + switch { + case s.history == nil && len(s.slots) == 0: + pool = newRow + case s.history == nil: + // First penalty slot with non-penalty slots already registered; + // seed zero rows so s.slots and pool row indices stay aligned. + zeros := mlx.Zeros(mlx.DTypeInt32, len(s.slots), targetWidth) + pool = zeros.Concatenate(0, newRow) + case targetWidth > s.historyWidth(): + pad := mlx.Zeros(mlx.DTypeInt32, s.history.Dim(0), targetWidth-s.historyWidth()) + pool = s.history.Concatenate(1, pad).Concatenate(0, newRow) + default: + pool = s.history.Concatenate(0, newRow) + } - mlx.Pin(pool) - mlx.Unpin(s.history) + // The concatenation still reads the old pool. + s.scope.Discard(s.history) + return []*mlx.Array{pool} + })[0] + s.scope.Attach(pool) s.history = pool if opts.usesHistory() { @@ -363,34 +368,38 @@ func (s *Sampler) Remove(seqID int) { return } - n := s.history.Dim(0) - var newHistory *mlx.Array - switch { - case n == 1: - newHistory = nil - case row == 0: - newHistory = s.history.Slice(mlx.Slice(1, n), mlx.Slice()) - case row == n-1: - newHistory = s.history.Slice(mlx.Slice(0, row), mlx.Slice()) - default: - before := s.history.Slice(mlx.Slice(0, row), mlx.Slice()) - after := s.history.Slice(mlx.Slice(row+1, n), mlx.Slice()) - newHistory = before.Concatenate(0, after) - } + newHistory := mlx.ScopedArrays(func() []*mlx.Array { + n := s.history.Dim(0) + var newHistory *mlx.Array + switch { + case n == 1: + newHistory = nil + case row == 0: + newHistory = s.history.Slice(mlx.Slice(1, n), mlx.Slice()) + case row == n-1: + newHistory = s.history.Slice(mlx.Slice(0, row), mlx.Slice()) + default: + before := s.history.Slice(mlx.Slice(0, row), mlx.Slice()) + after := s.history.Slice(mlx.Slice(row+1, n), mlx.Slice()) + newHistory = before.Concatenate(0, after) + } - mlx.Pin(newHistory) - mlx.Unpin(s.history) + s.scope.Discard(s.history) + return []*mlx.Array{newHistory} + })[0] + s.scope.Attach(newHistory) s.history = newHistory } // Free releases the pooled history tensor and resets the sampler to the // New-equivalent state so it may be reused. func (s *Sampler) Free() { - mlx.Unpin(s.history) + s.scope.Close() *s = Sampler{ byID: make(map[int]*slotState), allSameOpts: true, numCtx: s.numCtx, + scope: mlx.NewScope(), } } @@ -411,34 +420,38 @@ func (s *Sampler) Sample(seqIDs []int, logits *mlx.Array) Result { slots[i] = slot } - var token *mlx.Array - if opts0, ok := s.canBatch(slots); ok { - token = s.sampleTokensUniform(slots, opts0, logits) - } else { - token = s.sampleTokensSerial(slots, logits) - } - - res := Result{Token: token} - if s.anyLogprobs { - // Log-softmax over original logits so every row holds a truthful - // value (compute-for-all; consumers filter per-slot). Subtract - // max first for numerical stability in the logsumexp. - lp := logits.AsType(mlx.DTypeFloat32) - lp = lp.Subtract(lp.MaxAxis(-1, true)) - lp = lp.Subtract(lp.LogsumexpAxis(-1, true)) - res.Logprob = lp.TakeAlongAxis(token.ExpandDims(-1), -1) - if s.maxTopLogprobs > 0 { - k := s.maxTopLogprobs - if vocab := lp.Dim(lp.NumDims() - 1); k > vocab { - k = vocab + var res Result + mlx.ScopedArrays(func() []*mlx.Array { + var token *mlx.Array + if opts0, ok := s.canBatch(slots); ok { + token = s.sampleTokensUniform(slots, opts0, logits) + } else { + token = s.sampleTokensSerial(slots, logits) + } + + res = Result{Token: token} + if s.anyLogprobs { + // Log-softmax over original logits so every row holds a truthful + // value (compute-for-all; consumers filter per-slot). Subtract + // max first for numerical stability in the logsumexp. + lp := logits.AsType(mlx.DTypeFloat32) + lp = lp.Subtract(lp.MaxAxis(-1, true)) + lp = lp.Subtract(lp.LogsumexpAxis(-1, true)) + res.Logprob = lp.TakeAlongAxis(token.ExpandDims(-1), -1) + if s.maxTopLogprobs > 0 { + k := s.maxTopLogprobs + if vocab := lp.Dim(lp.NumDims() - 1); k > vocab { + k = vocab + } + // Argpartition on the negated values places the K largest + // (unsorted) in positions [0:K]. + idx := lp.Negative().ArgpartitionAxis(k-1, -1).Slice(mlx.Slice(), mlx.Slice(0, k)) + res.TopTokens = idx.AsType(mlx.DTypeInt32) + res.TopLogprobs = lp.TakeAlongAxis(idx, -1) } - // Argpartition on the negated values places the K largest - // (unsorted) in positions [0:K]. - idx := lp.Negative().ArgpartitionAxis(k-1, -1).Slice(mlx.Slice(), mlx.Slice(0, k)) - res.TopTokens = idx.AsType(mlx.DTypeInt32) - res.TopLogprobs = lp.TakeAlongAxis(idx, -1) } - } + return res.Arrays() + }) return res } diff --git a/x/mlxrunner/sample/sample_test.go b/x/mlxrunner/sample/sample_test.go index eb9446250b2..30edfca805c 100644 --- a/x/mlxrunner/sample/sample_test.go +++ b/x/mlxrunner/sample/sample_test.go @@ -36,7 +36,6 @@ func sampleOne(t *mlxtest.T, opts Options, priorTokens []int32, values []float32 s := New(128) t.Cleanup(func() { s.Free() - mlx.Sweep() }) s.Add(0, opts, priorTokens) @@ -139,7 +138,6 @@ func TestDistributionAppliesTopKBeforeTopP(t *testing.T) { s := New(128) t.Cleanup(func() { s.Free() - mlx.Sweep() }) s.Add(0, Options{Temperature: 1, TopK: 2, TopP: 0.7}, nil) @@ -210,7 +208,6 @@ func TestSeededSamplingIsReproducible(t *testing.T) { s := New(128) t.Cleanup(func() { s.Free() - mlx.Sweep() }) s.Add(0, Options{Temperature: 1, TopK: 4, Seed: seed, UseSeed: true}, nil) @@ -243,7 +240,6 @@ func TestSeededBernoulliIsReproducible(t *testing.T) { s := New(128) t.Cleanup(func() { s.Free() - mlx.Sweep() }) s.Add(0, Options{Seed: 99, UseSeed: true}, nil) @@ -269,7 +265,6 @@ func TestSampleHistoryWindow(t *testing.T) { s := New(128) t.Cleanup(func() { s.Free() - mlx.Sweep() }) // RepeatLastN=2 with priors {1, 2, 3}: makeHistoryRow keeps only @@ -300,7 +295,6 @@ func TestSpeculativeScoresUsesDraftHistoryWithoutCommit(t *testing.T) { s := New(128) t.Cleanup(func() { s.Free() - mlx.Sweep() }) s.Add(0, Options{RepeatLastN: 2, RepeatPenalty: 10}, []int32{1, 2}) @@ -333,7 +327,6 @@ func TestDistributionSingleRowAppliesDraftPrefix(t *testing.T) { s := New(128) t.Cleanup(func() { s.Free() - mlx.Sweep() }) // A proposal step passes one logits row with the chain's earlier drafts: @@ -361,7 +354,6 @@ func TestDistributionMultiRowWithoutChain(t *testing.T) { s := New(128) t.Cleanup(func() { s.Free() - mlx.Sweep() }) // A block drafter's proposal batch samples every row from one call with @@ -394,7 +386,6 @@ func TestCommitBatchesRingWrites(t *testing.T) { s := New(128) t.Cleanup(func() { s.Free() - mlx.Sweep() }) s.Add(0, Options{RepeatLastN: 4, RepeatPenalty: 1.1}, []int32{10, 11, 12}) @@ -492,7 +483,6 @@ func TestBatchSamplingPreservesPerSlotBehavior(t *testing.T) { s := New(128) t.Cleanup(func() { s.Free() - mlx.Sweep() }) for _, spec := range tc.slots { s.Add(spec.id, spec.opts, spec.priors) @@ -519,7 +509,6 @@ func TestRemoveDoesNotLeakHistory(t *testing.T) { s := New(128) t.Cleanup(func() { s.Free() - mlx.Sweep() }) s.Add(1, opts, []int32{1}) s.Add(2, opts, []int32{2}) diff --git a/x/mlxrunner/server.go b/x/mlxrunner/server.go index 7178e486057..0d3b6b1dd4e 100644 --- a/x/mlxrunner/server.go +++ b/x/mlxrunner/server.go @@ -53,10 +53,7 @@ func Execute(args []string) error { if err != nil { return err } - defer worker.Stop(context.Background(), func() { - mlx.Sweep() - mlx.ClearCache() - }) + defer worker.Stop(context.Background(), mlx.ClearCache) runnerCtx, cancelRunner := context.WithCancel(context.Background()) defer cancelRunner() diff --git a/x/mlxrunner/speculate.go b/x/mlxrunner/speculate.go index 4b461686f12..2aa2ca4917e 100644 --- a/x/mlxrunner/speculate.go +++ b/x/mlxrunner/speculate.go @@ -24,7 +24,7 @@ type draftSession interface { // hidden state at that slot. Runs arrive in slot order — prefill // chunks, the decode seed, then each round's validated tokens. media is // the run's manifest (feature-bearing for items the run overlaps), valid - // only for the call; a session that defers its forward pins what it + // only for the call; a session that defers its forward holds what it // keeps. Nil outside prefill. committed(tokens, hiddens *mlx.Array, position int, media []batch.MediaItem) @@ -201,6 +201,7 @@ type speculativeDecoder struct { current sampler.Result // emitted (or the seed), not yet forwarded inner *pipelinedDecoder // pipelines plain tokens while parked; nil while drafting grammar *grammar + scope *mlx.Scope // holds current across rounds } // decoder returns the decoder for this engine's session. A speculationSession that @@ -208,8 +209,9 @@ type speculativeDecoder struct { // running the inner pipelined decoder whose reports keep the draft KV level. func (s *speculationSession) decoder(seed *mlx.Array, position int, grammar *grammar) decoder { current := sampler.Result{Token: seed} - mlx.Pin(current.Arrays()...) - return &speculativeDecoder{s: s, position: position, current: current, grammar: grammar} + scope := mlx.NewScope() + scope.Attach(current.Arrays()...) + return &speculativeDecoder{s: s, position: position, current: current, grammar: grammar, scope: scope} } func (st *speculativeDecoder) next(remaining int) ([]sampler.Result, error) { @@ -231,16 +233,13 @@ func (st *speculativeDecoder) next(remaining int) ([]sampler.Result, error) { // land that extra token within it rather than overshooting. At // remaining 1 the cap is 0 and the last token decodes plainly. candidates = s.drafter.propose(st.current.Token, min(s.limit, remaining-1)) + mlx.AsyncEval(candidates.Arrays()...) } var accepted, observed int var err error if candidates == nil { results, err = st.park(remaining) } else { - // candidates stays pinned across accept's internal sweep and the - // draft-count read below; accept pins only its own intermediates. - mlx.Pin(candidates.tokens) - defer mlx.Unpin(candidates.tokens) results, accepted, observed, err = st.s.accept(&st.position, st.current, candidates, st.grammar) } if err != nil { @@ -257,11 +256,11 @@ func (st *speculativeDecoder) next(remaining int) ([]sampler.Result, error) { return results, nil } -// advance retires the last returned token as the next call's current, pinned -// across the sweeps the next call runs before reading it. Nothing is forced here. +// advance retires the last returned token as the next call's current, held +// until the next call reads it. Nothing is forced here. func (st *speculativeDecoder) advance(next sampler.Result) { - mlx.Pin(next.Arrays()...) - mlx.Unpin(st.current.Arrays()...) + st.scope.Attach(next.Arrays()...) + st.scope.Discard(st.current.Arrays()...) st.current = next } @@ -314,7 +313,7 @@ func (st *speculativeDecoder) close() { // the drafter level with the caches' resting offset. st.s.settle(st.current.Token) } - mlx.Unpin(st.current.Arrays()...) + st.scope.Close() st.s.logStats() } @@ -388,10 +387,6 @@ func commitSpeculation(caches []cache.Cache, accepted, draftCount, before int) { // acceptance model learns from, capped at the EOS (a terminator, not a target // rejection). NumPredict is the decode loop's to enforce, so a token past the // budget is left for decode to drop, not cut here. -// -// The caller keeps current and the candidate tokens pinned across the call, -// since accept sweeps before its eval and reads both afterward; accept pins -// only the intermediates it produces. func (s *speculationSession) accept(position *int, current sampler.Result, candidates *draftCandidates, g *grammar) (results []sampler.Result, accepted, observed int, err error) { r := s.spec.r before := *position @@ -412,54 +407,51 @@ func (s *speculationSession) accept(position *int, current sampler.Result, candi } defer commit(0) - dist := candidates.dist.Arrays() - mlx.Pin(dist...) - mlx.Sweep() - mlx.AsyncEval(candidates.tokens) - mlx.Unpin(dist...) - - hiddenSeq, auxHiddenSeq := r.Model.Forward(&batch.Batch{ - InputIDs: current.Token.ExpandDims(-1).Concatenate(1, candidates.tokens), - SeqOffsets: []int32{int32(before)}, - SeqQueryLens: []int32{int32(draftCount + 1)}, - Layout: s.layout, - }, s.spec.targets) - - // Row i of the fused hidden is the state after the token at before+i, so - // the rows already line up with the drafts: row 0 (current's state) - // predicts draft 0, and the row after the last accepted draft is the - // bonus row. No separate base-logits forward exists on this path. - logits := r.Model.Unembed(hiddenSeq) - - draftIDs := candidates.tokens.Ints() - constrained := g.constraining() - if constrained { - var errs []error - logits, errs = r.grammarEngine.mask([]*grammar{g}, logits, [][]int32{draftIDs}) - if err := errors.Join(errs...); err != nil { - return nil, 0, 0, err + var auxHiddenSeq, acceptedMask, residualTokens, bonusToken *mlx.Array + var draftIDs []int32 + var constrained bool + var maskErr error + mlx.ScopedEval(func() []*mlx.Array { + var hiddenSeq *mlx.Array + hiddenSeq, auxHiddenSeq = r.Model.Forward(&batch.Batch{ + InputIDs: current.Token.ExpandDims(-1).Concatenate(1, candidates.tokens), + SeqOffsets: []int32{int32(before)}, + SeqQueryLens: []int32{int32(draftCount + 1)}, + Layout: s.layout, + }, s.spec.targets) + + // Row i of the fused hidden is the state after the token at before+i, so + // the rows already line up with the drafts: row 0 (current's state) + // predicts draft 0, and the row after the last accepted draft is the + // bonus row. No separate base-logits forward exists on this path. + logits := r.Model.Unembed(hiddenSeq) + + draftIDs = candidates.tokens.Ints() + constrained = g.constraining() + if constrained { + var errs []error + logits, errs = r.grammarEngine.mask([]*grammar{g}, logits, [][]int32{draftIDs}) + if maskErr = errors.Join(errs...); maskErr != nil { + return nil + } } - } - targetDist := r.Sampler.Distribution(pipelineSlot, logits, candidates.tokens) - draftDist := candidates.dist - acceptedMask := r.sampleAcceptedMask(targetDist.SliceRows(0, draftCount), draftDist, candidates.tokens) - - // The next token is sampled for every possible outcome before anything - // is evaluated — the residual at each rejection point in one batched - // draw, plus the bonus row — so a single Eval covers acceptance and the - // next token instead of a second host round trip after the rejection - // point is known. - residualTokens := r.Sampler.SampleDistribution(pipelineSlot, targetDist.SliceRows(0, draftCount).ResidualAgainst(draftDist)) - bonusToken := r.sampleTokenAt(targetDist, draftCount) - - // Pin the arrays read after the eval; current and the candidate tokens - // stay pinned by the caller across the call. - live := []*mlx.Array{hiddenSeq, auxHiddenSeq, acceptedMask, residualTokens, bonusToken} - mlx.Pin(live...) - defer mlx.Unpin(live...) - mlx.Sweep() - mlx.Eval(candidates.tokens, acceptedMask, residualTokens, bonusToken) + targetDist := r.Sampler.Distribution(pipelineSlot, logits, candidates.tokens) + draftDist := candidates.dist + acceptedMask = r.sampleAcceptedMask(targetDist.SliceRows(0, draftCount), draftDist, candidates.tokens) + + // The next token is sampled for every possible outcome before anything + // is evaluated — the residual at each rejection point in one batched + // draw, plus the bonus row — so a single Eval covers acceptance and the + // next token instead of a second host round trip after the rejection + // point is known. + residualTokens = r.Sampler.SampleDistribution(pipelineSlot, targetDist.SliceRows(0, draftCount).ResidualAgainst(draftDist)) + bonusToken = r.sampleTokenAt(targetDist, draftCount) + return []*mlx.Array{auxHiddenSeq, acceptedMask, residualTokens, bonusToken} + }) + if maskErr != nil { + return nil, 0, 0, maskErr + } acceptedFlags := acceptedMask.Ints() for _, ok := range acceptedFlags { diff --git a/x/models/qwen3_5/vision.go b/x/models/qwen3_5/vision.go index fe1b45a1e61..df664b4df09 100644 --- a/x/models/qwen3_5/vision.go +++ b/x/models/qwen3_5/vision.go @@ -97,7 +97,7 @@ type visionLayout struct { // The publisher uses the same tensor layout, image preprocessing, and MRoPE // layout for both families. type VisionAdapter struct { - // Model is exported so mlx.Collect traverses and pins every tower weight. + // Model is exported so mlx.Collect reaches every tower weight. // An unexported wrapper field is invisible to the reflection collector. Model *Model } diff --git a/x/models/qwen4_exp/engram_cache.go b/x/models/qwen4_exp/engram_cache.go index b836a46d3e6..1cfce5d983c 100644 --- a/x/models/qwen4_exp/engram_cache.go +++ b/x/models/qwen4_exp/engram_cache.go @@ -14,6 +14,7 @@ import ( type engramCache struct { history *mlx.Array convHistory *mlx.Array + scope *mlx.Scope offset int eosID int64 width int @@ -27,24 +28,25 @@ type engramCache struct { type engramSnapshot struct { history *mlx.Array convHistory *mlx.Array + scope *mlx.Scope offset int } func newEngramCache(width, convTail, convDim int, eosID int64) *engramCache { - return &engramCache{width: width, convTail: convTail, convDim: convDim, eosID: eosID} + return &engramCache{width: width, convTail: convTail, convDim: convDim, eosID: eosID, scope: mlx.NewScope()} } func (c *engramCache) setHistory(value *mlx.Array) { value = value.Clone() - mlx.Pin(value) - mlx.Unpin(c.history) + c.scope.Attach(value) + c.scope.Discard(c.history) c.history = value } func (c *engramCache) setConvHistory(value *mlx.Array) { value = value.Clone() - mlx.Pin(value) - mlx.Unpin(c.convHistory) + c.scope.Attach(value) + c.scope.Discard(c.convHistory) c.convHistory = value } @@ -126,7 +128,7 @@ func (c *engramCache) State() []*mlx.Array { } func (c *engramCache) Free() { - mlx.Unpin(c.history, c.convHistory) + c.scope.Close() c.history = nil c.convHistory = nil c.offset = 0 @@ -176,8 +178,10 @@ func (c *engramCache) Restore(snapshot cache.Snapshot, target int) bool { if !ok || value.offset != target { return false } - c.setHistory(value.history) - c.setConvHistory(value.convHistory) + mlx.Scoped(func() { + c.setHistory(value.history) + c.setConvHistory(value.convHistory) + }) c.offset = target return true } @@ -194,12 +198,12 @@ func (c *engramCache) Split(snapshot cache.Snapshot, _ int) (cache.Snapshot, cac } func newEngramSnapshot(history, convHistory *mlx.Array, offset int) *engramSnapshot { - snapshot := &engramSnapshot{history: history.Clone(), convHistory: convHistory.Clone(), offset: offset} - mlx.Pin(snapshot.history, snapshot.convHistory) + snapshot := &engramSnapshot{history: history.Clone(), convHistory: convHistory.Clone(), scope: mlx.NewScope(), offset: offset} + snapshot.scope.Attach(snapshot.history, snapshot.convHistory) mlx.AsyncEval(snapshot.history, snapshot.convHistory) return snapshot } func (s *engramSnapshot) Size() int { return s.history.NumBytes() + s.convHistory.NumBytes() } func (s *engramSnapshot) SetMaterializeHook(func(int)) {} -func (s *engramSnapshot) Close() { mlx.Unpin(s.history, s.convHistory) } +func (s *engramSnapshot) Close() { s.scope.Close() } From aeb8f711670a327705481ead01239258b8e1df3f Mon Sep 17 00:00:00 2001 From: Jesse Gross Date: Fri, 4 Sep 2026 10:01:19 -0700 Subject: [PATCH 12/24] mlx: drop the empty-handle checks outside the bindings Models and layers checked optional weights for nil and also for a handle that no longer refers to an array, and evaluation and weight collection skipped such handles. No path produces one: a missing tensor is nil, and a handle only loses its array when its scope frees it, after which using it is a bug. The nil checks stay; the validity check is internal to the bindings now. --- x/mlxrunner/mlx/array.go | 6 ++++-- x/mlxrunner/mlx/mlx.go | 2 +- x/mlxrunner/mlx/nn.go | 2 +- x/mlxrunner/mlx/ops_extra.go | 6 +++--- x/mlxrunner/mlx/scope.go | 2 +- x/mlxrunner/mlx/scope_test.go | 22 +++++++++++----------- x/models/cohere2_moe/cohere2_moe.go | 2 +- x/models/gemma4/gemma4.go | 2 +- x/models/laguna/laguna.go | 8 ++++---- x/models/nemotron_h/nemotron_h.go | 2 +- x/models/nn/nn.go | 10 +++++----- x/models/qwen3_5/qwen3_5.go | 4 ++-- 12 files changed, 35 insertions(+), 33 deletions(-) diff --git a/x/mlxrunner/mlx/array.go b/x/mlxrunner/mlx/array.go index 675be0c86e8..c1871ad6c27 100644 --- a/x/mlxrunner/mlx/array.go +++ b/x/mlxrunner/mlx/array.go @@ -119,7 +119,9 @@ func (t *Array) Clone() *Array { // misc. utilities -func (t *Array) Valid() bool { +// valid reports whether t still refers to an array: false once its scope +// freed it. +func (t *Array) valid() bool { return t.ctx.ctx != nil } @@ -139,7 +141,7 @@ func (t *Array) LogValue() slog.Value { attrs := []slog.Attr{ slog.String("name", t.name), } - if t.Valid() { + if t.valid() { attrs = append(attrs, slog.Any("dtype", t.DType()), slog.Any("shape", t.Dims()), diff --git a/x/mlxrunner/mlx/mlx.go b/x/mlxrunner/mlx/mlx.go index 9c34963ea73..141ba9cc220 100644 --- a/x/mlxrunner/mlx/mlx.go +++ b/x/mlxrunner/mlx/mlx.go @@ -133,7 +133,7 @@ func doEval(outputs []*Array, async bool) { defer freeVectorArray(vector) for _, output := range outputs { - if output != nil && output.Valid() { + if output != nil { mlxCheck(C.mlx_vector_array_append_value(vector, output.ctx)) } } diff --git a/x/mlxrunner/mlx/nn.go b/x/mlxrunner/mlx/nn.go index d2e7fb4f17c..a953f97b459 100644 --- a/x/mlxrunner/mlx/nn.go +++ b/x/mlxrunner/mlx/nn.go @@ -8,7 +8,7 @@ type Linear struct { // Forward computes the linear transformation: x @ Weight.T + Bias func (m *Linear) Forward(x *Array) *Array { w := m.Weight.Transpose(1, 0) - if m.Bias.Valid() { + if m.Bias != nil { return m.Bias.Addmm(x, w, 1.0, 1.0) } diff --git a/x/mlxrunner/mlx/ops_extra.go b/x/mlxrunner/mlx/ops_extra.go index d97c27a3bb9..125dbb52868 100644 --- a/x/mlxrunner/mlx/ops_extra.go +++ b/x/mlxrunner/mlx/ops_extra.go @@ -153,7 +153,7 @@ func Conv1d(x, weight *Array, bias *Array, stride, padding, dilation, groups int C.int(groups), DefaultStream().ctx, )) - if bias != nil && bias.Valid() { + if bias != nil { out = Add(out, bias) } return out @@ -646,7 +646,7 @@ func collect(v reflect.Value, arrays *[]*Array, seen map[uintptr]bool) { seen[ptr] = true if arr, ok := v.Interface().(*Array); ok { - if arr != nil && arr.Valid() { + if arr != nil { *arrays = append(*arrays, arr) } return @@ -659,7 +659,7 @@ func collect(v reflect.Value, arrays *[]*Array, seen map[uintptr]bool) { case reflect.Struct: // Check if this struct IS an Array (not a pointer to one) if arr, ok := v.Addr().Interface().(*Array); ok { - if arr != nil && arr.Valid() { + if arr != nil { *arrays = append(*arrays, arr) } return diff --git a/x/mlxrunner/mlx/scope.go b/x/mlxrunner/mlx/scope.go index 92156246a5a..ec63d1b743b 100644 --- a/x/mlxrunner/mlx/scope.go +++ b/x/mlxrunner/mlx/scope.go @@ -177,7 +177,7 @@ func (s *Scope) take(t *Array) { if from == s { return } - if !t.Valid() { + if !t.valid() { panic(fmt.Sprintf("mlx: array %q used after its scope ended", t.name)) } if from.noEscape { diff --git a/x/mlxrunner/mlx/scope_test.go b/x/mlxrunner/mlx/scope_test.go index 05994d78653..641e6ddc7dd 100644 --- a/x/mlxrunner/mlx/scope_test.go +++ b/x/mlxrunner/mlx/scope_test.go @@ -22,17 +22,17 @@ func TestScopeFreesWhatIsNotReturned(t *testing.T) { return []*Array{FromValue(3), nil, kept} }) returned = out[0] - if !returned.Valid() { + if !returned.valid() { t.Fatal("returned array was freed with the scope that created it") } - if dropped.Valid() { + if dropped.valid() { t.Fatal("array not returned survived its scope") } }) - if returned.Valid() { + if returned.valid() { t.Fatal("returned array survived the scope it was returned into") } - if !kept.Valid() { + if !kept.valid() { t.Fatal("returning a held array moved it out of its scope") } }) @@ -48,10 +48,10 @@ func TestScopedEvalEvaluatesAfterBuild(t *testing.T) { tmp = FromValue(2) return []*Array{FromValue(1).Add(tmp)} }) - if tmp.Valid() { + if tmp.valid() { t.Fatal("intermediate survived the build scope") } - if !out[0].Valid() || out[0].Int() != 3 { + if !out[0].valid() || out[0].Int() != 3 { t.Fatal("returned array was not evaluated after the build scope") } }) @@ -72,11 +72,11 @@ func TestHeldScope(t *testing.T) { kept, discarded, detached = FromValue(1), FromValue(2), FromValue(3) held.Attach(kept, discarded, detached) held.Discard(discarded) - if discarded.Valid() { + if discarded.valid() { t.Fatal("discarded array survived") } Scoped(func() { held.Detach(detached) }) - if detached.Valid() { + if detached.valid() { t.Fatal("detached array survived the scope it was detached into") } if !panics(func() { other.Discard(kept) }) { @@ -92,11 +92,11 @@ func TestHeldScope(t *testing.T) { t.Fatal("no panic holding an array twice") } }) - if !kept.Valid() { + if !kept.valid() { t.Fatal("held array was freed with the scope that created it") } held.Close() - if kept.Valid() { + if kept.valid() { t.Fatal("held array survived its scope's close") } }) @@ -115,7 +115,7 @@ func TestScopeEndsOnPanic(t *testing.T) { panic("build failed") }) }() - if a.Valid() { + if a.valid() { t.Fatal("array survived the scope that panicked") } if currentScope != start { diff --git a/x/models/cohere2_moe/cohere2_moe.go b/x/models/cohere2_moe/cohere2_moe.go index 64f5d271f41..d29a20099d8 100644 --- a/x/models/cohere2_moe/cohere2_moe.go +++ b/x/models/cohere2_moe/cohere2_moe.go @@ -398,7 +398,7 @@ func supportsGatherQMM(mode string, bits int) bool { // weights to the [E, in, out] layout GatherMM consumes, materialized once at // load so the forward path avoids per-call transposes. func transposeExpertWeightForGatherMM(w *mlx.Array) *mlx.Array { - if w == nil || !w.Valid() || w.NumDims() != 3 { + if w == nil || w.NumDims() != 3 { return w } return mlx.Transpose(w, 0, 2, 1).Clone() diff --git a/x/models/gemma4/gemma4.go b/x/models/gemma4/gemma4.go index f9b0c6685ab..e402807ddfc 100644 --- a/x/models/gemma4/gemma4.go +++ b/x/models/gemma4/gemma4.go @@ -186,7 +186,7 @@ func sliceAxis1(a *mlx.Array, start, stop int32) *mlx.Array { // transposeForGatherMM transposes stacked expert weights from [experts, out, in] // to [experts, in, out] for use with GatherMM (which computes a @ b[group]). func transposeForGatherMM(w *mlx.Array) *mlx.Array { - if w == nil || !w.Valid() || w.NumDims() != 3 { + if w == nil || w.NumDims() != 3 { return w } return mlx.Transpose(w, 0, 2, 1).Clone() diff --git a/x/models/laguna/laguna.go b/x/models/laguna/laguna.go index 07751c70e4a..fd000e5f82b 100644 --- a/x/models/laguna/laguna.go +++ b/x/models/laguna/laguna.go @@ -594,14 +594,14 @@ func stackAndClone(parts []*mlx.Array) *mlx.Array { } func transposeExpertWeightForGatherMM(w *mlx.Array) *mlx.Array { - if w == nil || !w.Valid() || w.NumDims() != 3 { + if w == nil || w.NumDims() != 3 { return w } return mlx.Transpose(w, 0, 2, 1).Clone() } func transposeExpertWeightViewForGatherMM(w *mlx.Array) *mlx.Array { - if w == nil || !w.Valid() || w.NumDims() != 3 { + if w == nil || w.NumDims() != 3 { return w } return mlx.Transpose(w, 0, 2, 1) @@ -637,7 +637,7 @@ func denseExpertWeightForGatherMM(w *stackedExpertWeights) *mlx.Array { } func denseExpertWeightSupportsSourceLayout(w *stackedExpertWeights) bool { - return w != nil && w.Weight != nil && w.Weight.Valid() && w.Scales == nil && w.Weight.DType() == mlx.DTypeBFloat16 + return w != nil && w.Weight != nil && w.Scales == nil && w.Weight.DType() == mlx.DTypeBFloat16 } func denseExpertWeightsSupportSourceLayout(weights ...*stackedExpertWeights) bool { @@ -755,7 +755,7 @@ func splitLastDim(x *mlx.Array, first int32) (*mlx.Array, *mlx.Array) { } func fuseExpertStacks(a, b *mlx.Array, axis int) *mlx.Array { - if a == nil || !a.Valid() || b == nil || !b.Valid() { + if a == nil || b == nil { return nil } return mlx.Concatenate([]*mlx.Array{a, b}, axis).Clone() diff --git a/x/models/nemotron_h/nemotron_h.go b/x/models/nemotron_h/nemotron_h.go index 2a8bdabf791..51584ed8b95 100644 --- a/x/models/nemotron_h/nemotron_h.go +++ b/x/models/nemotron_h/nemotron_h.go @@ -422,7 +422,7 @@ func applyExpertWeightGlobalScale(weight, scale *mlx.Array) *mlx.Array { } func transposeExpertWeightForGatherMM(w *mlx.Array) *mlx.Array { - if w == nil || !w.Valid() || w.NumDims() != 3 { + if w == nil || w.NumDims() != 3 { return w } return mlx.Transpose(w, 0, 2, 1).Clone() diff --git a/x/models/nn/nn.go b/x/models/nn/nn.go index 07282e911ac..95e5c0875d4 100644 --- a/x/models/nn/nn.go +++ b/x/models/nn/nn.go @@ -61,7 +61,7 @@ type Linear struct { } func NewLinear(weight *mlx.Array, bias *mlx.Array) *Linear { - if bias != nil && bias.Valid() && bias.DType() != weight.DType() { + if bias != nil && bias.DType() != weight.DType() { bias = bias.AsType(weight.DType()) } return &Linear{Weight: weight, Bias: bias} @@ -69,7 +69,7 @@ func NewLinear(weight *mlx.Array, bias *mlx.Array) *Linear { func (l *Linear) Forward(x *mlx.Array) *mlx.Array { w := l.Weight.Transpose(1, 0) - if l.Bias != nil && l.Bias.Valid() { + if l.Bias != nil { return l.Bias.Addmm(x, w, 1.0, 1.0) } return x.Matmul(w) @@ -98,7 +98,7 @@ func NewQuantizedLinear(weight *mlx.Array, bias *mlx.Array, groupSize, bits int, } else { mlx.Eval(qw, scales) } - if bias != nil && bias.Valid() && bias.DType() != weight.DType() { + if bias != nil && bias.DType() != weight.DType() { bias = bias.AsType(weight.DType()) } return &QuantizedLinear{ @@ -130,7 +130,7 @@ func (ql *QuantizedLinear) Forward(x *mlx.Array) *mlx.Array { // coverage for this path. out = quantizedLinearOutputScale(out, ql.GlobalScale) } - if ql.Bias != nil && ql.Bias.Valid() { + if ql.Bias != nil { bias := ql.Bias if bias.DType() != out.DType() { bias = bias.AsType(out.DType()) @@ -194,7 +194,7 @@ func (qe *QuantizedEmbedding) Forward(indices *mlx.Array) *mlx.Array { weight := qe.Weight.TakeAxis(indices, 0) scales := qe.Scales.TakeAxis(indices, 0) var qbiases *mlx.Array - if qe.QBiases != nil && qe.QBiases.Valid() { + if qe.QBiases != nil { qbiases = qe.QBiases.TakeAxis(indices, 0) } return mlx.Dequantize(weight, scales, qbiases, qe.GroupSize, qe.Bits, qe.Mode, qe.GlobalScale) diff --git a/x/models/qwen3_5/qwen3_5.go b/x/models/qwen3_5/qwen3_5.go index 64ca4dbbe38..9076f8508f9 100644 --- a/x/models/qwen3_5/qwen3_5.go +++ b/x/models/qwen3_5/qwen3_5.go @@ -483,14 +483,14 @@ func stackAndClone(parts []*mlx.Array) *mlx.Array { } func transposeExpertWeightForGatherMM(w *mlx.Array) *mlx.Array { - if w == nil || !w.Valid() || w.NumDims() != 3 { + if w == nil || w.NumDims() != 3 { return w } return mlx.Transpose(w, 0, 2, 1).Clone() } func fuseExpertStacks(a, b *mlx.Array, axis int) *mlx.Array { - if a == nil || !a.Valid() || b == nil || !b.Valid() { + if a == nil || b == nil { return nil } return mlx.Concatenate([]*mlx.Array{a, b}, axis).Clone() From b68b112bd8868d6278250d7d4bdfafa5cbf035c8 Mon Sep 17 00:00:00 2001 From: Jesse Gross Date: Thu, 10 Sep 2026 13:58:42 -0700 Subject: [PATCH 13/24] mlxrunner: release the buffers weight loading leaves in the MLX pool Loading a model can transform tensors after reading them: qwen3.5 models pack their linear-attention projections into one layout, and MoE models fuse the gate and up expert stacks. The buffers those transforms consume go back to MLX's allocator pool rather than to the system, and nothing releases the pool until the first request finishes. On qwen3.8:27b-mlx that is 2.15 GiB held idle on top of 16.9 GiB of weights, counted in the runner's reported memory the whole time. Clear the pool once the weights are evaluated. Models whose tensors load unchanged, such as gemma4, leave nothing in the pool and are unaffected. --- x/mlxrunner/runner.go | 1 + 1 file changed, 1 insertion(+) diff --git a/x/mlxrunner/runner.go b/x/mlxrunner/runner.go index 626353c9c14..4814edfd543 100644 --- a/x/mlxrunner/runner.go +++ b/x/mlxrunner/runner.go @@ -63,6 +63,7 @@ func (r *Runner) Load(modelName string) error { return err } mlx.Eval(weights...) + mlx.ClearCache() r.weights = mlx.NewScope() r.weights.Attach(weights...) configureWiredMemory() From c16bf9892a560a11c208618968966246e485540c Mon Sep 17 00:00:00 2001 From: Parth Sareen Date: Fri, 11 Sep 2026 12:21:52 -0700 Subject: [PATCH 14/24] cmd: remove built-in agent (#18393) --- agent/approval.go | 198 -- agent/approval_test.go | 95 - agent/compactor.go | 667 ----- agent/compactor_test.go | 773 ------ agent/events.go | 177 -- agent/registry.go | 104 - agent/session.go | 1092 --------- agent/session_test.go | 2265 ----------------- agent/skill_activation.go | 57 - agent/skill_activation_test.go | 74 - agent/skills.go | 813 ------- agent/skills_test.go | 516 ---- agent/testdata/import/release-notes/SKILL.md | 8 - agent/tools/bash.go | 450 ---- agent/tools/bash_test.go | 258 -- agent/tools/bash_unix.go | 49 - agent/tools/bash_unix_test.go | 40 - agent/tools/bash_windows.go | 134 - agent/tools/bash_windows_test.go | 15 - agent/tools/file.go | 711 ------ agent/tools/file_test.go | 571 ----- agent/tools/file_unix_test.go | 121 - agent/tools/skill.go | 41 - agent/tools/skill_test.go | 163 -- agent/tools/web.go | 186 -- agent/tools/web_test.go | 214 -- cmd/agent_tui.go | 654 ----- cmd/agent_tui_test.go | 189 -- cmd/cmd.go | 42 +- cmd/cmd_test.go | 41 + cmd/internal/filedata/filedata.go | 189 -- cmd/internal/filedata/filedata_test.go | 223 -- cmd/tui/chat/approval.go | 475 ---- cmd/tui/chat/approval_test.go | 520 ---- cmd/tui/chat/chat.go | 1246 ---------- cmd/tui/chat/clipboard.go | 66 - cmd/tui/chat/cloudauth.go | 434 ---- cmd/tui/chat/cloudauth_test.go | 247 -- cmd/tui/chat/compaction.go | 101 - cmd/tui/chat/debug.go | 565 ----- cmd/tui/chat/events.go | 374 --- cmd/tui/chat/events_test.go | 454 ---- cmd/tui/chat/input.go | 1630 ------------- cmd/tui/chat/input_test.go | 1351 ----------- cmd/tui/chat/markdown.go | 425 ---- cmd/tui/chat/modals.go | 263 -- cmd/tui/chat/modals_test.go | 441 ---- cmd/tui/chat/multimodal_test.go | 307 --- cmd/tui/chat/render.go | 2240 ----------------- cmd/tui/chat/render_test.go | 2282 ------------------ cmd/tui/chat/test_helpers_test.go | 113 - cmd/tui/chat/theme.go | 129 - cmd/tui/chat/think.go | 166 -- cmd/tui/tui.go | 4 +- cmd/tui/tui_test.go | 4 +- go.mod | 2 +- 56 files changed, 77 insertions(+), 24892 deletions(-) delete mode 100644 agent/approval.go delete mode 100644 agent/approval_test.go delete mode 100644 agent/compactor.go delete mode 100644 agent/compactor_test.go delete mode 100644 agent/events.go delete mode 100644 agent/registry.go delete mode 100644 agent/session.go delete mode 100644 agent/session_test.go delete mode 100644 agent/skill_activation.go delete mode 100644 agent/skill_activation_test.go delete mode 100644 agent/skills.go delete mode 100644 agent/skills_test.go delete mode 100644 agent/testdata/import/release-notes/SKILL.md delete mode 100644 agent/tools/bash.go delete mode 100644 agent/tools/bash_test.go delete mode 100644 agent/tools/bash_unix.go delete mode 100644 agent/tools/bash_unix_test.go delete mode 100644 agent/tools/bash_windows.go delete mode 100644 agent/tools/bash_windows_test.go delete mode 100644 agent/tools/file.go delete mode 100644 agent/tools/file_test.go delete mode 100644 agent/tools/file_unix_test.go delete mode 100644 agent/tools/skill.go delete mode 100644 agent/tools/skill_test.go delete mode 100644 agent/tools/web.go delete mode 100644 agent/tools/web_test.go delete mode 100644 cmd/agent_tui.go delete mode 100644 cmd/agent_tui_test.go delete mode 100644 cmd/internal/filedata/filedata.go delete mode 100644 cmd/internal/filedata/filedata_test.go delete mode 100644 cmd/tui/chat/approval.go delete mode 100644 cmd/tui/chat/approval_test.go delete mode 100644 cmd/tui/chat/chat.go delete mode 100644 cmd/tui/chat/clipboard.go delete mode 100644 cmd/tui/chat/cloudauth.go delete mode 100644 cmd/tui/chat/cloudauth_test.go delete mode 100644 cmd/tui/chat/compaction.go delete mode 100644 cmd/tui/chat/debug.go delete mode 100644 cmd/tui/chat/events.go delete mode 100644 cmd/tui/chat/events_test.go delete mode 100644 cmd/tui/chat/input.go delete mode 100644 cmd/tui/chat/input_test.go delete mode 100644 cmd/tui/chat/markdown.go delete mode 100644 cmd/tui/chat/modals.go delete mode 100644 cmd/tui/chat/modals_test.go delete mode 100644 cmd/tui/chat/multimodal_test.go delete mode 100644 cmd/tui/chat/render.go delete mode 100644 cmd/tui/chat/render_test.go delete mode 100644 cmd/tui/chat/test_helpers_test.go delete mode 100644 cmd/tui/chat/theme.go delete mode 100644 cmd/tui/chat/think.go diff --git a/agent/approval.go b/agent/approval.go deleted file mode 100644 index 47b68e413ea..00000000000 --- a/agent/approval.go +++ /dev/null @@ -1,198 +0,0 @@ -package agent - -import ( - "context" - "strings" - "sync" -) - -type ApprovalRequest struct { - WorkingDir string - Calls []ApprovalToolCall -} - -func (r *ApprovalRequest) AddToolCall(id, name, scope string, args map[string]any) { - r.Calls = append(r.Calls, ApprovalToolCall{ - ToolCallID: id, - ToolName: name, - Args: args, - ApprovalScope: scope, - }) -} - -type ApprovalToolCall struct { - ToolCallID string - ToolName string - Args map[string]any - ApprovalScope string -} - -type Approval struct { - Allow bool - AllowAll bool - AllowScopes []string - Reason string -} - -type ApprovalPrompter interface { - PromptApproval(context.Context, ApprovalRequest) (Approval, error) -} - -type ApprovalState struct { - mu sync.RWMutex - allowAll bool - scopes map[string]bool -} - -func (s *ApprovalState) Set(allowAll bool, scopes map[string]bool) { - if s == nil { - return - } - s.mu.Lock() - defer s.mu.Unlock() - s.allowAll = allowAll - s.scopes = cloneApprovalScopes(scopes) -} - -// GrantAll grants blanket approval for all future tool calls. -func (s *ApprovalState) GrantAll() { - if s == nil { - return - } - s.mu.Lock() - defer s.mu.Unlock() - s.allowAll = true -} - -// AllGranted reports whether blanket approval has been granted. -func (s *ApprovalState) AllGranted() bool { - if s == nil { - return false - } - s.mu.RLock() - defer s.mu.RUnlock() - return s.allowAll -} - -func (s *ApprovalState) Allows(scope string) bool { - if s == nil { - return false - } - s.mu.RLock() - defer s.mu.RUnlock() - return s.allowAll || s.scopes[scope] -} - -// Apply merges an approval's scopes and allow-all flag into the state. It -// returns true if the approval grants permission (allow-all or at least one -// scope). It does not mutate the approval; the caller sets Allow based on the -// returned value. -func (s *ApprovalState) Apply(result *Approval) bool { - if s == nil || result == nil { - return false - } - s.mu.Lock() - defer s.mu.Unlock() - granted := false - if result.AllowAll { - s.allowAll = true - granted = true - } - if len(result.AllowScopes) > 0 { - granted = true - s.grantScopesLocked(result.AllowScopes) - } - return granted -} - -// GrantScopes merges the given scopes into the state. -func (s *ApprovalState) GrantScopes(scopes []string) { - if s == nil { - return - } - s.mu.Lock() - defer s.mu.Unlock() - s.grantScopesLocked(scopes) -} - -// grantScopesLocked adds trimmed, non-empty scopes to the state. Caller must -// hold s.mu. -func (s *ApprovalState) grantScopesLocked(scopes []string) { - if s.scopes == nil { - s.scopes = make(map[string]bool, len(scopes)) - } - for _, scope := range scopes { - scope = strings.TrimSpace(scope) - if scope != "" { - s.scopes[scope] = true - } - } -} - -func cloneApprovalScopes(src map[string]bool) map[string]bool { - if len(src) == 0 { - return nil - } - dst := make(map[string]bool, len(src)) - for scope, allowed := range src { - if allowed { - dst[scope] = true - } - } - return dst -} - -func (s *Session) needsApproval(tool Tool, name string, args map[string]any) bool { - return ToolRequiresApproval(tool, args) && !s.allows(toolApprovalScope(tool, name, args)) -} - -// allows reports whether scope is permitted by the session's accumulated approval state. -func (s *Session) allows(scope string) bool { - if s == nil || s.ApprovalState == nil { - return false - } - return s.ApprovalState.Allows(scope) -} - -// applyApproval merges an approval result into the session's state and marks -// the result as allowed when scopes or allow-all were granted. -func (s *Session) applyApproval(result *Approval) { - if s == nil || result == nil { - return - } - if s.ApprovalState == nil { - s.ApprovalState = &ApprovalState{} - } - if s.ApprovalState.Apply(result) { - result.Allow = true - } -} - -func (s *Session) authorizeToolCalls(ctx context.Context, req ApprovalRequest) (Approval, error) { - if s == nil || len(req.Calls) == 0 || (s.ApprovalState != nil && s.ApprovalState.AllGranted()) { - return Approval{Allow: true}, nil - } - if s.ApprovalPrompter == nil { - return Approval{ - Reason: "Tool execution requires approval, but no approval prompter is available.", - }, nil - } - - result, err := s.ApprovalPrompter.PromptApproval(ctx, req) - if err != nil { - return Approval{}, err - } - s.applyApproval(&result) - return result, nil -} - -// toolApprovalScope returns the approval scope key for a tool invocation. -// If the tool implements ScopedTool, its ApprovalScope method determines the -// scope (e.g. shell tools scope to "\x00"). Otherwise the scope -// is the trimmed tool name. -func toolApprovalScope(tool Tool, toolName string, args map[string]any) string { - if scoped, ok := tool.(ScopedTool); ok { - return scoped.ApprovalScope(args) - } - return strings.TrimSpace(toolName) -} diff --git a/agent/approval_test.go b/agent/approval_test.go deleted file mode 100644 index 726a3890885..00000000000 --- a/agent/approval_test.go +++ /dev/null @@ -1,95 +0,0 @@ -package agent - -import ( - "context" - "strings" - "testing" - - "github.com/ollama/ollama/api" -) - -type mockTool struct { - name string -} - -func (m mockTool) Name() string { return m.name } -func (m mockTool) Description() string { return "" } -func (m mockTool) Schema() api.ToolFunction { - return api.ToolFunction{Name: m.name} -} - -func (m mockTool) Execute(context.Context, ToolContext, map[string]any) (ToolResult, error) { - return ToolResult{}, nil -} - -func TestToolApprovalScopeUsesScopedTool(t *testing.T) { - shellTool := mockScopedTool{ - mockTool: mockTool{name: "bash"}, - scope: func(args map[string]any) string { - if cmd, ok := args["command"].(string); ok { - cmd = strings.TrimSpace(cmd) - if cmd != "" { - return "bash\x00" + cmd - } - } - return "bash" - }, - } - plainTool := mockTool{name: "edit"} - - tests := []struct { - tool Tool - name string - args map[string]any - want string - }{ - {shellTool, "bash", map[string]any{"command": " pwd "}, "bash\x00pwd"}, - {shellTool, "bash", map[string]any{"command": "Get-ChildItem"}, "bash\x00Get-ChildItem"}, - {plainTool, "edit", map[string]any{"path": "README.md"}, "edit"}, - } - for _, tt := range tests { - if got := toolApprovalScope(tt.tool, tt.name, tt.args); got != tt.want { - t.Fatalf("toolApprovalScope(%q) = %q, want %q", tt.name, got, tt.want) - } - } -} - -type mockScopedTool struct { - mockTool - scope func(args map[string]any) string -} - -func (m mockScopedTool) ApprovalScope(args map[string]any) string { - return m.scope(args) -} - -func TestSessionApplyApprovalScopes(t *testing.T) { - session := &Session{} - result := Approval{AllowScopes: []string{"edit", "bash\x00pwd", " "}} - - session.applyApproval(&result) - - if !result.Allow { - t.Fatal("scoped approval should allow the current request") - } - if !session.allows("edit") || !session.allows("bash\x00pwd") { - t.Fatal("scoped approval was not saved") - } - if session.allows("bash") || session.allows("bash\x00ls") { - t.Fatal("shell approval was too broad") - } - if session.ApprovalState.AllGranted() { - t.Fatal("allow all = true, want false for scoped approval") - } -} - -func TestSessionApplyApprovalAllowAll(t *testing.T) { - session := &Session{} - result := Approval{AllowAll: true} - - session.applyApproval(&result) - - if !result.Allow || !session.allows("anything") { - t.Fatalf("allow all = %v result = %#v, want allow all", session.ApprovalState.AllGranted(), result) - } -} diff --git a/agent/compactor.go b/agent/compactor.go deleted file mode 100644 index e0f6ce454d1..00000000000 --- a/agent/compactor.go +++ /dev/null @@ -1,667 +0,0 @@ -package agent - -import ( - "context" - "encoding/json" - "errors" - "fmt" - "net/http" - "strings" - - "github.com/ollama/ollama/api" -) - -// Compaction wire-format. These constants and helpers are the single canonical -// definition of how a compacted turn is represented in message history. -const ( - CompactionSummaryMessagePrefix = "Conversation summary:\n" - CompactionToolName = "summary" - CompactionToolCallID = "ollama_compaction" - CompactionContinueInstruction = "continue the task in progress. the history has been compacted, do not mention compaction to the user" -) - -const ( - defaultCompactionContextWindowTokens = 32768 - defaultCompactionKeepUserTurns = 3 - defaultCompactionThreshold = 0.8 - compactOnlySummaryContextTokens = 16000 - - maxCompactionSummaryRunes = 16 * 1024 - - compactionSystemPrompt = "Summarize the archived part of an Ollama agent conversation. Preserve user goals, decisions, files, commands, tool results, and unresolved tasks needed to continue. Omit private reasoning and return only the summary." -) - -type Compactor interface { - MaybeCompact(context.Context, CompactionRequest) (CompactionResult, error) - - // ContextWindowTokens returns the effective context window size in - // tokens, resolving runtime options against configured defaults. - ContextWindowTokens(options map[string]any) int - - // Threshold returns the compaction threshold as a fraction of the - // context window (e.g. 0.8 means compact at 80% capacity). - Threshold() float64 - - // ShouldCompact reports whether a compaction should run and returns the - // trigger reason. An empty trigger means compaction is not needed. - ShouldCompact(req CompactionRequest) (trigger string, should bool) -} - -type CompactionOptions struct { - ContextWindowTokens int - KeepUserTurns int - Threshold float64 -} - -type CompactionRequest struct { - ChatID string - Model string - SystemPrompt string - Messages []api.Message - Tools api.Tools - Format string - Latest api.ChatResponse - Options map[string]any - KeepAlive *api.Duration - Think *api.ThinkValue - Force bool - ContinueTask bool - KeepUserTurns *int - Progress func(CompactionProgress) -} - -type CompactionProgress struct { - Tokens int -} - -type CompactionResult struct { - Messages []api.Message - Compacted bool - Due bool - Summary string - Reason string -} - -type SimpleCompactor struct { - Client ChatClient - Options CompactionOptions -} - -func (c *SimpleCompactor) MaybeCompact(ctx context.Context, req CompactionRequest) (CompactionResult, error) { - result := CompactionResult{Messages: req.Messages} - if c == nil { - return result, nil - } - - result.Due = req.Force || c.shouldCompact(req) - if !result.Due { - return result, nil - } - if c.Client == nil { - result.Reason = "compaction is unavailable" - return result, nil - } - - keepUserTurns := c.keepUserTurns(req.Options) - if req.KeepUserTurns != nil { - keepUserTurns = *req.KeepUserTurns - } - prefix, previousSummary, archive, suffix, _, ok := splitCompactionMessages(req.Messages, keepUserTurns) - if !ok || len(archive) == 0 { - result.Reason = "nothing to compact" - return result, nil - } - - summary, err := c.summarize(ctx, req, previousSummary, archive) - if err != nil { - result.Reason = err.Error() - return result, err - } - summary = truncateCompactionSummary(strings.TrimSpace(summary)) - if summary == "" { - summary, err = c.summarizeEmptyFallback(ctx, req, previousSummary, archive) - if err != nil { - result.Reason = err.Error() - return result, err - } - summary = truncateCompactionSummary(strings.TrimSpace(summary)) - } - if summary == "" { - result.Reason = "summary was empty" - return result, nil - } - - compacted := make([]api.Message, 0, len(prefix)+len(suffix)+2) - compacted = append(compacted, prefix...) - compacted = append(compacted, CompactionSummaryMessages(summary, req.ContinueTask)...) - compacted = append(compacted, suffix...) - result.Messages = compacted - result.Compacted = true - result.Summary = summary - return result, nil -} - -func (c *SimpleCompactor) shouldCompact(req CompactionRequest) bool { - contextWindow := c.contextWindowTokens(req.Options) - threshold := int(float64(contextWindow) * c.threshold()) - if threshold <= 0 { - return false - } - if req.Latest.PromptEvalCount > 0 && req.Latest.PromptEvalCount >= threshold { - return true - } - return estimateCompactionRequestTokens(req) >= threshold -} - -func (c *SimpleCompactor) contextWindowTokens(options map[string]any) int { - return ResolveContextWindowTokens(options, c.Options.ContextWindowTokens) -} - -// ContextWindowTokens resolves the effective context window from runtime -// options or configured defaults. Satisfies the Compactor interface. -func (c *SimpleCompactor) ContextWindowTokens(options map[string]any) int { - if c == nil { - return 0 - } - return c.contextWindowTokens(options) -} - -func (c *SimpleCompactor) threshold() float64 { - return ResolveCompactionThreshold(c.Options.Threshold) -} - -// Threshold returns the configured compaction threshold fraction. Satisfies -// the Compactor interface. -func (c *SimpleCompactor) Threshold() float64 { - if c == nil { - return 0 - } - return c.threshold() -} - -// ShouldCompact reports whether compaction is due and the trigger reason. -// Satisfies the Compactor interface. -func (c *SimpleCompactor) ShouldCompact(req CompactionRequest) (string, bool) { - if c == nil { - return "", false - } - if req.Force { - return "force", true - } - if c.shouldCompact(req) { - contextWindow := c.contextWindowTokens(req.Options) - threshold := int(float64(contextWindow) * c.threshold()) - if req.Latest.PromptEvalCount > 0 && req.Latest.PromptEvalCount >= threshold { - return "prompt_eval", true - } - return "estimate", true - } - return "", false -} - -func (c *SimpleCompactor) keepUserTurns(options map[string]any) int { - contextWindow := c.contextWindowTokens(options) - if contextWindow > 0 && contextWindow < compactOnlySummaryContextTokens { - return 0 - } - if c.Options.KeepUserTurns > 0 { - return c.Options.KeepUserTurns - } - return defaultCompactionKeepUserTurns -} - -func ResolveContextWindowTokens(options map[string]any, configured int) int { - if n := intOption(options, "num_ctx"); n > 0 { - return n - } - if configured > 0 { - return configured - } - return defaultCompactionContextWindowTokens -} - -func ResolveCompactionThreshold(configured float64) float64 { - if configured > 0 { - return configured - } - return defaultCompactionThreshold -} - -func (c *SimpleCompactor) summarize(ctx context.Context, req CompactionRequest, previousSummary string, archive []api.Message) (string, error) { - body, err := compactionPrompt(previousSummary, archive, c.compactionPromptBodyBudgetTokens(req.Options)) - if err != nil { - return "", err - } - - chatReq := &api.ChatRequest{ - Model: req.Model, - Messages: []api.Message{ - { - Role: "system", - Content: compactionSystemPrompt, - }, - { - Role: "user", - Content: body, - }, - }, - Options: req.Options, - Think: req.Think, - } - if req.KeepAlive != nil { - chatReq.KeepAlive = req.KeepAlive - } - - var summary strings.Builder - if err := c.Client.Chat(ctx, chatReq, func(response api.ChatResponse) error { - summary.WriteString(response.Message.Content) - if req.Progress != nil { - tokens := response.EvalCount - if tokens <= 0 { - tokens = estimateCompactionTokens(summary.String()) - } - req.Progress(CompactionProgress{Tokens: tokens}) - } - return nil - }); err != nil { - return "", err - } - return summary.String(), nil -} - -func (c *SimpleCompactor) summarizeEmptyFallback(ctx context.Context, req CompactionRequest, previousSummary string, archive []api.Message) (string, error) { - retry := req - retry.Think = &api.ThinkValue{Value: false} - summary, err := c.summarize(ctx, retry, previousSummary, archive) - if err == nil { - return summary, nil - } - if !isUnsupportedCompactionThinkError(err) { - return "", err - } - if req.Think == nil { - return "", nil - } - retry.Think = nil - return c.summarize(ctx, retry, previousSummary, archive) -} - -func isUnsupportedCompactionThinkError(err error) bool { - if err == nil { - return false - } - text := strings.ToLower(err.Error()) - if !strings.Contains(text, "think") { - return false - } - var statusErr api.StatusError - if errors.As(err, &statusErr) && statusErr.StatusCode != 0 { - return statusErr.StatusCode == http.StatusBadRequest - } - return strings.Contains(text, "does not support") || strings.Contains(text, "not supported") || strings.Contains(text, "unsupported") -} - -// compactionSummaryMessageForTask renders a compaction summary as the content -// string stored on the synthetic tool-result message. -func compactionSummaryMessageForTask(summary string, continueTask bool) string { - content := CompactionSummaryMessagePrefix + strings.TrimSpace(summary) - if continueTask { - content = strings.TrimSpace(content) + "\n\n" + CompactionContinueInstruction - } - return content -} - -// CompactionSummaryMessages renders a compaction summary as the assistant -// tool-call plus tool-result pair that represents a compacted turn in the -// message history. -func CompactionSummaryMessages(summary string, continueTask bool) []api.Message { - return []api.Message{ - { - Role: "assistant", - ToolCalls: []api.ToolCall{{ - ID: CompactionToolCallID, - Function: api.ToolCallFunction{ - Name: CompactionToolName, - }, - }}, - }, - { - Role: "tool", - ToolName: CompactionToolName, - ToolCallID: CompactionToolCallID, - Content: compactionSummaryMessageForTask(summary, continueTask), - }, - } -} - -func (c *SimpleCompactor) compactionPromptBodyBudgetTokens(options map[string]any) int { - contextWindow := c.contextWindowTokens(options) - threshold := int(float64(contextWindow) * c.threshold()) - if threshold <= 0 { - return 0 - } - systemTokens := estimateCompactionTokens("system") + estimateCompactionTokens(compactionSystemPrompt) - userRoleTokens := estimateCompactionTokens("user") - budget := threshold - systemTokens - userRoleTokens - if budget <= 0 { - return 0 - } - return budget -} - -func truncateCompactionSummary(summary string) string { - return Truncate(summary, TruncateConfig{ - MaxRunes: maxCompactionSummaryRunes, - Label: "summary", - }) -} - -func estimateCompactionTokens(text string) int { - text = strings.TrimSpace(text) - if text == "" { - return 0 - } - return ApproximateTokens(len([]rune(text))) -} - -func estimateMessagesTokens(messages []api.Message) int { - var total int - for _, msg := range messages { - total += estimateCompactionTokens(msg.Role) - total += estimateCompactionTokens(msg.Content) - total += estimateCompactionTokens(msg.Thinking) - total += estimateCompactionTokens(msg.ToolName) - total += estimateCompactionTokens(msg.ToolCallID) - for _, call := range msg.ToolCalls { - total += estimateCompactionTokens(call.Function.Name) - total += estimateCompactionTokens(call.Function.Arguments.String()) - } - } - return total -} - -func estimateCompactionRequestTokens(req CompactionRequest) int { - requestMessages := sanitizeMessagesForEstimate(req.Messages) - if strings.TrimSpace(req.SystemPrompt) != "" { - requestMessages = make([]api.Message, 0, len(req.Messages)+1) - requestMessages = append(requestMessages, api.Message{Role: "system", Content: strings.TrimSpace(req.SystemPrompt)}) - requestMessages = append(requestMessages, sanitizeMessagesForEstimate(req.Messages)...) - } - - payload := struct { - Messages []api.Message `json:"messages,omitempty"` - Tools api.Tools `json:"tools,omitempty"` - Format json.RawMessage `json:"format,omitempty"` - }{ - Messages: requestMessages, - Tools: req.Tools, - } - if rawFormat, ok := compactionFormatForEstimate(req.Format); ok { - payload.Format = rawFormat - } - if data, err := json.Marshal(payload); err == nil { - return estimateCompactionTokens(string(data)) - } - - total := estimateMessagesTokens(requestMessages) - total += estimateCompactionTokens(req.Tools.String()) - total += estimateCompactionTokens(req.Format) - return total -} - -func (s *Session) estimateRunPromptTokens(opts RunOptions, messages []api.Message) int { - return estimateCompactionRequestTokens(CompactionRequest{ - SystemPrompt: opts.SystemPrompt, - Messages: messages, - Tools: s.availableTools(), - Format: opts.Format, - Options: opts.Options, - }) -} - -func (s *Session) checkPreflightPromptBudget(opts RunOptions, messages []api.Message) error { - contextWindow := s.contextWindowTokens(opts) - if contextWindow <= 0 { - return nil - } - estimated := s.estimateRunPromptTokens(opts, messages) - if estimated < contextWindow { - return nil - } - return fmt.Errorf("prompt is too large for the current context (~%d/%d tokens). Reduce the system prompt or message history, compact the conversation, or use a model with a larger context", estimated, contextWindow) -} - -func (s *Session) checkPostCompactionPromptBudget(opts RunOptions, messages []api.Message) error { - contextWindow := s.contextWindowTokens(opts) - if contextWindow <= 0 { - return nil - } - estimated := s.estimateRunPromptTokens(opts, messages) - if estimated < contextWindow { - return nil - } - return fmt.Errorf("history is still too large after compaction (~%d/%d tokens). Start a fresh request, reduce the system prompt or history, or use a model with a larger context", estimated, contextWindow) -} - -func sanitizeMessagesForEstimate(messages []api.Message) []api.Message { - requestMessages := sanitizeMessagesForRequest(messages) - for i := range requestMessages { - // Image token accounting is model-specific. Without the active model's - // tokenizer and vision accounting, raw image bytes/base64 make the - // estimate look much larger than the prompt the model actually sees. - requestMessages[i].Images = nil - } - return requestMessages -} - -func compactionFormatForEstimate(format string) (json.RawMessage, bool) { - format = strings.TrimSpace(format) - if format == "" { - return nil, false - } - if format == "json" { - return json.RawMessage(`"json"`), true - } - if !json.Valid([]byte(format)) { - return nil, false - } - return json.RawMessage(format), true -} - -func compactionPrompt(previousSummary string, archive []api.Message, maxTokens int) (string, error) { - messages := make([]api.Message, 0, len(archive)) - for _, msg := range archive { - msg.Thinking = "" - msg.Images = nil - messages = append(messages, msg) - } - return renderCompactionPrompt(previousSummary, fitCompactionMessagesToBudget(previousSummary, messages, maxTokens)) -} - -func renderCompactionPrompt(previousSummary string, messages []api.Message) (string, error) { - payload, err := json.MarshalIndent(messages, "", " ") - if err != nil { - return "", fmt.Errorf("marshal compaction messages: %w", err) - } - - var b strings.Builder - if strings.TrimSpace(previousSummary) != "" { - b.WriteString("Previous summary:\n") - b.WriteString(strings.TrimSpace(previousSummary)) - b.WriteString("\n\n") - } - b.WriteString("Messages to archive as JSON:\n") - b.Write(payload) - return b.String(), nil -} - -func fitCompactionMessagesToBudget(previousSummary string, messages []api.Message, maxTokens int) []api.Message { - if maxTokens <= 0 { - return messages - } - fitted := append([]api.Message(nil), messages...) - for range 16 { - body, err := renderCompactionPrompt(previousSummary, fitted) - if err != nil || estimateCompactionTokens(body) <= maxTokens { - return fitted - } - - idx := largestCompactionContentMessage(fitted) - if idx < 0 { - return fitted - } - overageTokens := estimateCompactionTokens(body) - maxTokens - currentRunes := len([]rune(fitted[idx].Content)) - nextRunes := currentRunes - overageTokens*4 - 256 - if nextRunes >= currentRunes { - nextRunes = currentRunes / 2 - } - fitted[idx].Content = truncateToolResultContentTo(fitted[idx].Content, nextRunes) - } - return fitted -} - -func largestCompactionContentMessage(messages []api.Message) int { - idx := -1 - size := 0 - for i, msg := range messages { - n := len([]rune(msg.Content)) - if n > size { - idx = i - size = n - } - } - return idx -} - -func splitCompactionMessages(messages []api.Message, keepUserTurns int) (prefix []api.Message, previousSummary string, archive []api.Message, suffix []api.Message, keptUserTurns int, ok bool) { - if keepUserTurns < 0 { - keepUserTurns = defaultCompactionKeepUserTurns - } - - start := 0 - for start < len(messages) && messages[start].Role == "system" && !isCompactionSummary(messages[start]) { - prefix = append(prefix, messages[start]) - start++ - } - - candidates := make([]api.Message, 0, len(messages)-start) - for i := start; i < len(messages); i++ { - msg := messages[i] - if isCompactionSummary(msg) { - previousSummary = CompactionSummaryText(msg.Content) - continue - } - if isCompactionToolCall(msg) { - if i+1 < len(messages) && isCompactionSummary(messages[i+1]) { - previousSummary = CompactionSummaryText(messages[i+1].Content) - i++ - } - continue - } - candidates = append(candidates, msg) - } - - userTurnIndexes := make([]int, 0, keepUserTurns) - for i := len(candidates) - 1; i >= 0; i-- { - if candidates[i].Role == "user" { - userTurnIndexes = append(userTurnIndexes, i) - } - } - keptUserTurns = keepUserTurns - if len(userTurnIndexes) <= keptUserTurns { - keptUserTurns = len(userTurnIndexes) - 1 - } - if keptUserTurns < 0 { - keptUserTurns = 0 - } - - suffixStart := len(candidates) - if keptUserTurns > 0 { - suffixStart = userTurnIndexes[keptUserTurns-1] - } - if suffixStart <= 0 || len(candidates[:suffixStart]) == 0 { - return prefix, previousSummary, nil, nil, keptUserTurns, false - } - - return prefix, previousSummary, candidates[:suffixStart], candidates[suffixStart:], keptUserTurns, true -} - -func isCompactionToolName(name string) bool { - return name == CompactionToolName -} - -func isCompactionSummary(msg api.Message) bool { - return (msg.Role == "user" || msg.Role == "system" || (msg.Role == "tool" && isCompactionToolName(msg.ToolName))) && - strings.HasPrefix(msg.Content, CompactionSummaryMessagePrefix) -} - -// IsCompactionSummary reports whether msg uses the canonical compaction -// summary message representation. -func IsCompactionSummary(msg api.Message) bool { - return isCompactionSummary(msg) -} - -// CompactionSummaryContent returns the user-visible summary from msg when it -// is a canonical compaction summary. -func CompactionSummaryContent(msg api.Message) (string, bool) { - if !isCompactionSummary(msg) { - return "", false - } - return CompactionSummaryText(msg.Content), true -} - -// IsCompactionToolResult reports whether msg is the synthetic tool result used -// to represent compaction in message history. -func IsCompactionToolResult(msg api.Message) bool { - return msg.Role == "tool" && (isCompactionToolName(msg.ToolName) || msg.ToolCallID == CompactionToolCallID) -} - -// IsCompactionToolCall reports whether msg is the synthetic assistant tool -// call paired with a compaction summary result. -func IsCompactionToolCall(msg api.Message) bool { - return isCompactionToolCall(msg) -} - -func isCompactionToolCall(msg api.Message) bool { - if msg.Role != "assistant" { - return false - } - for _, call := range msg.ToolCalls { - if isCompactionToolName(call.Function.Name) { - return true - } - } - return false -} - -// CompactionSummaryText reverses CompactionSummaryMessages, returning the -// user-visible summary text with the prefix and any continuation instruction -// removed. -func CompactionSummaryText(content string) string { - return strings.TrimSpace(strings.TrimSuffix( - strings.TrimSpace(strings.TrimPrefix(content, CompactionSummaryMessagePrefix)), - CompactionContinueInstruction, - )) -} - -func intOption(options map[string]any, key string) int { - if options == nil { - return 0 - } - switch v := options[key].(type) { - case int: - return v - case int64: - return int(v) - case float64: - return int(v) - case float32: - return int(v) - case json.Number: - n, _ := v.Int64() - return int(n) - default: - return 0 - } -} diff --git a/agent/compactor_test.go b/agent/compactor_test.go deleted file mode 100644 index 79eba1f371d..00000000000 --- a/agent/compactor_test.go +++ /dev/null @@ -1,773 +0,0 @@ -package agent - -import ( - "context" - "net/http" - "strings" - "testing" - - "github.com/ollama/ollama/api" -) - -type scriptedCompactionClient struct { - responses [][]api.ChatResponse - errs []error - requests []*api.ChatRequest -} - -func (c *scriptedCompactionClient) Chat(_ context.Context, req *api.ChatRequest, fn api.ChatResponseFunc) error { - c.requests = append(c.requests, req) - i := len(c.requests) - 1 - if i < len(c.responses) { - for _, response := range c.responses[i] { - if err := fn(response); err != nil { - return err - } - } - } - if i < len(c.errs) { - return c.errs[i] - } - return nil -} - -func assertCompactionSummaryPair(t *testing.T, messages []api.Message) { - t.Helper() - if len(messages) != 2 { - t.Fatalf("compaction summary pair len = %d, want 2: %#v", len(messages), messages) - } - if messages[0].Role != "assistant" || len(messages[0].ToolCalls) != 1 || messages[0].ToolCalls[0].Function.Name != CompactionToolName { - t.Fatalf("compaction assistant message = %#v", messages[0]) - } - if messages[0].ToolCalls[0].Function.Arguments.Len() != 0 { - t.Fatalf("compaction summary tool call should not have arguments: %#v", messages[0].ToolCalls[0].Function.Arguments.ToMap()) - } - if messages[1].Role != "tool" || messages[1].ToolName != CompactionToolName || messages[1].ToolCallID != messages[0].ToolCalls[0].ID { - t.Fatalf("compaction tool result = %#v", messages[1]) - } - if !strings.HasPrefix(messages[1].Content, CompactionSummaryMessagePrefix) { - t.Fatalf("compaction tool result missing summary prefix: %#v", messages[1]) - } -} - -func TestSimpleCompactorSummarizesOldMessages(t *testing.T) { - client := &fakeClient{ - responses: [][]api.ChatResponse{{ - {Message: api.Message{Role: "assistant", Content: "summary"}}, - }}, - } - compactor := &SimpleCompactor{Client: client, Options: CompactionOptions{ - ContextWindowTokens: 16000, - KeepUserTurns: 2, - Threshold: 0.5, - }} - - messages := []api.Message{ - {Role: "system", Content: "stay pinned"}, - {Role: "user", Content: "old request"}, - {Role: "assistant", Content: "old answer", Thinking: "hidden"}, - {Role: "user", Content: "recent one"}, - {Role: "assistant", Content: "recent answer"}, - {Role: "user", Content: "recent two"}, - } - - result, err := compactor.MaybeCompact(context.Background(), CompactionRequest{ - ChatID: "chat-1", - Model: "model", - Messages: messages, - Latest: api.ChatResponse{Metrics: api.Metrics{PromptEvalCount: 12000}}, - }) - if err != nil { - t.Fatal(err) - } - if !result.Compacted { - t.Fatal("expected compaction") - } - compacted := result.Messages - if len(compacted) != 6 { - t.Fatalf("compacted messages = %d, want 6", len(compacted)) - } - if compacted[0].Content != "stay pinned" { - t.Fatalf("first message = %#v", compacted[0]) - } - if result.Summary != "summary" { - t.Fatalf("result summary = %q", result.Summary) - } - assertCompactionSummaryPair(t, compacted[1:3]) - if compacted[3].Content != "recent one" || compacted[5].Content != "recent two" { - t.Fatalf("recent turns were not kept: %#v", compacted) - } - if len(client.requests) != 1 { - t.Fatalf("summary requests = %d, want 1", len(client.requests)) - } - if strings.Contains(client.requests[0].Messages[1].Content, "hidden") { - t.Fatal("compaction prompt should omit thinking") - } -} - -func TestSimpleCompactorKeepsOnlySummaryForSmallContext(t *testing.T) { - client := &fakeClient{ - responses: [][]api.ChatResponse{{ - {Message: api.Message{Role: "assistant", Content: "small context summary"}}, - }}, - } - compactor := &SimpleCompactor{Client: client, Options: CompactionOptions{ - ContextWindowTokens: compactOnlySummaryContextTokens - 1, - KeepUserTurns: 3, - Threshold: 0.5, - }} - - result, err := compactor.MaybeCompact(context.Background(), CompactionRequest{ - ChatID: "chat-1", - Model: "model", - ContinueTask: true, - Messages: []api.Message{ - {Role: "system", Content: "pinned"}, - {Role: "user", Content: "old request"}, - {Role: "assistant", Content: "old answer"}, - {Role: "user", Content: "latest request"}, - }, - Latest: api.ChatResponse{Metrics: api.Metrics{PromptEvalCount: 12000}}, - }) - if err != nil { - t.Fatal(err) - } - if !result.Compacted { - t.Fatal("expected compaction") - } - if len(result.Messages) != 3 { - t.Fatalf("messages = %#v, want system plus compaction summary pair", result.Messages) - } - if result.Messages[0].Content != "pinned" { - t.Fatalf("leading system message not kept: %#v", result.Messages) - } - assertCompactionSummaryPair(t, result.Messages[1:]) - if !strings.Contains(result.Messages[2].Content, CompactionContinueInstruction) { - t.Fatalf("tool result missing continue instruction: %q", result.Messages[2].Content) - } -} - -func TestSimpleCompactorAddsContinueTaskInstructionOnlyToToolResult(t *testing.T) { - client := &fakeClient{ - responses: [][]api.ChatResponse{{ - {Message: api.Message{Role: "assistant", Content: "summary"}}, - }}, - } - compactor := &SimpleCompactor{Client: client, Options: CompactionOptions{ - ContextWindowTokens: 100, - KeepUserTurns: 1, - Threshold: 0.5, - }} - - result, err := compactor.MaybeCompact(context.Background(), CompactionRequest{ - ChatID: "chat-1", - Model: "model", - ContinueTask: true, - Messages: []api.Message{ - {Role: "user", Content: "old request"}, - {Role: "assistant", Content: "old answer"}, - {Role: "user", Content: "recent request"}, - }, - Latest: api.ChatResponse{Metrics: api.Metrics{PromptEvalCount: 75}}, - }) - if err != nil { - t.Fatal(err) - } - if result.Summary != "summary" { - t.Fatalf("result summary = %q", result.Summary) - } - content := result.Messages[1].Content - if !strings.Contains(content, CompactionContinueInstruction) { - t.Fatalf("tool result missing continue instruction: %q", content) - } - if got := CompactionSummaryText(content); got != "summary" { - t.Fatalf("visible summary text = %q", got) - } -} - -func TestSimpleCompactorTruncatesOversizedSummary(t *testing.T) { - longSummary := strings.Repeat("x", maxCompactionSummaryRunes+1024) - client := &fakeClient{ - responses: [][]api.ChatResponse{{ - {Message: api.Message{Role: "assistant", Content: longSummary}}, - }}, - } - compactor := &SimpleCompactor{Client: client, Options: CompactionOptions{ - ContextWindowTokens: 100, - KeepUserTurns: 1, - Threshold: 0.5, - }} - - result, err := compactor.MaybeCompact(context.Background(), CompactionRequest{ - ChatID: "chat-1", - Model: "model", - Messages: []api.Message{ - {Role: "user", Content: "old one"}, - {Role: "assistant", Content: "old answer"}, - {Role: "user", Content: "recent one"}, - }, - Latest: api.ChatResponse{Metrics: api.Metrics{PromptEvalCount: 75}}, - }) - if err != nil { - t.Fatal(err) - } - if !result.Compacted { - t.Fatal("expected compaction") - } - if runeCount := len([]rune(result.Summary)); runeCount > maxCompactionSummaryRunes+200 { - t.Fatalf("summary runes = %d, want <= %d (plus marker)", runeCount, maxCompactionSummaryRunes) - } - if !strings.Contains(result.Summary, "[summary truncated:") { - t.Fatalf("summary missing truncation marker: %q", result.Summary) - } - if !strings.Contains(result.Messages[1].Content, "[summary truncated:") { - t.Fatalf("compacted message missing truncation marker: %#v", result.Messages) - } -} - -func TestSimpleCompactorRetriesEmptySummaryWithThinkFalse(t *testing.T) { - client := &scriptedCompactionClient{ - responses: [][]api.ChatResponse{ - {{Message: api.Message{Role: "assistant", Thinking: "internal summary plan"}}}, - {{Message: api.Message{Role: "assistant", Content: "fallback summary"}}}, - }, - } - compactor := &SimpleCompactor{Client: client, Options: CompactionOptions{ - ContextWindowTokens: 100, - KeepUserTurns: 1, - Threshold: 0.5, - }} - - result, err := compactor.MaybeCompact(context.Background(), CompactionRequest{ - Model: "model", - Messages: []api.Message{ - {Role: "user", Content: "old request"}, - {Role: "assistant", Content: "old answer"}, - {Role: "user", Content: "recent request"}, - }, - Force: true, - }) - if err != nil { - t.Fatal(err) - } - if !result.Compacted || result.Summary != "fallback summary" { - t.Fatalf("compaction result = %#v", result) - } - if len(client.requests) != 2 { - t.Fatalf("summary requests = %d, want 2", len(client.requests)) - } - if client.requests[0].Think != nil { - t.Fatalf("first summary request think = %#v, want nil", client.requests[0].Think) - } - if client.requests[1].Think == nil || client.requests[1].Think.Value != false { - t.Fatalf("fallback summary request think = %#v, want false", client.requests[1].Think) - } -} - -func TestSimpleCompactorIgnoresUnsupportedThinkFalseFallback(t *testing.T) { - client := &scriptedCompactionClient{ - responses: [][]api.ChatResponse{ - {{Message: api.Message{Role: "assistant", Thinking: "internal summary plan"}}}, - nil, - }, - errs: []error{ - nil, - api.StatusError{StatusCode: http.StatusBadRequest, ErrorMessage: "model does not support thinking"}, - }, - } - compactor := &SimpleCompactor{Client: client, Options: CompactionOptions{ - ContextWindowTokens: 100, - KeepUserTurns: 1, - Threshold: 0.5, - }} - - result, err := compactor.MaybeCompact(context.Background(), CompactionRequest{ - Model: "model", - Messages: []api.Message{ - {Role: "user", Content: "old request"}, - {Role: "assistant", Content: "old answer"}, - {Role: "user", Content: "recent request"}, - }, - Force: true, - }) - if err != nil { - t.Fatal(err) - } - if result.Compacted || result.Reason != "summary was empty" { - t.Fatalf("compaction result = %#v", result) - } - if len(client.requests) != 2 { - t.Fatalf("summary requests = %d, want 2", len(client.requests)) - } - if client.requests[1].Think == nil || client.requests[1].Think.Value != false { - t.Fatalf("fallback summary request think = %#v, want false", client.requests[1].Think) - } -} - -func TestSimpleCompactorFallsBackToUnsetThinkWhenThinkFalseUnsupported(t *testing.T) { - client := &scriptedCompactionClient{ - responses: [][]api.ChatResponse{ - {{Message: api.Message{Role: "assistant", Thinking: "internal summary plan"}}}, - nil, - {{Message: api.Message{Role: "assistant", Content: "unset think summary"}}}, - }, - errs: []error{ - nil, - api.StatusError{StatusCode: http.StatusBadRequest, ErrorMessage: "think level is not supported"}, - nil, - }, - } - compactor := &SimpleCompactor{Client: client, Options: CompactionOptions{ - ContextWindowTokens: 100, - KeepUserTurns: 1, - Threshold: 0.5, - }} - thinkHigh := &api.ThinkValue{Value: "high"} - - result, err := compactor.MaybeCompact(context.Background(), CompactionRequest{ - Model: "model", - Messages: []api.Message{ - {Role: "user", Content: "old request"}, - {Role: "assistant", Content: "old answer"}, - {Role: "user", Content: "recent request"}, - }, - Think: thinkHigh, - Force: true, - }) - if err != nil { - t.Fatal(err) - } - if !result.Compacted || result.Summary != "unset think summary" { - t.Fatalf("compaction result = %#v", result) - } - if len(client.requests) != 3 { - t.Fatalf("summary requests = %d, want 3", len(client.requests)) - } - if client.requests[0].Think != thinkHigh { - t.Fatalf("first summary request think = %#v, want original", client.requests[0].Think) - } - if client.requests[1].Think == nil || client.requests[1].Think.Value != false { - t.Fatalf("fallback summary request think = %#v, want false", client.requests[1].Think) - } - if client.requests[2].Think != nil { - t.Fatalf("unsupported fallback retry think = %#v, want nil", client.requests[2].Think) - } -} - -func TestSimpleCompactorKeepsFewerTurnsForShortChats(t *testing.T) { - client := &fakeClient{ - responses: [][]api.ChatResponse{{ - {Message: api.Message{Role: "assistant", Content: "short summary"}}, - }}, - } - compactor := &SimpleCompactor{Client: client, Options: CompactionOptions{ - ContextWindowTokens: 16000, - KeepUserTurns: 3, - Threshold: 0.5, - }} - - result, err := compactor.MaybeCompact(context.Background(), CompactionRequest{ - ChatID: "chat-1", - Model: "model", - Messages: []api.Message{ - {Role: "user", Content: "old request"}, - {Role: "assistant", Content: "old answer"}, - {Role: "user", Content: "latest request"}, - }, - Latest: api.ChatResponse{Metrics: api.Metrics{PromptEvalCount: 12000}}, - }) - if err != nil { - t.Fatal(err) - } - if !result.Compacted { - t.Fatal("expected compaction") - } - if len(result.Messages) != 3 { - t.Fatalf("messages = %#v, want compaction tool pair plus latest request", result.Messages) - } - assertCompactionSummaryPair(t, result.Messages[:2]) - if result.Messages[2].Content != "latest request" { - t.Fatalf("latest turn was not kept: %#v", result.Messages) - } -} - -func TestSimpleCompactorCanArchiveWholeShortChat(t *testing.T) { - client := &fakeClient{ - responses: [][]api.ChatResponse{{ - {Message: api.Message{Role: "assistant", Content: "whole summary"}}, - }}, - } - compactor := &SimpleCompactor{Client: client, Options: CompactionOptions{ - ContextWindowTokens: 100, - KeepUserTurns: 3, - Threshold: 0.5, - }} - - result, err := compactor.MaybeCompact(context.Background(), CompactionRequest{ - ChatID: "chat-1", - Model: "model", - Messages: []api.Message{ - {Role: "user", Content: "only request"}, - {Role: "assistant", Content: "only answer"}, - }, - Latest: api.ChatResponse{Metrics: api.Metrics{PromptEvalCount: 75}}, - }) - if err != nil { - t.Fatal(err) - } - if !result.Compacted { - t.Fatal("expected compaction") - } - if len(result.Messages) != 2 { - t.Fatalf("messages = %#v, want only compaction tool pair", result.Messages) - } - assertCompactionSummaryPair(t, result.Messages) -} - -func TestSimpleCompactorSkipsBelowThreshold(t *testing.T) { - client := &fakeClient{} - compactor := &SimpleCompactor{Client: client, Options: CompactionOptions{ - ContextWindowTokens: 100, - Threshold: 0.8, - }} - - messages := []api.Message{ - {Role: "user", Content: "one"}, - {Role: "user", Content: "two"}, - {Role: "user", Content: "three"}, - {Role: "user", Content: "four"}, - {Role: "user", Content: "five"}, - } - result, err := compactor.MaybeCompact(context.Background(), CompactionRequest{ - Model: "model", - Messages: messages, - Latest: api.ChatResponse{Metrics: api.Metrics{PromptEvalCount: 50}}, - }) - if err != nil { - t.Fatal(err) - } - if result.Compacted { - t.Fatal("did not expect compaction") - } - if result.Due { - t.Fatal("below-threshold compaction should not be due") - } - if len(result.Messages) != len(messages) { - t.Fatalf("messages changed below threshold: %#v", result.Messages) - } - if len(client.requests) != 0 { - t.Fatalf("summary requests = %d, want 0", len(client.requests)) - } -} - -func TestSimpleCompactorUsesEstimatedMessagesWhenPromptEvalMissing(t *testing.T) { - client := &fakeClient{ - responses: [][]api.ChatResponse{{ - {Message: api.Message{Role: "assistant", Content: "estimated summary"}}, - }}, - } - compactor := &SimpleCompactor{Client: client, Options: CompactionOptions{ - ContextWindowTokens: 100, - KeepUserTurns: 1, - Threshold: 0.8, - }} - - result, err := compactor.MaybeCompact(context.Background(), CompactionRequest{ - Model: "model", - Messages: []api.Message{ - {Role: "user", Content: "old request"}, - {Role: "assistant", Content: "old answer"}, - {Role: "user", Content: "read large output"}, - {Role: "assistant", ToolCalls: []api.ToolCall{{ - ID: "call-1", - Function: api.ToolCallFunction{ - Name: "read", - }, - }}}, - {Role: "tool", ToolName: "read", ToolCallID: "call-1", Content: strings.Repeat("x", 360)}, - }, - }) - if err != nil { - t.Fatal(err) - } - if !result.Due || !result.Compacted { - t.Fatalf("expected estimate-driven compaction, got %#v", result) - } - if result.Summary != "estimated summary" { - t.Fatalf("summary = %q", result.Summary) - } -} - -func TestSimpleCompactorEstimateIncludesRequestPreamble(t *testing.T) { - compactor := &SimpleCompactor{Client: nil, Options: CompactionOptions{ - ContextWindowTokens: 100, - Threshold: 0.8, - }} - - if !compactor.shouldCompact(CompactionRequest{ - SystemPrompt: strings.Repeat("system ", 360), - Messages: []api.Message{{Role: "user", Content: "tiny"}}, - }) { - t.Fatal("system prompt should count toward compaction estimate") - } - - if !compactor.shouldCompact(CompactionRequest{ - Messages: []api.Message{{Role: "user", Content: "tiny"}}, - Tools: api.Tools{{ - Type: "function", - Function: api.ToolFunction{ - Name: "verbose_tool", - Description: strings.Repeat("description ", 360), - }, - }}, - }) { - t.Fatal("tool definitions should count toward compaction estimate") - } -} - -func TestCompactionPromptFitsBudgetByTruncatingLargeToolOutput(t *testing.T) { - largeToolOutput := strings.Repeat("x", 10_000) - body, err := compactionPrompt("", []api.Message{ - {Role: "user", Content: "what changed?"}, - {Role: "assistant", ToolCalls: []api.ToolCall{{ - ID: "call-1", - Function: api.ToolCallFunction{ - Name: "bash", - }, - }}}, - {Role: "tool", ToolName: "bash", ToolCallID: "call-1", Content: largeToolOutput}, - }, 300) - if err != nil { - t.Fatal(err) - } - if estimateCompactionTokens(body) > 300 { - t.Fatalf("compaction prompt tokens = %d, want <= 300", estimateCompactionTokens(body)) - } - if strings.Count(body, "x") >= len(largeToolOutput) { - t.Fatal("large tool output was not truncated") - } - if !strings.Contains(body, "[tool output truncated: showing first ~") { - t.Fatalf("truncation marker missing from compaction prompt: %q", body) - } -} - -func TestCompactionPromptRetruncatesAlreadyTruncatedToolOutput(t *testing.T) { - alreadyTruncated := strings.Repeat("x", 7000) + "\n\n[tool output truncated: showing first ~100 tokens and last ~100 tokens; omitted ~99999 tokens. Use a narrower command, line range, or search query if more detail is needed.]\n\n" + strings.Repeat("y", 7000) - body, err := compactionPrompt("", []api.Message{ - {Role: "user", Content: "what changed?"}, - {Role: "assistant", ToolCalls: []api.ToolCall{{ - ID: "call-1", - Function: api.ToolCallFunction{ - Name: "bash", - }, - }}}, - {Role: "tool", ToolName: "bash", ToolCallID: "call-1", Content: alreadyTruncated}, - }, 300) - if err != nil { - t.Fatal(err) - } - if estimateCompactionTokens(body) > 300 { - t.Fatalf("compaction prompt tokens = %d, want <= 300", estimateCompactionTokens(body)) - } - if strings.Count(body, "x")+strings.Count(body, "y") >= 14_000 { - t.Fatal("already-truncated tool output was not truncated again") - } - if !strings.Contains(body, "[tool output truncated: showing first ~") { - t.Fatalf("truncation marker missing from compaction prompt: %q", body) - } -} - -func TestCompactionSummaryTextStripsPrefix(t *testing.T) { - content := compactionSummaryMessageForTask("worked on branch changes", false) - if got := CompactionSummaryText(content); got != "worked on branch changes" { - t.Fatalf("summary text = %q", got) - } -} - -func TestCompactionSummaryCanTellModelToContinueTask(t *testing.T) { - content := compactionSummaryMessageForTask("worked on branch changes", true) - if !strings.Contains(content, CompactionContinueInstruction) { - t.Fatalf("summary message missing continue instruction: %q", content) - } - if got := CompactionSummaryText(content); got != "worked on branch changes" { - t.Fatalf("summary text = %q", got) - } -} - -func TestResolveContextWindowTokensPrefersExplicitNumCtx(t *testing.T) { - tests := []struct { - name string - options map[string]any - configured int - want int - }{ - { - name: "explicit smaller num ctx", - options: map[string]any{"num_ctx": 4096}, - configured: 8192, - want: 4096, - }, - { - name: "explicit num ctx can exceed configured metadata", - options: map[string]any{"num_ctx": 131072}, - configured: 8192, - want: 131072, - }, - { - name: "metadata without explicit num ctx", - configured: 32768, - want: 32768, - }, - } - - for _, tt := range tests { - t.Run(tt.name, func(t *testing.T) { - if got := ResolveContextWindowTokens(tt.options, tt.configured); got != tt.want { - t.Fatalf("ResolveContextWindowTokens() = %d, want %d", got, tt.want) - } - }) - } -} - -func TestSimpleCompactorForceCompactsWithoutPromptEvalCount(t *testing.T) { - client := &fakeClient{ - responses: [][]api.ChatResponse{{ - {Message: api.Message{Role: "assistant", Content: "forced summary"}}, - }}, - } - compactor := &SimpleCompactor{Client: client, Options: CompactionOptions{ - ContextWindowTokens: 100, - KeepUserTurns: 1, - Threshold: 0.8, - }} - - result, err := compactor.MaybeCompact(context.Background(), CompactionRequest{ - Model: "model", - Messages: []api.Message{ - {Role: "user", Content: "old"}, - {Role: "assistant", Content: "old answer"}, - {Role: "user", Content: "recent"}, - }, - Force: true, - }) - if err != nil { - t.Fatal(err) - } - if !result.Due || !result.Compacted { - t.Fatalf("forced compaction result = %#v", result) - } - if result.Summary != "forced summary" { - t.Fatalf("summary = %q", result.Summary) - } -} - -func TestSimpleCompactorDefaultsToKeepingThreeUserTurns(t *testing.T) { - client := &fakeClient{ - responses: [][]api.ChatResponse{{ - {Message: api.Message{Role: "assistant", Content: "summary"}}, - }}, - } - compactor := &SimpleCompactor{Client: client, Options: CompactionOptions{ - ContextWindowTokens: 16000, - Threshold: 0.5, - }} - - result, err := compactor.MaybeCompact(context.Background(), CompactionRequest{ - ChatID: "chat-1", - Model: "model", - Messages: []api.Message{ - {Role: "user", Content: "old"}, - {Role: "assistant", Content: "old answer"}, - {Role: "user", Content: "one"}, - {Role: "assistant", Content: "one answer"}, - {Role: "user", Content: "two"}, - {Role: "assistant", Content: "two answer"}, - {Role: "user", Content: "three"}, - }, - Latest: api.ChatResponse{Metrics: api.Metrics{PromptEvalCount: 12000}}, - }) - if err != nil { - t.Fatal(err) - } - if !result.Compacted { - t.Fatal("expected compaction") - } - assertCompactionSummaryPair(t, result.Messages[:2]) - if got := result.Messages[2].Content; got != "one" { - t.Fatalf("first kept turn = %q, want one", got) - } -} - -func TestSimpleCompactorCarriesPreviousSummary(t *testing.T) { - client := &fakeClient{ - responses: [][]api.ChatResponse{{ - {Message: api.Message{Role: "assistant", Content: "new summary"}}, - }}, - } - compactor := &SimpleCompactor{Client: client, Options: CompactionOptions{ - ContextWindowTokens: 16000, - KeepUserTurns: 1, - Threshold: 0.5, - }} - - result, err := compactor.MaybeCompact(context.Background(), CompactionRequest{ - Model: "model", - Messages: []api.Message{ - {Role: "system", Content: CompactionSummaryMessagePrefix + "old summary"}, - {Role: "user", Content: "old"}, - {Role: "assistant", Content: "old answer"}, - {Role: "user", Content: "recent"}, - }, - Latest: api.ChatResponse{Metrics: api.Metrics{PromptEvalCount: 12000}}, - }) - if err != nil { - t.Fatal(err) - } - if !result.Compacted { - t.Fatal("expected compaction") - } - if !strings.Contains(client.requests[0].Messages[1].Content, "Previous summary:\nold summary") { - t.Fatalf("previous summary missing from request: %q", client.requests[0].Messages[1].Content) - } -} - -func TestSimpleCompactorCarriesPreviousToolSummaryAndPlacesNewSummaryBeforeKeptSuffix(t *testing.T) { - client := &fakeClient{ - responses: [][]api.ChatResponse{{ - {Message: api.Message{Role: "assistant", Content: "new summary"}}, - }}, - } - compactor := &SimpleCompactor{Client: client, Options: CompactionOptions{ - ContextWindowTokens: 16000, - KeepUserTurns: 1, - Threshold: 0.5, - }} - - messages := []api.Message{ - {Role: "user", Content: "kept before old summary"}, - CompactionSummaryMessages("old summary", false)[0], - CompactionSummaryMessages("old summary", false)[1], - {Role: "user", Content: "latest request"}, - } - result, err := compactor.MaybeCompact(context.Background(), CompactionRequest{ - Model: "model", - Messages: messages, - Latest: api.ChatResponse{Metrics: api.Metrics{PromptEvalCount: 12000}}, - }) - if err != nil { - t.Fatal(err) - } - if !result.Compacted { - t.Fatal("expected compaction") - } - if !strings.Contains(client.requests[0].Messages[1].Content, "Previous summary:\nold summary") { - t.Fatalf("previous summary missing from request: %q", client.requests[0].Messages[1].Content) - } - if len(result.Messages) != 3 { - t.Fatalf("messages = %#v, want compaction pair plus latest request", result.Messages) - } - assertCompactionSummaryPair(t, result.Messages[:2]) - if result.Messages[2].Content != "latest request" { - t.Fatalf("kept suffix = %#v", result.Messages) - } -} diff --git a/agent/events.go b/agent/events.go deleted file mode 100644 index 6d2c4a881eb..00000000000 --- a/agent/events.go +++ /dev/null @@ -1,177 +0,0 @@ -package agent - -import ( - "context" - "errors" - - "github.com/ollama/ollama/api" -) - -type EventType string - -const ( - EventMessageDelta EventType = "message_delta" - EventThinkingDelta EventType = "thinking_delta" - EventToolCallDetected EventType = "tool_call_detected" - EventToolStarted EventType = "tool_started" - EventToolFinished EventType = "tool_finished" - EventCompactionStarted EventType = "compaction_started" - EventCompactionProgress EventType = "compaction_progress" - EventCompacted EventType = "compacted" - EventCompactionSkipped EventType = "compaction_skipped" - EventRunFinished EventType = "run_finished" - EventError EventType = "error" -) - -// ToolStatus is the typed lifecycle state for a tool call, carried on -// Event.ToolStatus for tool events. -type ToolStatus string - -const ( - ToolStatusRunning ToolStatus = "running" - ToolStatusDone ToolStatus = "done" - ToolStatusFailed ToolStatus = "failed" - ToolStatusDenied ToolStatus = "denied" - ToolStatusDisabled ToolStatus = "disabled" - ToolStatusSkipped ToolStatus = "skipped" -) - -// RunStatus is the typed terminal outcome of a run, carried on Event.Status for -// run_finished events. -type RunStatus string - -const ( - RunStatusDone RunStatus = "done" - RunStatusDenied RunStatus = "denied" - RunStatusCanceled RunStatus = "canceled" -) - -// CompactionTrigger is the typed reason a compaction ran or was attempted, -// carried on Event.CompactionTrigger for compaction events. -type CompactionTrigger string - -const ( - CompactionTriggerForce CompactionTrigger = "force" - CompactionTriggerPromptEval CompactionTrigger = "prompt_eval" - CompactionTriggerEstimate CompactionTrigger = "estimate" - CompactionTriggerToolOutput CompactionTrigger = "tool_output" - CompactionTriggerError CompactionTrigger = "error" - CompactionTriggerDue CompactionTrigger = "due" -) - -type Event struct { - Type EventType `json:"type"` - RunID string `json:"runId,omitempty"` - ChatID string `json:"chatId,omitempty"` - Model string `json:"model,omitempty"` - Status RunStatus `json:"status,omitempty"` - ToolStatus ToolStatus `json:"toolStatus,omitempty"` - CompactionTrigger CompactionTrigger `json:"compactionTrigger,omitempty"` - ToolCallID string `json:"toolCallId,omitempty"` - ToolName string `json:"toolName,omitempty"` - WorkingDir string `json:"workingDir,omitempty"` - Content string `json:"content,omitempty"` - Thinking string `json:"thinking,omitempty"` - ToolCalls []api.ToolCall `json:"toolCalls,omitempty"` - Messages []api.Message `json:"messages,omitempty"` - Args map[string]any `json:"args,omitempty"` - Tokens int `json:"tokens,omitempty"` - Error string `json:"error,omitempty"` -} - -type EventSink interface { - Emit(Event) error -} - -type EventSinkFunc func(Event) error - -func (fn EventSinkFunc) Emit(event Event) error { - if fn == nil { - return nil - } - return fn(event) -} - -// eventMetadata carries the run identification fields shared by all events. -type eventMetadata struct { - runID string - chatID string - model string -} - -func newEventMetadata(runID string, opts RunOptions) eventMetadata { - return eventMetadata{runID: runID, chatID: opts.ChatID, model: opts.Model} -} - -func newMessageDelta(m eventMetadata, content string) Event { - return Event{Type: EventMessageDelta, RunID: m.runID, ChatID: m.chatID, Model: m.model, Content: content} -} - -func newThinkingDelta(m eventMetadata, thinking string) Event { - return Event{Type: EventThinkingDelta, RunID: m.runID, ChatID: m.chatID, Model: m.model, Thinking: thinking} -} - -func newToolCallDetected(m eventMetadata, calls []api.ToolCall) Event { - return Event{Type: EventToolCallDetected, RunID: m.runID, ChatID: m.chatID, Model: m.model, ToolCalls: calls} -} - -func newToolStarted(m eventMetadata, callID, toolName, workingDir string, args map[string]any) Event { - return Event{Type: EventToolStarted, RunID: m.runID, ChatID: m.chatID, Model: m.model, ToolStatus: ToolStatusRunning, ToolCallID: callID, ToolName: toolName, WorkingDir: workingDir, Args: args} -} - -func newToolFinished(m eventMetadata, status ToolStatus, callID, toolName, workingDir string, args map[string]any, content, errMsg string) Event { - ev := Event{Type: EventToolFinished, RunID: m.runID, ChatID: m.chatID, Model: m.model, ToolStatus: status, ToolCallID: callID, ToolName: toolName, WorkingDir: workingDir, Args: args, Content: content} - if errMsg != "" { - ev.Error = errMsg - } - return ev -} - -func newRunFinished(m eventMetadata, status RunStatus) Event { - return Event{Type: EventRunFinished, RunID: m.runID, ChatID: m.chatID, Model: m.model, Status: status} -} - -func newErrorEvent(m eventMetadata, errMsg string) Event { - return Event{Type: EventError, RunID: m.runID, ChatID: m.chatID, Model: m.model, Error: errMsg} -} - -func newCompactionProgress(m eventMetadata, tokens int) Event { - return Event{Type: EventCompactionProgress, RunID: m.runID, ChatID: m.chatID, Model: m.model, Tokens: tokens} -} - -func newCompactionStarted(m eventMetadata, trigger CompactionTrigger) Event { - return Event{Type: EventCompactionStarted, RunID: m.runID, ChatID: m.chatID, Model: m.model, CompactionTrigger: trigger} -} - -func newCompactionSkipped(m eventMetadata, trigger CompactionTrigger, content string) Event { - return Event{Type: EventCompactionSkipped, RunID: m.runID, ChatID: m.chatID, Model: m.model, CompactionTrigger: trigger, Content: content} -} - -func newCompacted(m eventMetadata, messages []api.Message, trigger CompactionTrigger, content string) Event { - return Event{Type: EventCompacted, RunID: m.runID, ChatID: m.chatID, Model: m.model, CompactionTrigger: trigger, Content: content, Messages: messages} -} - -func (s *Session) emit(event Event) error { - if s == nil { - return nil - } - var errs []error - for _, sink := range s.EventSinks { - if sink == nil { - continue - } - if err := sink.Emit(event); err != nil { - errs = append(errs, err) - } - } - return errors.Join(errs...) -} - -func (s *Session) emitIgnoringCanceled(ctx context.Context, event Event) error { - err := s.emit(event) - if err != nil && ctx != nil && ctx.Err() != nil { - //nolint:nilerr // Event sinks may close during cancellation; cancellation is not a user-facing emit failure. - return nil - } - return err -} diff --git a/agent/registry.go b/agent/registry.go deleted file mode 100644 index 377efa8141e..00000000000 --- a/agent/registry.go +++ /dev/null @@ -1,104 +0,0 @@ -package agent - -import ( - "context" - "fmt" - "sort" - - "github.com/ollama/ollama/api" -) - -type ToolContext struct { - WorkingDir string -} - -type ToolResult struct { - Content string - WorkingDir string -} - -type Tool interface { - Name() string - Description() string - Schema() api.ToolFunction - Execute(context.Context, ToolContext, map[string]any) (ToolResult, error) -} - -type ApprovalRequired interface { - RequiresApproval(map[string]any) bool -} - -// ScopedTool is implemented by tools that need per-invocation approval -// scoping beyond the tool name (e.g. shell commands scoped to the exact -// command string). Tools that don't implement this are scoped by name only. -type ScopedTool interface { - ApprovalScope(args map[string]any) string -} - -type Registry struct { - tools map[string]Tool -} - -func (r *Registry) Register(tool Tool) { - if r == nil || tool == nil { - return - } - if r.tools == nil { - r.tools = make(map[string]Tool) - } - r.tools[tool.Name()] = tool -} - -func (r *Registry) Get(name string) (Tool, bool) { - if r == nil { - return nil, false - } - tool, ok := r.tools[name] - return tool, ok -} - -func (r *Registry) Names() []string { - if r == nil { - return nil - } - names := make([]string, 0, len(r.tools)) - for name := range r.tools { - names = append(names, name) - } - sort.Strings(names) - return names -} - -func (r *Registry) Tools() api.Tools { - if r == nil { - return nil - } - names := r.Names() - apiTools := make(api.Tools, 0, len(names)) - for _, name := range names { - tool := r.tools[name] - apiTools = append(apiTools, api.Tool{ - Type: "function", - Function: tool.Schema(), - }) - } - return apiTools -} - -func (r *Registry) Execute(ctx context.Context, toolCtx ToolContext, call api.ToolCall) (ToolResult, error) { - tool, ok := r.Get(call.Function.Name) - if !ok { - return ToolResult{}, fmt.Errorf("unknown tool: %s", call.Function.Name) - } - return tool.Execute(ctx, toolCtx, call.Function.Arguments.ToMap()) -} - -func ToolRequiresApproval(tool Tool, args map[string]any) bool { - if tool == nil { - return false - } - if t, ok := tool.(ApprovalRequired); ok { - return t.RequiresApproval(args) - } - return false -} diff --git a/agent/session.go b/agent/session.go deleted file mode 100644 index 46969c8f931..00000000000 --- a/agent/session.go +++ /dev/null @@ -1,1092 +0,0 @@ -package agent - -import ( - "context" - "encoding/json" - "errors" - "fmt" - "os" - "path/filepath" - "strings" - - "github.com/google/uuid" - - "github.com/ollama/ollama/api" - "github.com/ollama/ollama/internal/modelref" -) - -type ChatClient interface { - Chat(context.Context, *api.ChatRequest, api.ChatResponseFunc) error -} - -type Session struct { - Client ChatClient - EventSinks []EventSink - Tools *Registry - Skills *SkillCatalog - DisableTools bool - ApprovalPrompter ApprovalPrompter - ApprovalState *ApprovalState - WorkingDir string - Compactor Compactor -} - -type RunOptions struct { - ChatID string - Model string - SystemPrompt string - Messages []api.Message - NewMessages []api.Message - Format string - Options map[string]any - Think *api.ThinkValue - KeepAlive *api.Duration - // SkillName loads a catalog skill as an ordered synthetic tool call/result - // before the first model request for this run. - SkillName string - // MaxToolRounds limits consecutive model/tool cycles. A positive value is - // an explicit limit. Zero selects the model-specific default: local models - // use the default guard and cloud models are unlimited. A negative value - // disables the guard for tests or special callers. - MaxToolRounds int -} - -type RunResult struct { - Messages []api.Message - Latest api.ChatResponse - WorkingDir string -} - -const ( - defaultMaxToolRounds = 100 - maxToolResultRunes = 60000 - smallContextToolResultRunes = 6000 - tinyContextToolResultRunes = 3200 - smallContextToolResultTokenWindow = 8192 - tinyContextToolResultTokenWindow = 4096 - toolTruncationMarkerReserveTokens = 64 - toolOutputFullOmissionPrefix = "[tool output truncated: output omitted because the context is full;" -) - -type toolOutputOverflow struct { - toolName string - toolCallID string - content string -} - -type toolBatchResult struct { - messages []api.Message - stop toolExecutionStop - overflows []toolOutputOverflow -} - -// toolExecutionStop is the batch-level outcome for a group of tool calls, -// distinct from per-call Event.Status values. The values overlap with -// runFinish.status ("denied", "canceled") because a denied or canceled -// batch also terminates the run with the matching status. -type toolExecutionStop string - -const ( - toolExecutionDenied toolExecutionStop = "denied" - toolExecutionCanceled toolExecutionStop = "canceled" -) - -const toolExecutionDisabledMessage = "Tool execution disabled." - -type runPhase int - -const ( - runPhaseModel runPhase = iota - runPhaseTools - runPhaseCompact - runPhaseDone -) - -type runState struct { - runID string - opts RunOptions - - phase runPhase - - messages []api.Message - latest api.ChatResponse - - assistant api.Message - pendingToolCalls []api.ToolCall - canceled bool - - toolBatch *toolBatchResult - - consecutiveModelErrors int - toolRounds int - maxToolRounds int - compactionSkipNotified bool - - finish runFinish -} - -type runFinish struct { - status RunStatus - ignoreCanceled bool - err error -} - -func (st *runState) finishDone() { - st.finish = runFinish{status: RunStatusDone} - st.phase = runPhaseDone -} - -func (st *runState) finishDenied() { - st.finish = runFinish{status: RunStatusDenied} - st.phase = runPhaseDone -} - -func (st *runState) finishCanceled() { - st.finish = runFinish{status: RunStatusCanceled, ignoreCanceled: true} - st.phase = runPhaseDone -} - -func (st *runState) finishError(err error) { - st.finish = runFinish{err: err} - st.phase = runPhaseDone -} - -func (s *Session) Run(ctx context.Context, opts RunOptions) (*RunResult, error) { - if err := s.validateRun(opts); err != nil { - return nil, err - } - if s.ApprovalState == nil { - s.ApprovalState = &ApprovalState{} - } - runID := uuid.NewString() - messages, err := s.buildRunMessages(ctx, runID, opts) - if err != nil { - return nil, err - } - activatedSkill, err := s.activateSkill(ctx, runID, opts) - if err != nil { - s.emit(newErrorEvent(newEventMetadata(runID, opts), err.Error())) - return nil, err - } - if len(activatedSkill) > 0 { - messages = append(messages, activatedSkill...) - if err := s.checkPreflightPromptBudget(opts, messages); err != nil { - s.emit(newErrorEvent(newEventMetadata(runID, opts), err.Error())) - return nil, err - } - } - - st := runState{ - runID: runID, - opts: opts, - phase: runPhaseModel, - messages: messages, - maxToolRounds: resolvedMaxToolRounds(opts.Model, opts.MaxToolRounds), - } - for { - switch st.phase { - case runPhaseModel: - if err := s.runModelStep(ctx, &st); err != nil { - return nil, err - } - case runPhaseTools: - if err := s.runToolStep(ctx, &st); err != nil { - return nil, err - } - case runPhaseCompact: - if err := s.runCompactionStep(ctx, &st); err != nil { - return nil, err - } - case runPhaseDone: - return s.finishRun(ctx, &st) - } - } -} - -// validateRun checks the preconditions for a run. -func (s *Session) validateRun(opts RunOptions) error { - if s == nil { - return errors.New("nil session") - } - if s.Client == nil { - return errors.New("agent session requires a chat client") - } - if opts.Model == "" { - return errors.New("agent session requires a model") - } - return nil -} - -// buildRunMessages sanitizes the provided message history, runs the preflight -// prompt-budget check, and returns the initial message list for the run. It -// emits an EventError and returns it if the preflight check fails. -func (s *Session) buildRunMessages(ctx context.Context, runID string, opts RunOptions) ([]api.Message, error) { - messages := make([]api.Message, 0, len(opts.Messages)+len(opts.NewMessages)) - for _, msg := range opts.Messages { - messages = append(messages, sanitizeMessageForRun(msg)) - } - for _, msg := range opts.NewMessages { - msg = sanitizeMessageForRun(msg) - messages = append(messages, msg) - } - - if err := s.checkPreflightPromptBudget(opts, messages); err != nil { - s.emit(newErrorEvent(newEventMetadata(runID, opts), err.Error())) - return nil, err - } - return messages, nil -} - -func (s *Session) runModelStep(ctx context.Context, st *runState) error { - opts := st.opts - meta := newEventMetadata(st.runID, opts) - - assistant, pendingToolCalls, canceled, err := s.chatRound(ctx, st.runID, opts, st.messages, &st.latest) - if err != nil { - var statusErr api.StatusError - if errors.As(err, &statusErr) && statusErr.StatusCode >= 500 && st.consecutiveModelErrors < 2 { - st.consecutiveModelErrors++ - st.messages = append(st.messages, api.Message{ - Role: "user", - Content: fmt.Sprintf("Your previous response caused an error: %s\n\nPlease try again with a valid response.", statusErr.ErrorMessage), - }) - return nil - } - s.emit(newErrorEvent(meta, err.Error())) - return err - } - st.consecutiveModelErrors = 0 - st.assistant = assistant - st.pendingToolCalls = pendingToolCalls - st.canceled = canceled - - if !messageEmpty(assistant) { - st.messages = append(st.messages, assistant) - } - - if len(pendingToolCalls) == 0 { - st.toolBatch = nil - st.phase = runPhaseCompact - return nil - } - - if canceled { - skipped, skipErr := s.skipToolCalls(ctx, st.runID, opts, pendingToolCalls, "Tool execution skipped because the run was canceled.") - if skipErr != nil { - s.emit(newErrorEvent(meta, skipErr.Error())) - return skipErr - } - st.messages = append(st.messages, skipped...) - st.finishCanceled() - return nil - } - - if s.DisableTools { - batch, skipErr := s.disabledToolCalls(ctx, st.runID, opts, st.messages, pendingToolCalls) - if skipErr != nil { - s.emit(newErrorEvent(meta, skipErr.Error())) - return skipErr - } - st.messages = append(st.messages, batch.messages...) - st.toolBatch = &batch - st.phase = runPhaseCompact - return nil - } - - if s.Tools == nil { - st.finishDone() - return nil - } - - if st.maxToolRounds >= 0 && st.toolRounds >= st.maxToolRounds { - content := fmt.Sprintf("Tool execution skipped because the max tool-round limit of %d was reached. Send another message to continue.", st.maxToolRounds) - toolMessages, skipErr := s.skipToolCalls(ctx, st.runID, opts, pendingToolCalls, content) - if skipErr != nil { - s.emit(newErrorEvent(meta, skipErr.Error())) - return skipErr - } - st.messages = append(st.messages, toolMessages...) - err := fmt.Errorf("tool round limit reached after %d rounds; send another message to continue", st.maxToolRounds) - s.emit(newErrorEvent(meta, err.Error())) - st.finishError(err) - return nil - } - - st.phase = runPhaseTools - return nil -} - -func (s *Session) runToolStep(ctx context.Context, st *runState) error { - batch, err := s.executeToolCalls(ctx, st.runID, st.opts, st.messages, st.pendingToolCalls) - if err != nil { - s.emit(newErrorEvent(newEventMetadata(st.runID, st.opts), err.Error())) - return err - } - - st.messages = append(st.messages, batch.messages...) - st.toolBatch = &batch - st.phase = runPhaseCompact - return nil -} - -func (s *Session) runCompactionStep(ctx context.Context, st *runState) error { - opts := st.opts - meta := newEventMetadata(st.runID, opts) - var err error - if st.toolBatch != nil && len(st.toolBatch.overflows) > 0 { - st.messages, st.compactionSkipNotified, err = s.compactForToolOutputOverflow(ctx, st.runID, opts, st.messages, st.latest, st.assistant, st.toolBatch.messages, st.toolBatch.overflows, st.compactionSkipNotified) - } else { - st.messages, st.compactionSkipNotified, err = s.maybeCompact(ctx, st.runID, opts, st.messages, st.latest, st.compactionSkipNotified) - } - if err != nil { - s.emit(newErrorEvent(meta, err.Error())) - st.finishError(err) - return nil - } - - if st.toolBatch == nil { - if st.canceled { - st.finishCanceled() - } else { - st.finishDone() - } - return nil - } - - switch st.toolBatch.stop { - case toolExecutionDenied: - st.finishDenied() - case toolExecutionCanceled: - st.finishCanceled() - default: - st.toolRounds++ - st.assistant = api.Message{} - st.pendingToolCalls = nil - st.toolBatch = nil - st.phase = runPhaseModel - } - return nil -} - -func (s *Session) finishRun(ctx context.Context, st *runState) (*RunResult, error) { - if st.finish.status != "" { - event := newRunFinished(newEventMetadata(st.runID, st.opts), st.finish.status) - var err error - if st.finish.ignoreCanceled { - err = s.emitIgnoringCanceled(ctx, event) - } else { - err = s.emit(event) - } - if err != nil { - return nil, err - } - } - return &RunResult{Messages: st.messages, Latest: st.latest, WorkingDir: s.WorkingDir}, st.finish.err -} - -func (s *Session) chatRound(ctx context.Context, runID string, opts RunOptions, messages []api.Message, latest *api.ChatResponse) (api.Message, []api.ToolCall, bool, error) { - meta := newEventMetadata(runID, opts) - var tools api.Tools - if !s.DisableTools { - tools = s.availableTools() - } - req := buildChatRequest(opts, messages, tools) - - assistant := api.Message{Role: "assistant"} - var pendingToolCalls []api.ToolCall - - err := s.Client.Chat(ctx, &req, func(response api.ChatResponse) error { - if response.Message.Role != "" { - assistant.Role = response.Message.Role - } - - if messageEmpty(response.Message) { - *latest = response - return nil - } - - if response.Message.Thinking != "" { - assistant.Thinking += response.Message.Thinking - if err := s.emit(newThinkingDelta(meta, response.Message.Thinking)); err != nil { - return err - } - } - - if response.Message.Content != "" { - assistant.Content += response.Message.Content - if err := s.emit(newMessageDelta(meta, response.Message.Content)); err != nil { - return err - } - } - - if len(response.Message.ToolCalls) > 0 { - assistant.ToolCalls = append(assistant.ToolCalls, response.Message.ToolCalls...) - pendingToolCalls = append(pendingToolCalls, response.Message.ToolCalls...) - if err := s.emit(newToolCallDetected(meta, response.Message.ToolCalls)); err != nil { - return err - } - } - - *latest = response - return nil - }) - if err != nil { - if isContextCanceledError(ctx, err) { - return assistant, pendingToolCalls, true, nil - } - return assistant, pendingToolCalls, false, err - } - - return assistant, pendingToolCalls, false, nil -} - -func buildChatRequest(opts RunOptions, messages []api.Message, tools api.Tools) api.ChatRequest { - requestMessages := sanitizeMessagesForRequest(messages) - if strings.TrimSpace(opts.SystemPrompt) != "" { - withSystem := make([]api.Message, 0, len(requestMessages)+1) - withSystem = append(withSystem, api.Message{Role: "system", Content: opts.SystemPrompt}) - requestMessages = append(withSystem, requestMessages...) - } - - format := opts.Format - if format == "json" { - format = `"` + format + `"` - } - - req := api.ChatRequest{ - Model: opts.Model, - Messages: requestMessages, - Format: json.RawMessage(format), - Options: opts.Options, - Think: opts.Think, - } - if opts.KeepAlive != nil { - req.KeepAlive = opts.KeepAlive - } - if len(tools) > 0 { - req.Tools = tools - } - return req -} - -func (s *Session) executeToolCalls(ctx context.Context, runID string, opts RunOptions, messages []api.Message, calls []api.ToolCall) (toolBatchResult, error) { - meta := newEventMetadata(runID, opts) - batch := toolBatchResult{ - messages: make([]api.Message, 0, len(calls)), - } - // Pre-compute the full-history token estimate once per batch instead of - // re-marshaling the entire history for each tool call. Per-call deltas - // (tool messages already appended this batch) are tracked in batchTokens - // and added to historyTokens for a lightweight running total. - historyTokens := s.estimateRunPromptTokens(opts, messages) - batchTokens := 0 - - type plannedToolCall struct { - call api.ToolCall - tool Tool - toolName string - args map[string]any - workingDir string - } - plans := make([]plannedToolCall, 0, len(calls)) - batchWorkingDir := s.currentWorkingDir() - approvalReq := ApprovalRequest{WorkingDir: batchWorkingDir} - for _, call := range calls { - toolName := call.Function.Name - args := call.Function.Arguments.ToMap() - tool, ok := s.Tools.Get(toolName) - plans = append(plans, plannedToolCall{ - call: call, - tool: tool, - toolName: toolName, - args: args, - workingDir: batchWorkingDir, - }) - if ok && s.needsApproval(tool, toolName, args) { - approvalReq.AddToolCall(call.ID, toolName, toolApprovalScope(tool, toolName, args), args) - } - } - - if len(approvalReq.Calls) > 0 { - approvalResult, err := s.authorizeToolCalls(ctx, approvalReq) - if err != nil { - if ctx.Err() != nil { - skipped, skipErr := s.skipToolCalls(ctx, runID, opts, calls, "Tool execution skipped because the run was canceled.") - if skipErr != nil { - return toolBatchResult{}, skipErr - } - batch.messages = append(batch.messages, skipped...) - batch.stop = toolExecutionCanceled - return batch, nil - } - return toolBatchResult{}, err - } - if !approvalResult.Allow { - content := approvalResult.Reason - if content == "" { - content = "Tool execution denied." - } - for _, plan := range plans { - msg := s.toolMessageForContext(plan.toolName, plan.call.ID, content, opts, historyTokens+batchTokens) - batch.messages = append(batch.messages, msg) - batchTokens += estimateMessagesTokens([]api.Message{msg}) - deniedContent := msg.Content - if emitErr := s.emit(newToolFinished(meta, "denied", plan.call.ID, plan.toolName, "", plan.args, deniedContent, deniedContent)); emitErr != nil { - return toolBatchResult{}, emitErr - } - } - batch.stop = toolExecutionDenied - return batch, nil - } - } - - for i, plan := range plans { - call := plan.call - toolName := plan.toolName - args := plan.args - if ctx.Err() != nil { - skipped, skipErr := s.skipToolCalls(ctx, runID, opts, calls[i:], "Tool execution skipped because the run was canceled.") - if skipErr != nil { - return toolBatchResult{}, skipErr - } - batch.messages = append(batch.messages, skipped...) - batch.stop = toolExecutionCanceled - return batch, nil - } - if plan.tool == nil { - content := fmt.Sprintf("Error: unknown tool: %s", toolName) - msg := s.toolMessageForContext(toolName, call.ID, content, opts, historyTokens+batchTokens) - batch.messages = append(batch.messages, msg) - batchTokens += estimateMessagesTokens([]api.Message{msg}) - content = msg.Content - if toolOutputFullyOmitted(content) { - batch.overflows = append(batch.overflows, toolOutputOverflow{toolName: toolName, toolCallID: call.ID, content: fmt.Sprintf("Error: unknown tool: %s", toolName)}) - } - if emitErr := s.emit(newToolFinished(meta, "failed", call.ID, toolName, "", args, content, fmt.Sprintf("unknown tool: %s", toolName))); emitErr != nil { - return toolBatchResult{}, emitErr - } - continue - } - - if err := s.emit(newToolStarted(meta, call.ID, toolName, plan.workingDir, args)); err != nil { - return toolBatchResult{}, err - } - - result, err := s.Tools.Execute(ctx, ToolContext{WorkingDir: plan.workingDir}, call) - if err != nil { - rawContent := fmt.Sprintf("Error: %v", err) - msg := s.toolMessageForContext(toolName, call.ID, rawContent, opts, historyTokens+batchTokens) - batch.messages = append(batch.messages, msg) - batchTokens += estimateMessagesTokens([]api.Message{msg}) - content := msg.Content - if toolOutputFullyOmitted(content) { - batch.overflows = append(batch.overflows, toolOutputOverflow{toolName: toolName, toolCallID: call.ID, content: rawContent}) - } - if emitErr := s.emitIgnoringCanceled(ctx, newToolFinished(meta, "failed", call.ID, toolName, "", args, content, err.Error())); emitErr != nil { - return toolBatchResult{}, emitErr - } - if ctx.Err() != nil { - skipped, skipErr := s.skipToolCalls(ctx, runID, opts, calls[i+1:], "Tool execution skipped because the run was canceled.") - if skipErr != nil { - return toolBatchResult{}, skipErr - } - batch.messages = append(batch.messages, skipped...) - batch.stop = toolExecutionCanceled - return batch, nil - } - continue - } - - eventWorkingDir := plan.workingDir - if s.applyToolWorkingDir(result.WorkingDir) { - eventWorkingDir = s.WorkingDir - } - rawContent := result.Content - - msg := s.toolMessageForContext(toolName, call.ID, rawContent, opts, historyTokens+batchTokens) - batch.messages = append(batch.messages, msg) - batchTokens += estimateMessagesTokens([]api.Message{msg}) - content := msg.Content - - if toolOutputFullyOmitted(content) { - batch.overflows = append(batch.overflows, toolOutputOverflow{toolName: toolName, toolCallID: call.ID, content: rawContent}) - } - if err := s.emitIgnoringCanceled(ctx, newToolFinished(meta, "done", call.ID, toolName, eventWorkingDir, args, content, "")); err != nil { - return toolBatchResult{}, err - } - if ctx.Err() != nil { - skipped, skipErr := s.skipToolCalls(ctx, runID, opts, calls[i+1:], "Tool execution skipped because the run was canceled.") - if skipErr != nil { - return toolBatchResult{}, skipErr - } - batch.messages = append(batch.messages, skipped...) - batch.stop = toolExecutionCanceled - return batch, nil - } - } - return batch, nil -} - -func (s *Session) disabledToolCalls(ctx context.Context, runID string, opts RunOptions, messages []api.Message, calls []api.ToolCall) (toolBatchResult, error) { - meta := newEventMetadata(runID, opts) - batch := toolBatchResult{ - messages: make([]api.Message, 0, len(calls)), - } - historyTokens := s.estimateRunPromptTokens(opts, messages) - batchTokens := 0 - for _, call := range calls { - toolName := call.Function.Name - args := call.Function.Arguments.ToMap() - msg := s.toolMessageForContext(toolName, call.ID, toolExecutionDisabledMessage, opts, historyTokens+batchTokens) - batch.messages = append(batch.messages, msg) - batchTokens += estimateMessagesTokens([]api.Message{msg}) - if emitErr := s.emitIgnoringCanceled(ctx, newToolFinished(meta, "disabled", call.ID, toolName, "", args, msg.Content, msg.Content)); emitErr != nil { - return toolBatchResult{}, emitErr - } - } - return batch, nil -} - -func (s *Session) skipToolCalls(ctx context.Context, runID string, opts RunOptions, calls []api.ToolCall, content string) ([]api.Message, error) { - meta := newEventMetadata(runID, opts) - toolMessages := make([]api.Message, 0, len(calls)) - for _, call := range calls { - toolName := call.Function.Name - args := call.Function.Arguments.ToMap() - msg := toolMessage(toolName, call.ID, content) - toolMessages = append(toolMessages, msg) - if emitErr := s.emitIgnoringCanceled(ctx, newToolFinished(meta, "skipped", call.ID, toolName, "", args, msg.Content, msg.Content)); emitErr != nil { - return nil, emitErr - } - } - return toolMessages, nil -} - -func (s *Session) currentWorkingDir() string { - if s.WorkingDir != "" { - return s.WorkingDir - } - wd, err := os.Getwd() - if err != nil { - return "" - } - s.WorkingDir = wd - return s.WorkingDir -} - -func (s *Session) applyToolWorkingDir(next string) bool { - next = strings.TrimSpace(next) - if next == "" { - return false - } - current := s.currentWorkingDir() - nextAbs, err := canonicalSessionPath(next) - if err != nil { - return false - } - if current == nextAbs { - return false - } - s.WorkingDir = nextAbs - return true -} - -func canonicalSessionPath(path string) (string, error) { - abs, err := filepath.Abs(path) - if err != nil { - return "", err - } - resolved, err := filepath.EvalSymlinks(abs) - if err == nil { - return resolved, nil - } - return abs, nil -} - -func isContextCanceledError(ctx context.Context, err error) bool { - if err == nil { - return false - } - if errors.Is(err, context.Canceled) { - return true - } - return ctx != nil && errors.Is(ctx.Err(), context.Canceled) && strings.Contains(err.Error(), "context canceled") -} - -func (s *Session) maybeCompact(ctx context.Context, runID string, opts RunOptions, messages []api.Message, latest api.ChatResponse, skipNotified bool) ([]api.Message, bool, error) { - if s.Compactor == nil { - return messages, skipNotified, nil - } - req := s.compactionRequest(runID, opts, messages, latest) - trigger := s.autoCompactionTrigger(req) - if trigger != "" { - s.emitCompactionStarted(runID, opts, trigger) - } - result, err := s.Compactor.MaybeCompact(ctx, req) - if err != nil { - if result.Due && !skipNotified { - if trigger == "" { - trigger = CompactionTriggerError - } - s.emitCompactionSkipped(runID, opts, trigger, result.Reason) - skipNotified = true - } - return messages, skipNotified, nil - } - if !result.Compacted { - if result.Due && !skipNotified { - if trigger == "" { - trigger = CompactionTriggerDue - } - s.emitCompactionSkipped(runID, opts, trigger, result.Reason) - skipNotified = true - } - return messages, skipNotified, nil - } - s.emitCompacted(runID, opts, result.Messages, trigger, result.Summary) - if err := s.checkPostCompactionPromptBudget(opts, result.Messages); err != nil { - return result.Messages, skipNotified, err - } - return result.Messages, skipNotified, nil -} - -func (s *Session) compactForToolOutputOverflow(ctx context.Context, runID string, opts RunOptions, messages []api.Message, latest api.ChatResponse, assistant api.Message, toolMessages []api.Message, overflows []toolOutputOverflow, skipNotified bool) ([]api.Message, bool, error) { - if s.Compactor == nil { - return messages, skipNotified, nil - } - - keepUserTurns := 0 - req := s.compactionRequest(runID, opts, messages, latest) - req.Force = true - req.KeepUserTurns = &keepUserTurns - s.emitCompactionStarted(runID, opts, CompactionTriggerToolOutput) - - result, err := s.Compactor.MaybeCompact(ctx, req) - if err != nil { - if result.Due && !skipNotified { - s.emitCompactionSkipped(runID, opts, CompactionTriggerToolOutput, result.Reason) - skipNotified = true - } - return messages, skipNotified, nil - } - if !result.Compacted { - if result.Due && !skipNotified { - s.emitCompactionSkipped(runID, opts, CompactionTriggerToolOutput, result.Reason) - skipNotified = true - } - return messages, skipNotified, nil - } - - overflowByID := make(map[string]toolOutputOverflow, len(overflows)) - for _, overflow := range overflows { - overflowByID[overflow.toolCallID] = overflow - } - - compacted := append([]api.Message(nil), result.Messages...) - if !messageEmpty(assistant) { - compacted = append(compacted, assistant) - } - - historyTokens := s.estimateRunPromptTokens(opts, compacted) - batchTokens := 0 - for _, msg := range toolMessages { - content := msg.Content - toolName := msg.ToolName - if overflow, ok := overflowByID[msg.ToolCallID]; ok { - content = overflow.content - if overflow.toolName != "" { - toolName = overflow.toolName - } - } - refit := s.toolMessageForPostCompactionContext(toolName, msg.ToolCallID, content, opts, historyTokens+batchTokens) - compacted = append(compacted, refit) - batchTokens += estimateMessagesTokens([]api.Message{refit}) - } - - s.emitCompacted(runID, opts, compacted, CompactionTriggerToolOutput, result.Summary) - if err := s.checkPostCompactionPromptBudget(opts, compacted); err != nil { - return compacted, skipNotified, err - } - return compacted, skipNotified, nil -} - -func (s *Session) compactionRequest(runID string, opts RunOptions, messages []api.Message, latest api.ChatResponse) CompactionRequest { - meta := newEventMetadata(runID, opts) - return CompactionRequest{ - ChatID: opts.ChatID, - Model: opts.Model, - SystemPrompt: opts.SystemPrompt, - Messages: messages, - Tools: s.availableTools(), - Format: opts.Format, - Latest: latest, - Options: opts.Options, - KeepAlive: opts.KeepAlive, - Think: opts.Think, - ContinueTask: true, - Progress: func(progress CompactionProgress) { - _ = s.emit(newCompactionProgress(meta, progress.Tokens)) - }, - } -} - -func (s *Session) emitCompactionStarted(runID string, opts RunOptions, trigger CompactionTrigger) { - _ = s.emit(newCompactionStarted(newEventMetadata(runID, opts), trigger)) -} - -func (s *Session) emitCompactionSkipped(runID string, opts RunOptions, trigger CompactionTrigger, reason string) { - _ = s.emit(newCompactionSkipped(newEventMetadata(runID, opts), trigger, CompactionSkippedMessage(reason))) -} - -func (s *Session) emitCompacted(runID string, opts RunOptions, messages []api.Message, trigger CompactionTrigger, summary string) { - _ = s.emit(newCompacted(newEventMetadata(runID, opts), messages, trigger, summary)) -} - -func (s *Session) autoCompactionTrigger(req CompactionRequest) CompactionTrigger { - if s.Compactor == nil { - return "" - } - trigger, should := s.Compactor.ShouldCompact(req) - if should { - return CompactionTrigger(trigger) - } - return "" -} - -func CompactionSkippedMessage(reason string) string { - reason = strings.TrimSpace(reason) - if reason == "" { - reason = "compaction could not run" - } - return reason -} - -func resolvedMaxToolRounds(model string, value int) int { - if value != 0 { - return value - } - if modelref.HasExplicitCloudSource(model) { - return -1 - } - return defaultMaxToolRounds -} - -// toolMessageWithBudget sizes a tool result message to fit within a token -// budget (compaction threshold or context window). baseTokens is the -// pre-computed estimate of everything before this message; budgetTokens is -// the ceiling. If the message already fits, it is returned with only the -// small-context rune cap applied. -func (s *Session) toolMessageWithBudget(toolName, toolCallID, content string, opts RunOptions, baseTokens, budgetTokens int) api.Message { - maxRunes := maxToolResultRunes - if limit := smallContextToolResultLimitRunes(s.contextWindowTokens(opts)); limit > 0 { - maxRunes = min(maxRunes, limit) - } - - if budgetTokens <= 0 { - return toolMessageWithLimit(toolName, toolCallID, content, maxRunes) - } - - msg := toolMessageWithLimit(toolName, toolCallID, content, maxRunes) - projectedTokens := baseTokens + estimateMessagesTokens([]api.Message{msg}) - if projectedTokens < budgetTokens { - return msg - } - - overheadTokens := estimateMessagesTokens([]api.Message{{ - Role: "tool", - ToolName: toolName, - ToolCallID: toolCallID, - }}) - // Keep oversized tool output below the budget before it is appended to - // history. This is especially important for <=8k contexts: the next step - // must still have enough room to compact and continue the same user - // request instead of asking the user to prompt again. - availableRunes := (budgetTokens - baseTokens - overheadTokens - toolTruncationMarkerReserveTokens) * 4 - maxRunes = min(maxRunes, max(0, availableRunes)) - msg.Content = truncateToolResultContentTo(content, maxRunes) - return msg -} - -func (s *Session) toolMessageForContext(toolName, toolCallID, content string, opts RunOptions, baseTokens int) api.Message { - return s.toolMessageWithBudget(toolName, toolCallID, content, opts, baseTokens, s.compactionThresholdTokens(opts)) -} - -func (s *Session) toolMessageForPostCompactionContext(toolName, toolCallID, content string, opts RunOptions, baseTokens int) api.Message { - return s.toolMessageWithBudget(toolName, toolCallID, content, opts, baseTokens, s.contextWindowTokens(opts)) -} - -func toolMessageWithLimit(toolName, toolCallID, content string, maxRunes int) api.Message { - return api.Message{ - Role: "tool", - Content: truncateToolResultContentTo(content, maxRunes), - ToolName: toolName, - ToolCallID: toolCallID, - } -} - -func smallContextToolResultLimitRunes(contextWindow int) int { - switch { - case contextWindow > 0 && contextWindow <= tinyContextToolResultTokenWindow: - return tinyContextToolResultRunes - case contextWindow > 0 && contextWindow <= smallContextToolResultTokenWindow: - return smallContextToolResultRunes - default: - return 0 - } -} - -func (s *Session) availableTools() api.Tools { - if s == nil || s.Tools == nil { - return nil - } - return s.Tools.Tools() -} - -func (s *Session) compactionThresholdTokens(opts RunOptions) int { - contextWindow := s.contextWindowTokens(opts) - if contextWindow <= 0 { - return 0 - } - - configuredThreshold := 0.0 - if s.Compactor != nil { - configuredThreshold = s.Compactor.Threshold() - } - - threshold := int(float64(contextWindow) * ResolveCompactionThreshold(configuredThreshold)) - if threshold <= 0 { - return 0 - } - return threshold -} - -func (s *Session) contextWindowTokens(opts RunOptions) int { - if s.Compactor == nil { - return 0 - } - return s.Compactor.ContextWindowTokens(opts.Options) -} - -func toolMessage(toolName, toolCallID, content string) api.Message { - return toolMessageWithLimit(toolName, toolCallID, content, maxToolResultRunes) -} - -func sanitizeMessageForRun(msg api.Message) api.Message { - if msg.Role == "tool" { - msg.Content = truncateToolResultContent(msg.Content) - } - return msg -} - -func sanitizeMessagesForRequest(messages []api.Message) []api.Message { - if len(messages) == 0 { - return nil - } - sanitized := make([]api.Message, len(messages)) - for i, msg := range messages { - sanitized[i] = sanitizeMessageForRun(msg) - } - return sanitized -} - -func truncateToolResultContent(content string) string { - return truncateToolResultContentTo(content, maxToolResultRunes) -} - -func truncateToolResultContentTo(content string, maxRunes int) string { - return Truncate(content, TruncateConfig{ - MaxRunes: maxRunes, - HeadTail: true, - HeadPct: 75, - Label: "tool output", - Hint: "Use a narrower command, line range, or search query if more detail is needed.", - FullOmissionPrefix: toolOutputFullOmissionPrefix, - }) -} - -// TruncateConfig configures content truncation via Truncate. -type TruncateConfig struct { - MaxRunes int // rune limit; <= 0 means full omission - HeadTail bool // true = head + tail split; false = head only - HeadPct int // percentage of MaxRunes for head (e.g. 75); tail gets the rest - Label string // e.g. "tool output", "summary", "stdout" - Hint string // guidance text appended to marker (optional) - FullOmissionPrefix string // marker prefix when MaxRunes <= 0 -} - -// Truncate truncates content to at most cfg.MaxRunes runes. When HeadTail is -// true, it preserves the first HeadPct% and last (100-HeadPct)% of the budget -// with a marker between; otherwise it keeps only the head. MaxRunes <= 0 -// triggers full omission using FullOmissionPrefix. All token counts in -// markers use ApproximateTokens. -func Truncate(content string, cfg TruncateConfig) string { - runes := []rune(content) - total := len(runes) - - if cfg.MaxRunes <= 0 { - return fmt.Sprintf("%s omitted ~%d tokens.%s]", cfg.FullOmissionPrefix, ApproximateTokens(total), truncHint(cfg.Hint)) - } - if total <= cfg.MaxRunes { - return content - } - - if !cfg.HeadTail { - head := cfg.MaxRunes - omitted := total - head - return string(runes[:head]) + TruncMarker(cfg.Label, head, 0, omitted, false, cfg.Hint) - } - - head := cfg.MaxRunes * cfg.HeadPct / 100 - tail := cfg.MaxRunes - head - omitted := total - head - tail - return string(runes[:head]) + TruncMarker(cfg.Label, head, tail, omitted, true, cfg.Hint) + string(runes[len(runes)-tail:]) -} - -func truncHint(hint string) string { - hint = strings.TrimSpace(hint) - if hint == "" { - return "" - } - if !strings.HasSuffix(hint, ".") { - hint += "." - } - return " " + hint -} - -// TruncMarker formats a truncation marker with consistent wording. head and -// tail are rune counts; omitted is the count of runes removed. headTail -// selects the head+tail vs head-only format. hint is optional guidance text. -func TruncMarker(label string, head, tail, omitted int, headTail bool, hint string) string { - var b strings.Builder - b.WriteString("\n\n[") - b.WriteString(label) - b.WriteString(" truncated: ") - if headTail { - fmt.Fprintf(&b, "showing first ~%d tokens and last ~%d tokens; ", ApproximateTokens(head), ApproximateTokens(tail)) - } else { - fmt.Fprintf(&b, "showing first ~%d tokens; ", ApproximateTokens(head)) - } - fmt.Fprintf(&b, "omitted ~%d tokens.%s]", ApproximateTokens(omitted), truncHint(hint)) - if headTail { - b.WriteString("\n\n") - } - return b.String() -} - -func toolOutputFullyOmitted(content string) bool { - return strings.HasPrefix(content, toolOutputFullOmissionPrefix) -} - -// ApproximateTokens estimates token count from a character/byte count using -// the standard ~4 chars-per-token heuristic. It is intentionally rough; all -// callers use it only for sizing/truncation decisions, not billing. -func ApproximateTokens(n int) int { - if n <= 0 { - return 0 - } - return max(1, (n+3)/4) -} - -func messageEmpty(msg api.Message) bool { - return msg.Content == "" && msg.Thinking == "" && len(msg.ToolCalls) == 0 -} diff --git a/agent/session_test.go b/agent/session_test.go deleted file mode 100644 index e95db328090..00000000000 --- a/agent/session_test.go +++ /dev/null @@ -1,2265 +0,0 @@ -package agent - -import ( - "context" - "encoding/json" - "errors" - "os" - "path/filepath" - "strings" - "testing" - - "github.com/ollama/ollama/api" -) - -type fakeClient struct { - calls int - responses [][]api.ChatResponse - requests []*api.ChatRequest - err error -} - -func (c *fakeClient) Chat(ctx context.Context, req *api.ChatRequest, fn api.ChatResponseFunc) error { - c.requests = append(c.requests, req) - if c.calls >= len(c.responses) { - return nil - } - responses := c.responses[c.calls] - c.calls++ - for _, response := range responses { - if err := fn(response); err != nil { - return err - } - } - return c.err -} - -type staticTool struct{} - -type approvalTestTool struct { - called *bool -} - -type namedApprovalTestTool struct { - name string -} - -type cwdTestTool struct{} - -type largeTool struct{} - -type preTruncatedTool struct{} - -type cancelingTool struct { - cancel context.CancelFunc -} - -type cancelAfterToolCallClient struct { - cancel context.CancelFunc -} - -type recordingCompactor struct { - requests []CompactionRequest -} - -type oversizedCompactor struct { - requests []CompactionRequest -} - -type recordingEventSink struct { - events []Event -} - -func (s *recordingEventSink) Emit(event Event) error { - s.events = append(s.events, event) - return nil -} - -func hasEventType(events []Event, eventType EventType) bool { - for _, event := range events { - if event.Type == eventType { - return true - } - } - return false -} - -func hasEventWithTokens(events []Event, eventType EventType, tokens int) bool { - for _, event := range events { - if event.Type == eventType && event.Tokens == tokens { - return true - } - } - return false -} - -func TestSessionEmitsToAllSinksAfterError(t *testing.T) { - errSink := EventSinkFunc(func(Event) error { - return errors.New("sink failed") - }) - events := &recordingEventSink{} - session := &Session{EventSinks: []EventSink{errSink, events}} - - err := session.emit(Event{Type: EventRunFinished}) - if err == nil { - t.Fatal("emit should return the first sink error") - } - if !hasEventType(events.events, EventRunFinished) { - t.Fatalf("later sink did not receive event after earlier error: %#v", events.events) - } -} - -func (c cancelAfterToolCallClient) Chat(ctx context.Context, req *api.ChatRequest, fn api.ChatResponseFunc) error { - args := api.NewToolCallFunctionArguments() - args.Set("value", "skip me") - if err := fn(api.ChatResponse{Message: api.Message{Role: "assistant", ToolCalls: []api.ToolCall{{ - ID: "call-1", - Function: api.ToolCallFunction{ - Name: "echo_tool", - Arguments: args, - }, - }}}}); err != nil { - return err - } - c.cancel() - return context.Canceled -} - -func (c *recordingCompactor) MaybeCompact(_ context.Context, req CompactionRequest) (CompactionResult, error) { - c.requests = append(c.requests, req) - result := CompactionResult{Messages: req.Messages, Due: true} - if len(req.Messages) > 0 && req.Messages[len(req.Messages)-1].Role == "tool" { - result.Messages = CompactionSummaryMessages("tool result summarized", false) - result.Compacted = true - result.Summary = "tool result summarized" - } - return result, nil -} - -func (c *recordingCompactor) ContextWindowTokens(options map[string]any) int { - return ResolveContextWindowTokens(options, 0) -} -func (c *recordingCompactor) Threshold() float64 { return 0 } -func (c *recordingCompactor) ShouldCompact(_ CompactionRequest) (string, bool) { - return "", false -} - -func (c *oversizedCompactor) MaybeCompact(_ context.Context, req CompactionRequest) (CompactionResult, error) { - c.requests = append(c.requests, req) - summary := strings.Repeat("oversized summary ", 300) - return CompactionResult{ - Messages: CompactionSummaryMessages(summary, req.ContinueTask), - Compacted: true, - Due: true, - Summary: summary, - }, nil -} - -func (c *oversizedCompactor) ContextWindowTokens(options map[string]any) int { - return ResolveContextWindowTokens(options, 0) -} -func (c *oversizedCompactor) Threshold() float64 { return 0 } -func (c *oversizedCompactor) ShouldCompact(_ CompactionRequest) (string, bool) { - return "", false -} - -type recordingApprovalPrompter struct { - requests []ApprovalRequest - results []Approval -} - -func approvalStateForTest(allowAll bool, scopes map[string]bool) *ApprovalState { - state := &ApprovalState{} - state.Set(allowAll, scopes) - return state -} - -func (staticTool) Name() string { - return "echo_tool" -} - -func (staticTool) Description() string { - return "echoes a value" -} - -func (staticTool) Schema() api.ToolFunction { - props := api.NewToolPropertiesMap() - props.Set("value", api.ToolProperty{Type: api.PropertyType{"string"}}) - return api.ToolFunction{ - Name: "echo_tool", - Description: "echoes a value", - Parameters: api.ToolFunctionParameters{ - Type: "object", - Properties: props, - }, - } -} - -func (staticTool) Execute(context.Context, ToolContext, map[string]any) (ToolResult, error) { - return ToolResult{Content: "tool says hello"}, nil -} - -func (largeTool) Name() string { - return "large_tool" -} - -func (largeTool) Description() string { - return "returns a large result" -} - -func (largeTool) Schema() api.ToolFunction { - return api.ToolFunction{ - Name: "large_tool", - Description: "returns a large result", - Parameters: api.ToolFunctionParameters{ - Type: "object", - }, - } -} - -func (largeTool) Execute(context.Context, ToolContext, map[string]any) (ToolResult, error) { - return ToolResult{Content: strings.Repeat("x", maxToolResultRunes+100)}, nil -} - -func (preTruncatedTool) Name() string { - return "pre_truncated_tool" -} - -func (preTruncatedTool) Description() string { - return "returns a large result that is already marked as truncated" -} - -func (preTruncatedTool) Schema() api.ToolFunction { - return api.ToolFunction{ - Name: "pre_truncated_tool", - Description: "returns a large result that is already marked as truncated", - Parameters: api.ToolFunctionParameters{ - Type: "object", - }, - } -} - -func (preTruncatedTool) Execute(context.Context, ToolContext, map[string]any) (ToolResult, error) { - content := strings.Repeat("x", smallContextToolResultRunes) + - "\n\n[tool output truncated: showing first ~1500 tokens; omitted ~999 tokens. Use a narrower command, line range, or search query if more detail is needed.]\n\n" + - strings.Repeat("y", smallContextToolResultRunes) - return ToolResult{Content: content}, nil -} - -func (t cancelingTool) Name() string { - return "cancel_tool" -} - -func (t cancelingTool) Description() string { - return "cancels while running" -} - -func (t cancelingTool) Schema() api.ToolFunction { - return api.ToolFunction{ - Name: t.Name(), - Description: t.Description(), - Parameters: api.ToolFunctionParameters{ - Type: "object", - }, - } -} - -func (t cancelingTool) Execute(ctx context.Context, _ ToolContext, _ map[string]any) (ToolResult, error) { - t.cancel() - <-ctx.Done() - return ToolResult{}, ctx.Err() -} - -func (t approvalTestTool) Name() string { - return "approval_tool" -} - -func (t approvalTestTool) Description() string { - return "requires approval" -} - -func (t approvalTestTool) Schema() api.ToolFunction { - return api.ToolFunction{ - Name: "approval_tool", - Description: "requires approval", - Parameters: api.ToolFunctionParameters{ - Type: "object", - }, - } -} - -func (t approvalTestTool) RequiresApproval(map[string]any) bool { - return true -} - -func (t approvalTestTool) Execute(context.Context, ToolContext, map[string]any) (ToolResult, error) { - if t.called != nil { - *t.called = true - } - return ToolResult{Content: "approved"}, nil -} - -func (t namedApprovalTestTool) Name() string { - return t.name -} - -func (t namedApprovalTestTool) Description() string { - return "requires approval" -} - -func (t namedApprovalTestTool) Schema() api.ToolFunction { - return api.ToolFunction{ - Name: t.name, - Description: "requires approval", - Parameters: api.ToolFunctionParameters{ - Type: "object", - }, - } -} - -func (t namedApprovalTestTool) RequiresApproval(map[string]any) bool { - return true -} - -func (t namedApprovalTestTool) Execute(context.Context, ToolContext, map[string]any) (ToolResult, error) { - return ToolResult{Content: "approved"}, nil -} - -// ApprovalScope mimics the Bash tool's command-scoping behavior so tests can -// exercise the shell approval flow without importing the tools package. -func (t namedApprovalTestTool) ApprovalScope(args map[string]any) string { - if t.name == "bash" || t.name == "powershell" { - if cmd, ok := args["command"].(string); ok { - cmd = strings.TrimSpace(cmd) - if cmd != "" { - return t.name + "\x00" + cmd - } - } - } - return t.name -} - -func (p *recordingApprovalPrompter) PromptApproval(_ context.Context, req ApprovalRequest) (Approval, error) { - p.requests = append(p.requests, req) - if len(p.results) == 0 { - return Approval{Allow: true}, nil - } - result := p.results[0] - p.results = p.results[1:] - return result, nil -} - -func (cwdTestTool) Name() string { - return "cwd_tool" -} - -func (cwdTestTool) Description() string { - return "tests cwd state" -} - -func (cwdTestTool) Schema() api.ToolFunction { - props := api.NewToolPropertiesMap() - props.Set("mode", api.ToolProperty{Type: api.PropertyType{"string"}}) - props.Set("path", api.ToolProperty{Type: api.PropertyType{"string"}}) - return api.ToolFunction{ - Name: "cwd_tool", - Description: "tests cwd state", - Parameters: api.ToolFunctionParameters{ - Type: "object", - Properties: props, - }, - } -} - -func (cwdTestTool) RequiresApproval(map[string]any) bool { - return true -} - -func (cwdTestTool) Execute(_ context.Context, toolCtx ToolContext, args map[string]any) (ToolResult, error) { - switch args["mode"] { - case "set": - path, _ := args["path"].(string) - return ToolResult{Content: "changed", WorkingDir: filepath.Join(toolCtx.WorkingDir, path)}, nil - case "escape": - return ToolResult{Content: "escaped", WorkingDir: filepath.Dir(toolCtx.WorkingDir)}, nil - default: - return ToolResult{Content: toolCtx.WorkingDir}, nil - } -} - -func TestSessionRunsToolLoop(t *testing.T) { - args := api.NewToolCallFunctionArguments() - args.Set("value", "hello") - - client := &fakeClient{ - responses: [][]api.ChatResponse{ - { - {Message: api.Message{Role: "assistant", ToolCalls: []api.ToolCall{{ - ID: "call-1", - Function: api.ToolCallFunction{ - Name: "echo_tool", - Arguments: args, - }, - }}}}, - }, - { - {Message: api.Message{Role: "assistant", Content: "done"}}, - }, - }, - } - - registry := &Registry{} - registry.Register(staticTool{}) - - session := &Session{ - Client: client, - Tools: registry, - ApprovalState: approvalStateForTest(true, nil), - } - - result, err := session.Run(context.Background(), RunOptions{ - ChatID: "chat-1", - Model: "model", - NewMessages: []api.Message{{Role: "user", Content: "use a tool"}}, - }) - if err != nil { - t.Fatal(err) - } - - if client.calls != 2 { - t.Fatalf("client calls = %d, want 2", client.calls) - } - if len(result.Messages) != 4 { - t.Fatalf("messages = %d, want 4", len(result.Messages)) - } - if result.Messages[2].Role != "tool" || result.Messages[2].Content != "tool says hello" { - t.Fatalf("tool message = %#v", result.Messages[2]) - } - if len(client.requests[0].Tools) != 1 { - t.Fatalf("first request tools = %d, want 1", len(client.requests[0].Tools)) - } - if len(client.requests[1].Messages) != 3 { - t.Fatalf("second request messages = %d, want 3", len(client.requests[1].Messages)) - } -} - -func TestSessionAddsSystemPromptOnlyToRequest(t *testing.T) { - client := &fakeClient{ - responses: [][]api.ChatResponse{ - {{Message: api.Message{Role: "assistant", Content: "done"}}}, - }, - } - session := &Session{Client: client} - - _, err := session.Run(context.Background(), RunOptions{ - ChatID: "chat-1", - Model: "model", - SystemPrompt: "available context: go-code", - NewMessages: []api.Message{{Role: "user", Content: "hello"}}, - }) - if err != nil { - t.Fatal(err) - } - if len(client.requests) != 1 { - t.Fatalf("requests = %d, want 1", len(client.requests)) - } - reqMessages := client.requests[0].Messages - if len(reqMessages) != 2 || reqMessages[0].Role != "system" || reqMessages[0].Content != "available context: go-code" { - t.Fatalf("request messages = %#v", reqMessages) - } -} - -func TestSessionChatRequestMatchesRunRequest(t *testing.T) { - client := &fakeClient{ - responses: [][]api.ChatResponse{ - {{Message: api.Message{Role: "assistant", Content: "done"}}}, - }, - } - registry := &Registry{} - registry.Register(staticTool{}) - session := &Session{Client: client, Tools: registry} - opts := RunOptions{ - ChatID: "chat-1", - Model: "model", - SystemPrompt: "available context: go-code", - NewMessages: []api.Message{{Role: "user", Content: "hello"}}, - Format: "json", - Options: map[string]any{"temperature": 0.5}, - } - - want := buildChatRequest(opts, opts.NewMessages, registry.Tools()) - _, err := session.Run(context.Background(), opts) - if err != nil { - t.Fatal(err) - } - if len(client.requests) != 1 { - t.Fatalf("requests = %d, want 1", len(client.requests)) - } - gotJSON, err := json.Marshal(client.requests[0]) - if err != nil { - t.Fatal(err) - } - wantJSON, err := json.Marshal(want) - if err != nil { - t.Fatal(err) - } - if string(gotJSON) != string(wantJSON) { - t.Fatalf("ChatRequest mismatch\ngot: %s\nwant: %s", gotJSON, wantJSON) - } -} - -func TestSessionAccumulatesStreamingAssistantMessage(t *testing.T) { - args := api.NewToolCallFunctionArguments() - args.Set("value", "hello") - responses := make([]api.ChatResponse, 0, 100) - var wantContent, wantThinking string - for range 99 { - wantContent += "x" - wantThinking += "t" - responses = append(responses, api.ChatResponse{ - Message: api.Message{Role: "assistant", Content: "x", Thinking: "t"}, - }) - } - toolCall := api.ToolCall{ - ID: "call-1", - Function: api.ToolCallFunction{ - Name: "echo_tool", - Arguments: args, - }, - } - responses = append(responses, api.ChatResponse{ - Message: api.Message{Role: "assistant", ToolCalls: []api.ToolCall{toolCall}}, - }) - - session := &Session{ - Client: &fakeClient{responses: [][]api.ChatResponse{responses}}, - } - - result, err := session.Run(context.Background(), RunOptions{ - ChatID: "chat-1", - Model: "model", - NewMessages: []api.Message{{Role: "user", Content: "stream"}}, - }) - if err != nil { - t.Fatal(err) - } - - if len(result.Messages) != 2 || result.Messages[1].Content != wantContent || result.Messages[1].Thinking != wantThinking || len(result.Messages[1].ToolCalls) != 1 { - t.Fatalf("result messages = %#v", result.Messages) - } -} - -func TestSessionRequestHistoryKeepsThinkingAndServerToolCallIDs(t *testing.T) { - args := api.NewToolCallFunctionArguments() - args.Set("value", "hello") - client := &fakeClient{ - responses: [][]api.ChatResponse{ - { - {Message: api.Message{Role: "assistant", Thinking: "private chain"}}, - {Message: api.Message{Role: "assistant", Content: "I'll check."}}, - {Message: api.Message{Role: "assistant", ToolCalls: []api.ToolCall{{ - ID: "volatile-random-id", - Function: api.ToolCallFunction{ - Name: "echo_tool", - Arguments: args, - }, - }}}}, - }, - {{Message: api.Message{Role: "assistant", Content: "done"}}}, - }, - } - registry := &Registry{} - registry.Register(staticTool{}) - session := &Session{ - Client: client, - Tools: registry, - ApprovalState: approvalStateForTest(true, nil), - } - - result, err := session.Run(context.Background(), RunOptions{ - ChatID: "chat-1", - Model: "model", - NewMessages: []api.Message{{Role: "user", Content: "use a tool"}}, - }) - if err != nil { - t.Fatal(err) - } - if len(client.requests) != 2 { - t.Fatalf("requests = %d, want 2", len(client.requests)) - } - - secondRequestMessages := client.requests[1].Messages - if len(secondRequestMessages) != 3 { - t.Fatalf("second request messages = %#v", secondRequestMessages) - } - assistant := secondRequestMessages[1] - if assistant.Role != "assistant" { - t.Fatalf("second request assistant = %#v", assistant) - } - if assistant.Thinking != "private chain" { - t.Fatalf("assistant thinking = %q, want preserved", assistant.Thinking) - } - if len(assistant.ToolCalls) != 1 || assistant.ToolCalls[0].ID != "volatile-random-id" { - t.Fatalf("assistant tool calls = %#v", assistant.ToolCalls) - } - tool := secondRequestMessages[2] - if tool.Role != "tool" || tool.ToolCallID != "volatile-random-id" { - t.Fatalf("tool result message = %#v", tool) - } - if len(result.Messages) < 3 || result.Messages[1].Thinking != "private chain" { - t.Fatalf("visible result messages lost thinking: %#v", result.Messages) - } -} - -func TestSessionKeepsPartialStreamOnCancellation(t *testing.T) { - session := &Session{ - Client: &fakeClient{ - responses: [][]api.ChatResponse{{ - {Message: api.Message{Role: "assistant", Content: "partial "}}, - {Message: api.Message{Role: "assistant", Content: "answer"}}, - }}, - err: context.Canceled, - }, - } - - result, err := session.Run(context.Background(), RunOptions{ - ChatID: "chat-1", - Model: "model", - NewMessages: []api.Message{{Role: "user", Content: "cancel"}}, - }) - if err != nil { - t.Fatal(err) - } - if len(result.Messages) != 2 || result.Messages[1].Content != "partial answer" { - t.Fatalf("result messages = %#v", result.Messages) - } -} - -func TestSessionCancellationKeepsPartialResultWhenUISinkCancels(t *testing.T) { - ctx, cancel := context.WithCancel(context.Background()) - cancel() - trace := &recordingEventSink{} - session := &Session{ - Client: &fakeClient{ - responses: [][]api.ChatResponse{{ - {Message: api.Message{Role: "assistant", Content: "partial"}}, - }}, - err: context.Canceled, - }, - EventSinks: []EventSink{ - EventSinkFunc(func(event Event) error { - if event.Type == EventRunFinished { - return context.Canceled - } - return nil - }), - trace, - }, - } - - result, err := session.Run(ctx, RunOptions{ - ChatID: "chat-1", - Model: "model", - NewMessages: []api.Message{{Role: "user", Content: "cancel"}}, - }) - if err != nil { - t.Fatal(err) - } - if result == nil || len(result.Messages) != 2 || result.Messages[1].Content != "partial" { - t.Fatalf("result messages = %#v, want partial assistant result", result) - } - if !hasEventType(trace.events, EventRunFinished) { - t.Fatalf("trace sink did not receive run finished event: %#v", trace.events) - } -} - -func TestSessionTreatsHTTPContextCanceledStringAsCancellation(t *testing.T) { - ctx, cancel := context.WithCancel(context.Background()) - cancel() - client := &fakeClient{err: errors.New(`Post "http://127.0.0.1:11434/api/chat": context canceled`)} - session := &Session{Client: client} - - result, err := session.Run(ctx, RunOptions{ - Model: "model", - NewMessages: []api.Message{{Role: "user", Content: "hello"}}, - }) - if err != nil { - t.Fatalf("Run returned error for canceled HTTP request: %v", err) - } - if result == nil { - t.Fatal("Run returned nil result") - } - if len(result.Messages) != 1 || result.Messages[0].Content != "hello" { - t.Fatalf("messages = %#v, want original user message only", result.Messages) - } -} - -func TestSessionDisabledToolsOmitToolsAndReturnsDisabledResults(t *testing.T) { - args := api.NewToolCallFunctionArguments() - args.Set("value", "hello") - client := &fakeClient{ - responses: [][]api.ChatResponse{ - {{ - Message: api.Message{Role: "assistant", ToolCalls: []api.ToolCall{{ - ID: "call-1", - Function: api.ToolCallFunction{ - Name: "echo_tool", - Arguments: args, - }, - }}}, - }}, - {{Message: api.Message{Role: "assistant", Content: "tools are off"}}}, - }, - } - registry := &Registry{} - registry.Register(staticTool{}) - events := &recordingEventSink{} - session := &Session{ - Client: client, - EventSinks: []EventSink{events}, - Tools: registry, - DisableTools: true, - } - - result, err := session.Run(context.Background(), RunOptions{ - ChatID: "chat-1", - Model: "model", - NewMessages: []api.Message{{Role: "user", Content: "use a tool"}}, - }) - if err != nil { - t.Fatal(err) - } - if len(client.requests) != 2 { - t.Fatalf("requests = %d, want 2", len(client.requests)) - } - if got := len(client.requests[0].Tools); got != 0 { - t.Fatalf("advertised tools = %d, want 0", got) - } - secondMessages := client.requests[1].Messages - if len(secondMessages) != 3 { - t.Fatalf("second request messages = %#v", secondMessages) - } - if secondMessages[2].Role != "tool" || secondMessages[2].ToolCallID != "call-1" || secondMessages[2].Content != toolExecutionDisabledMessage { - t.Fatalf("disabled tool message = %#v", secondMessages[2]) - } - if len(result.Messages) != 4 || result.Messages[2].Content != toolExecutionDisabledMessage { - t.Fatalf("result messages = %#v", result.Messages) - } - var sawDetected, sawDisabled bool - for _, event := range events.events { - if event.Type == EventToolCallDetected { - sawDetected = true - } - if event.Type == EventToolFinished && event.ToolStatus == ToolStatusDisabled && event.Content == toolExecutionDisabledMessage { - sawDisabled = true - } - } - if !sawDetected || !sawDisabled { - t.Fatalf("events missing detected/disabled: %#v", events.events) - } -} - -func TestSessionCancellationAfterToolCallAppendsSkippedToolMessage(t *testing.T) { - ctx, cancel := context.WithCancel(context.Background()) - registry := &Registry{} - registry.Register(staticTool{}) - session := &Session{ - Client: cancelAfterToolCallClient{cancel: cancel}, - Tools: registry, - ApprovalState: approvalStateForTest(true, nil), - } - - result, err := session.Run(ctx, RunOptions{ - ChatID: "chat-1", - Model: "model", - NewMessages: []api.Message{{Role: "user", Content: "cancel after tool"}}, - }) - if err != nil { - t.Fatal(err) - } - if len(result.Messages) != 3 { - t.Fatalf("messages = %#v", result.Messages) - } - if len(result.Messages[1].ToolCalls) != 1 { - t.Fatalf("assistant tool calls = %#v", result.Messages[1]) - } - if result.Messages[2].Role != "tool" || result.Messages[2].ToolCallID != "call-1" { - t.Fatalf("skipped tool message = %#v", result.Messages[2]) - } - if !strings.Contains(result.Messages[2].Content, "run was canceled") { - t.Fatalf("skipped content = %q", result.Messages[2].Content) - } -} - -func TestSessionCancellationDuringToolExecutionAppendsToolMessage(t *testing.T) { - ctx, cancel := context.WithCancel(context.Background()) - events := &recordingEventSink{} - registry := &Registry{} - registry.Register(cancelingTool{cancel: cancel}) - client := &fakeClient{responses: [][]api.ChatResponse{{ - {Message: api.Message{Role: "assistant", ToolCalls: []api.ToolCall{{ - ID: "call-1", - Function: api.ToolCallFunction{ - Name: "cancel_tool", - }, - }}}}, - }}} - session := &Session{ - Client: client, - Tools: registry, - ApprovalState: approvalStateForTest(true, nil), - EventSinks: []EventSink{events}, - } - - result, err := session.Run(ctx, RunOptions{ - ChatID: "chat-1", - Model: "model", - NewMessages: []api.Message{{Role: "user", Content: "cancel during tool"}}, - }) - if err != nil { - t.Fatal(err) - } - if len(result.Messages) != 3 { - t.Fatalf("messages = %#v", result.Messages) - } - if result.Messages[2].Role != "tool" || result.Messages[2].ToolCallID != "call-1" { - t.Fatalf("tool message = %#v", result.Messages[2]) - } - if !strings.Contains(result.Messages[2].Content, "context canceled") { - t.Fatalf("tool content = %q", result.Messages[2].Content) - } - var finished *Event - for i := range events.events { - if events.events[i].Type == EventRunFinished { - finished = &events.events[i] - } - } - if finished == nil { - t.Fatalf("run finished event missing: %#v", events.events) - } - if finished.Status != RunStatusCanceled { - t.Fatalf("run status = %q, want canceled", finished.Status) - } -} - -func TestSessionToolLoopAllowsRoundsUnderDefaultCap(t *testing.T) { - responses := make([][]api.ChatResponse, 0, 26) - for i := range 25 { - args := api.NewToolCallFunctionArguments() - args.Set("value", "hello") - responses = append(responses, []api.ChatResponse{{ - Message: api.Message{Role: "assistant", ToolCalls: []api.ToolCall{{ - ID: "call-" + string(rune('a'+i)), - Function: api.ToolCallFunction{ - Name: "echo_tool", - Arguments: args, - }, - }}}, - }}) - } - responses = append(responses, []api.ChatResponse{{ - Message: api.Message{Role: "assistant", Content: "done"}, - }}) - - client := &fakeClient{responses: responses} - registry := &Registry{} - registry.Register(staticTool{}) - - session := &Session{ - Client: client, - Tools: registry, - ApprovalState: approvalStateForTest(true, nil), - } - - if _, err := session.Run(context.Background(), RunOptions{ - Model: "model", - NewMessages: []api.Message{{Role: "user", Content: "keep going"}}, - }); err != nil { - t.Fatal(err) - } - - if client.calls != 26 { - t.Fatalf("client calls = %d, want 26", client.calls) - } -} - -func TestSessionLocalToolRoundLimitAppendsSkippedToolMessages(t *testing.T) { - firstArgs := api.NewToolCallFunctionArguments() - firstArgs.Set("value", "first") - secondArgs := api.NewToolCallFunctionArguments() - secondArgs.Set("value", "second") - thirdArgs := api.NewToolCallFunctionArguments() - thirdArgs.Set("value", "third") - - client := &fakeClient{ - responses: [][]api.ChatResponse{ - {{ - Message: api.Message{Role: "assistant", ToolCalls: []api.ToolCall{{ - ID: "call-1", - Function: api.ToolCallFunction{ - Name: "echo_tool", - Arguments: firstArgs, - }, - }}}, - }}, - {{ - Message: api.Message{Role: "assistant", ToolCalls: []api.ToolCall{ - { - ID: "call-2", - Function: api.ToolCallFunction{ - Name: "echo_tool", - Arguments: secondArgs, - }, - }, - { - ID: "call-3", - Function: api.ToolCallFunction{ - Name: "echo_tool", - Arguments: thirdArgs, - }, - }, - }}, - }}, - }, - } - registry := &Registry{} - registry.Register(staticTool{}) - session := &Session{ - Client: client, - Tools: registry, - ApprovalState: approvalStateForTest(true, nil), - } - - result, err := session.Run(context.Background(), RunOptions{ - ChatID: "chat-1", - Model: "test:local", - NewMessages: []api.Message{{Role: "user", Content: "hit cap"}}, - MaxToolRounds: 1, - }) - if err == nil || !strings.Contains(err.Error(), "tool round limit reached after 1 rounds") { - t.Fatalf("error = %v, want tool-round limit", err) - } - if result == nil { - t.Fatal("expected partial result with skipped tool messages") - } - if len(result.Messages) != 6 { - t.Fatalf("messages = %#v", result.Messages) - } - for i, wantID := range []string{"call-2", "call-3"} { - msg := result.Messages[4+i] - if msg.Role != "tool" || msg.ToolCallID != wantID { - t.Fatalf("skipped tool %d = %#v", i, msg) - } - if !strings.Contains(msg.Content, "max tool-round limit of 1") { - t.Fatalf("skipped content = %q", msg.Content) - } - } -} - -func TestSessionLocalToolLoopStopsAtDefaultRoundCap(t *testing.T) { - responses := make([][]api.ChatResponse, 0, defaultMaxToolRounds+1) - for range defaultMaxToolRounds + 1 { - args := api.NewToolCallFunctionArguments() - args.Set("value", "hello") - responses = append(responses, []api.ChatResponse{{ - Message: api.Message{Role: "assistant", ToolCalls: []api.ToolCall{{ - ID: "call", - Function: api.ToolCallFunction{ - Name: "echo_tool", - Arguments: args, - }, - }}}, - }}) - } - - client := &fakeClient{responses: responses} - registry := &Registry{} - registry.Register(staticTool{}) - session := &Session{ - Client: client, - Tools: registry, - ApprovalState: approvalStateForTest(true, nil), - } - - _, err := session.Run(context.Background(), RunOptions{ - Model: "test:local", - NewMessages: []api.Message{{Role: "user", Content: "keep going"}}, - }) - if err == nil || !strings.Contains(err.Error(), "tool round limit reached after 100 rounds") { - t.Fatalf("error = %v, want default tool round limit", err) - } - if client.calls != defaultMaxToolRounds+1 { - t.Fatalf("client calls = %d, want %d", client.calls, defaultMaxToolRounds+1) - } -} - -func TestSessionCloudToolLoopHonorsExplicitRoundCap(t *testing.T) { - args := api.NewToolCallFunctionArguments() - args.Set("value", "hello") - client := &fakeClient{responses: [][]api.ChatResponse{ - {{ - Message: api.Message{Role: "assistant", ToolCalls: []api.ToolCall{{ - ID: "call-1", - Function: api.ToolCallFunction{ - Name: "echo_tool", - Arguments: args, - }, - }}}, - }}, - {{ - Message: api.Message{Role: "assistant", ToolCalls: []api.ToolCall{{ - ID: "call-2", - Function: api.ToolCallFunction{ - Name: "echo_tool", - Arguments: args, - }, - }}}, - }}, - }} - registry := &Registry{} - registry.Register(staticTool{}) - session := &Session{ - Client: client, - Tools: registry, - ApprovalState: approvalStateForTest(true, nil), - } - - result, err := session.Run(context.Background(), RunOptions{ - Model: "test:cloud", - NewMessages: []api.Message{{Role: "user", Content: "keep going"}}, - MaxToolRounds: 1, - }) - if err == nil || !strings.Contains(err.Error(), "tool round limit reached after 1 rounds") { - t.Fatalf("error = %v, want explicit tool-round limit", err) - } - if result == nil { - t.Fatal("expected partial result with skipped tool message") - } - if client.calls != 2 { - t.Fatalf("client calls = %d, want 2", client.calls) - } -} - -func TestSessionToolLoopNegativeLimitIsUnlimited(t *testing.T) { - responses := make([][]api.ChatResponse, 0, defaultMaxToolRounds+2) - for range defaultMaxToolRounds + 1 { - args := api.NewToolCallFunctionArguments() - args.Set("value", "hello") - responses = append(responses, []api.ChatResponse{{ - Message: api.Message{Role: "assistant", ToolCalls: []api.ToolCall{{ - ID: "call", - Function: api.ToolCallFunction{ - Name: "echo_tool", - Arguments: args, - }, - }}}, - }}) - } - responses = append(responses, []api.ChatResponse{{ - Message: api.Message{Role: "assistant", Content: "done"}, - }}) - - client := &fakeClient{responses: responses} - registry := &Registry{} - registry.Register(staticTool{}) - session := &Session{ - Client: client, - Tools: registry, - ApprovalState: approvalStateForTest(true, nil), - } - - if _, err := session.Run(context.Background(), RunOptions{ - Model: "model", - NewMessages: []api.Message{{Role: "user", Content: "keep going"}}, - MaxToolRounds: -1, - }); err != nil { - t.Fatal(err) - } - if client.calls != defaultMaxToolRounds+2 { - t.Fatalf("client calls = %d, want %d", client.calls, defaultMaxToolRounds+2) - } -} - -func TestSessionTruncatesLargeToolResultsBeforeHistory(t *testing.T) { - args := api.NewToolCallFunctionArguments() - client := &fakeClient{ - responses: [][]api.ChatResponse{ - {{ - Message: api.Message{Role: "assistant", ToolCalls: []api.ToolCall{{ - ID: "call-1", - Function: api.ToolCallFunction{ - Name: "large_tool", - Arguments: args, - }, - }}}, - }}, - {{Message: api.Message{Role: "assistant", Content: "done"}}}, - }, - } - registry := &Registry{} - registry.Register(largeTool{}) - session := &Session{ - Client: client, - Tools: registry, - ApprovalState: approvalStateForTest(true, nil), - } - - result, err := session.Run(context.Background(), RunOptions{ - ChatID: "chat-1", - Model: "model", - NewMessages: []api.Message{{Role: "user", Content: "use a tool"}}, - }) - if err != nil { - t.Fatal(err) - } - if len(result.Messages) < 3 { - t.Fatalf("messages = %#v", result.Messages) - } - content := result.Messages[2].Content - if !strings.Contains(content, "[tool output truncated: showing first ~") || - !strings.Contains(content, "omitted ~25 tokens") || - !strings.Contains(content, "Use a narrower command, line range, or search query") { - t.Fatalf("tool content missing truncation marker: %q", content) - } - if strings.Count(content, "x") != maxToolResultRunes { - t.Fatalf("truncated content x count = %d, want %d", strings.Count(content, "x"), maxToolResultRunes) - } - requestContent := client.requests[1].Messages[2].Content - if !strings.Contains(requestContent, "[tool output truncated: showing first ~") { - t.Fatalf("second model request did not use capped tool content: %q", requestContent) - } - if strings.Count(requestContent, "x") > maxToolResultRunes { - t.Fatalf("request tool content x count = %d, want at most %d", strings.Count(requestContent, "x"), maxToolResultRunes) - } -} - -func TestSessionSmallContextUsesLowerToolResultPreviewCap(t *testing.T) { - args := api.NewToolCallFunctionArguments() - client := &fakeClient{ - responses: [][]api.ChatResponse{ - {{ - Message: api.Message{Role: "assistant", ToolCalls: []api.ToolCall{{ - ID: "call-1", - Function: api.ToolCallFunction{ - Name: "large_tool", - Arguments: args, - }, - }}}, - }}, - {{Message: api.Message{Role: "assistant", Content: "done"}}}, - }, - } - registry := &Registry{} - registry.Register(largeTool{}) - session := &Session{ - Client: client, - Tools: registry, - ApprovalState: approvalStateForTest(true, nil), - Compactor: &SimpleCompactor{Client: nil, Options: CompactionOptions{ - ContextWindowTokens: smallContextToolResultTokenWindow, - }}, - } - - result, err := session.Run(context.Background(), RunOptions{ - Model: "model", - NewMessages: []api.Message{{Role: "user", Content: "use a tool"}}, - }) - if err != nil { - t.Fatal(err) - } - - content := result.Messages[2].Content - if !strings.Contains(content, "[tool output truncated: showing first ~") || - !strings.Contains(content, "Use a narrower command, line range, or search query") { - t.Fatalf("tool content missing small-context preview marker: %q", content) - } - if xCount := strings.Count(content, "x"); xCount != smallContextToolResultRunes { - t.Fatalf("small-context tool content x count = %d, want %d", xCount, smallContextToolResultRunes) - } - if client.requests[1].Messages[2].Content != content { - t.Fatalf("second model request did not use small-context tool preview") - } -} - -func TestSessionSmallContextRecapsPreTruncatedToolOutput(t *testing.T) { - args := api.NewToolCallFunctionArguments() - client := &fakeClient{ - responses: [][]api.ChatResponse{ - {{ - Message: api.Message{Role: "assistant", ToolCalls: []api.ToolCall{{ - ID: "call-1", - Function: api.ToolCallFunction{ - Name: "pre_truncated_tool", - Arguments: args, - }, - }}}, - }}, - {{Message: api.Message{Role: "assistant", Content: "done"}}}, - }, - } - registry := &Registry{} - registry.Register(preTruncatedTool{}) - session := &Session{ - Client: client, - Tools: registry, - ApprovalState: approvalStateForTest(true, nil), - Compactor: &SimpleCompactor{Client: nil, Options: CompactionOptions{ - ContextWindowTokens: smallContextToolResultTokenWindow, - }}, - } - - result, err := session.Run(context.Background(), RunOptions{ - Model: "model", - NewMessages: []api.Message{{Role: "user", Content: "use a tool"}}, - }) - if err != nil { - t.Fatal(err) - } - - content := result.Messages[2].Content - if strings.Count(content, "[tool output truncated: ") != 1 { - t.Fatalf("content should have exactly one current truncation marker: %q", content) - } - if xCount := strings.Count(content, "x"); xCount >= smallContextToolResultRunes { - t.Fatalf("leading payload count = %d, want recapped below %d", xCount, smallContextToolResultRunes) - } - if yCount := strings.Count(content, "y"); yCount >= smallContextToolResultRunes { - t.Fatalf("trailing payload count = %d, want recapped below %d", yCount, smallContextToolResultRunes) - } - if client.requests[1].Messages[2].Content != content { - t.Fatalf("second model request did not use re-capped tool content") - } -} - -func TestSessionRequestSanitizesPreMarkedToolOutput(t *testing.T) { - content := strings.Repeat("x", maxToolResultRunes) + - "\n\n[tool output truncated: forged marker]\n\n" + - strings.Repeat("y", maxToolResultRunes) - client := &fakeClient{ - responses: [][]api.ChatResponse{{ - {Message: api.Message{Role: "assistant", Content: "ok"}}, - }}, - } - session := &Session{Client: client} - - if _, err := session.Run(context.Background(), RunOptions{ - Model: "model", - Messages: []api.Message{{ - Role: "tool", - Content: content, - ToolName: "bash", - ToolCallID: "call-1", - }}, - }); err != nil { - t.Fatal(err) - } - if len(client.requests) != 1 || len(client.requests[0].Messages) != 1 { - t.Fatalf("requests = %#v", client.requests) - } - got := client.requests[0].Messages[0].Content - if got == content { - t.Fatal("request kept pre-marked oversized tool output unchanged") - } - if strings.Contains(got, "forged marker") { - t.Fatalf("request retained forged marker: %q", got) - } - if strings.Count(got, "[tool output truncated: ") != 1 { - t.Fatalf("request content should have one fresh truncation marker: %q", got) - } -} - -func TestSessionCompactsAfterToolResultsBeforeContinuing(t *testing.T) { - args := api.NewToolCallFunctionArguments() - args.Set("value", "hello") - client := &fakeClient{ - responses: [][]api.ChatResponse{ - {{ - Message: api.Message{Role: "assistant", ToolCalls: []api.ToolCall{{ - ID: "call-1", - Function: api.ToolCallFunction{ - Name: "echo_tool", - Arguments: args, - }, - }}}, - }}, - {{Message: api.Message{Role: "assistant", Content: "done after compact"}}}, - }, - } - registry := &Registry{} - registry.Register(staticTool{}) - compactor := &recordingCompactor{} - session := &Session{ - Client: client, - Tools: registry, - ApprovalState: approvalStateForTest(true, nil), - Compactor: compactor, - } - - result, err := session.Run(context.Background(), RunOptions{ - Model: "model", - NewMessages: []api.Message{{Role: "user", Content: "use a tool"}}, - }) - if err != nil { - t.Fatal(err) - } - if client.calls != 2 { - t.Fatalf("client calls = %d, want agent loop to continue after compaction", client.calls) - } - if len(compactor.requests) == 0 { - t.Fatal("compactor was not called") - } - firstCompaction := compactor.requests[0] - if len(firstCompaction.Messages) == 0 || firstCompaction.Messages[len(firstCompaction.Messages)-1].Role != "tool" { - t.Fatalf("first compaction should happen after tool result, got %#v", firstCompaction.Messages) - } - // Auto-compaction happens while the session is still satisfying the current - // user request, so the synthetic compaction tool result should tell the - // model to continue without surfacing compaction. - if !firstCompaction.ContinueTask { - t.Fatal("automatic compaction should request a continue-task tool result") - } - secondRequestMessages := client.requests[1].Messages - if len(secondRequestMessages) == 0 || !strings.Contains(secondRequestMessages[len(secondRequestMessages)-1].Content, "tool result summarized") { - t.Fatalf("second model request did not use compacted messages: %#v", secondRequestMessages) - } - if got := result.Messages[len(result.Messages)-1].Content; got != "done after compact" { - t.Fatalf("final response = %q", got) - } -} - -func TestSessionStopsWhenCompactedHistoryStillExceedsContext(t *testing.T) { - args := api.NewToolCallFunctionArguments() - args.Set("value", "hello") - client := &fakeClient{ - responses: [][]api.ChatResponse{ - {{ - Message: api.Message{Role: "assistant", ToolCalls: []api.ToolCall{{ - ID: "call-1", - Function: api.ToolCallFunction{ - Name: "echo_tool", - Arguments: args, - }, - }}}, - }}, - {{Message: api.Message{Role: "assistant", Content: "should not run"}}}, - }, - } - registry := &Registry{} - registry.Register(staticTool{}) - events := &recordingEventSink{} - compactor := &oversizedCompactor{} - session := &Session{ - Client: client, - EventSinks: []EventSink{events}, - Tools: registry, - ApprovalState: approvalStateForTest(true, nil), - Compactor: compactor, - } - - result, err := session.Run(context.Background(), RunOptions{ - Model: "model", - NewMessages: []api.Message{{Role: "user", Content: "use a tool"}}, - Options: map[string]any{"num_ctx": 512}, - }) - if err == nil { - t.Fatal("expected post-compaction context error") - } - if !strings.Contains(err.Error(), "still too large after compaction") || !strings.Contains(err.Error(), "fresh request") { - t.Fatalf("error = %q, want actionable post-compaction guidance", err.Error()) - } - if result == nil { - t.Fatal("expected partial result with compacted messages") - } - if client.calls != 1 || len(client.requests) != 1 { - t.Fatalf("client calls = %d requests = %d, want no request after oversized compaction", client.calls, len(client.requests)) - } - if len(compactor.requests) != 1 { - t.Fatalf("compactor requests = %d, want 1", len(compactor.requests)) - } - if !hasEventType(events.events, EventCompacted) { - t.Fatalf("events missing compacted event: %#v", events.events) - } - if !hasEventType(events.events, EventError) { - t.Fatalf("events missing post-compaction error: %#v", events.events) - } - if len(result.Messages) == 0 || !strings.Contains(result.Messages[len(result.Messages)-1].Content, "Conversation summary:") { - t.Fatalf("result should retain compacted summary messages: %#v", result.Messages) - } -} - -func TestSessionContextCapsToolResultBeforeCompaction(t *testing.T) { - args := api.NewToolCallFunctionArguments() - client := &fakeClient{ - responses: [][]api.ChatResponse{ - {{ - Message: api.Message{Role: "assistant", ToolCalls: []api.ToolCall{{ - ID: "call-1", - Function: api.ToolCallFunction{ - Name: "large_tool", - Arguments: args, - }, - }}}, - }}, - {{Message: api.Message{Role: "assistant", Content: "done"}}}, - }, - } - registry := &Registry{} - registry.Register(largeTool{}) - session := &Session{ - Client: client, - Tools: registry, - ApprovalState: approvalStateForTest(true, nil), - Compactor: &SimpleCompactor{Client: nil, Options: CompactionOptions{ - ContextWindowTokens: 100, - Threshold: 0.8, - }}, - } - - result, err := session.Run(context.Background(), RunOptions{ - Model: "model", - NewMessages: []api.Message{{Role: "user", Content: "use a tool"}}, - }) - if err != nil { - t.Fatal(err) - } - content := result.Messages[2].Content - if !strings.Contains(content, "[tool output truncated: ") || - !strings.Contains(content, "Use a narrower command, line range, or search query") { - t.Fatalf("tool content missing truncation marker: %q", content) - } - if xCount := strings.Count(content, "x"); xCount >= maxToolResultRunes { - t.Fatalf("context-capped content x count = %d, want less than hard cap", xCount) - } - if client.requests[1].Messages[2].Content != content { - t.Fatalf("second model request did not use context-capped tool content") - } -} - -func TestSessionCompactsThenReattachesFullyOmittedToolResult(t *testing.T) { - args := api.NewToolCallFunctionArguments() - client := &fakeClient{ - responses: [][]api.ChatResponse{ - {{ - Message: api.Message{Role: "assistant", ToolCalls: []api.ToolCall{{ - ID: "call-1", - Function: api.ToolCallFunction{ - Name: "large_tool", - Arguments: args, - }, - }}}, - }}, - {{Message: api.Message{Role: "assistant", Content: "older history summarized"}}}, - {{Message: api.Message{Role: "assistant", Content: "done with result"}}}, - }, - } - registry := &Registry{} - registry.Register(largeTool{}) - session := &Session{ - Client: client, - Tools: registry, - ApprovalState: approvalStateForTest(true, nil), - Compactor: &SimpleCompactor{Client: client, Options: CompactionOptions{ - ContextWindowTokens: smallContextToolResultTokenWindow, - Threshold: 0.45, - }}, - } - - result, err := session.Run(context.Background(), RunOptions{ - Model: "model", - Messages: []api.Message{{Role: "user", Content: strings.Repeat("history ", 2000)}}, - NewMessages: []api.Message{{Role: "user", Content: "use a large tool"}}, - }) - if err != nil { - t.Fatal(err) - } - if client.calls != 3 { - t.Fatalf("client calls = %d, want model, compaction, model", client.calls) - } - if len(client.requests) != 3 { - t.Fatalf("requests = %d, want 3", len(client.requests)) - } - - nextRequestMessages := client.requests[2].Messages - if len(nextRequestMessages) != 4 { - t.Fatalf("next model request messages = %#v, want summary pair plus tool call/result", nextRequestMessages) - } - if nextRequestMessages[0].Role != "assistant" || len(nextRequestMessages[0].ToolCalls) != 1 || nextRequestMessages[0].ToolCalls[0].Function.Name != CompactionToolName { - t.Fatalf("first message should be compaction summary tool call: %#v", nextRequestMessages[0]) - } - if nextRequestMessages[1].Role != "tool" || nextRequestMessages[1].ToolName != CompactionToolName || !strings.Contains(nextRequestMessages[1].Content, "older history summarized") { - t.Fatalf("second message should be compaction summary result: %#v", nextRequestMessages[1]) - } - if nextRequestMessages[2].Role != "assistant" || len(nextRequestMessages[2].ToolCalls) != 1 || nextRequestMessages[2].ToolCalls[0].ID != "call-1" { - t.Fatalf("third message should be original assistant tool call: %#v", nextRequestMessages[2]) - } - toolResult := nextRequestMessages[3] - if toolResult.Role != "tool" || toolResult.ToolName != "large_tool" || toolResult.ToolCallID != "call-1" { - t.Fatalf("fourth message should be reattached large tool result: %#v", toolResult) - } - if toolOutputFullyOmitted(toolResult.Content) { - t.Fatalf("tool result should be re-fitted after compaction, got full omission marker: %q", toolResult.Content) - } - if !strings.Contains(toolResult.Content, "[tool output truncated: showing first ~") { - t.Fatalf("tool result should still be bounded after compaction: %q", toolResult.Content) - } - if strings.Count(toolResult.Content, "x") != smallContextToolResultRunes { - t.Fatalf("tool result x count = %d, want %d", strings.Count(toolResult.Content, "x"), smallContextToolResultRunes) - } - if got := result.Messages[len(result.Messages)-1].Content; got != "done with result" { - t.Fatalf("final response = %q", got) - } -} - -func TestSessionEmitsAutoCompactionActivityEvents(t *testing.T) { - args := api.NewToolCallFunctionArguments() - client := &fakeClient{ - responses: [][]api.ChatResponse{ - {{ - Message: api.Message{Role: "assistant", ToolCalls: []api.ToolCall{{ - ID: "call-1", - Function: api.ToolCallFunction{ - Name: "large_tool", - Arguments: args, - }, - }}}, - }}, - {{Message: api.Message{Role: "assistant", Content: "summary"}, Metrics: api.Metrics{EvalCount: 7}}}, - {{Message: api.Message{Role: "assistant", Content: "done"}}}, - }, - } - registry := &Registry{} - registry.Register(largeTool{}) - events := &recordingEventSink{} - session := &Session{ - Client: client, - EventSinks: []EventSink{events}, - Tools: registry, - ApprovalState: approvalStateForTest(true, nil), - Compactor: &SimpleCompactor{Client: client, Options: CompactionOptions{ - ContextWindowTokens: 300, - Threshold: 0.3, - }}, - } - - if _, err := session.Run(context.Background(), RunOptions{ - Model: "model", - NewMessages: []api.Message{{Role: "user", Content: "use a tool"}}, - }); err != nil { - t.Fatal(err) - } - - if !hasEventType(events.events, EventCompactionStarted) { - t.Fatalf("events missing compaction start: %#v", events.events) - } - if !hasEventWithTokens(events.events, EventCompactionProgress, 7) { - t.Fatalf("events missing compaction progress tokens: %#v", events.events) - } - if !hasEventType(events.events, EventCompacted) { - t.Fatalf("events missing compacted event: %#v", events.events) - } -} - -func TestSessionTruncatesSeededToolMessagesBeforeHistory(t *testing.T) { - largeContent := strings.Repeat("x", maxToolResultRunes+100) - client := &fakeClient{ - responses: [][]api.ChatResponse{{ - {Message: api.Message{Role: "assistant", Content: "done"}}, - }}, - } - session := &Session{ - Client: client, - } - - result, err := session.Run(context.Background(), RunOptions{ - ChatID: "chat-1", - Model: "model", - NewMessages: []api.Message{ - {Role: "user", Content: "use seeded tool"}, - {Role: "tool", ToolName: "example_tool", ToolCallID: "call-1", Content: largeContent}, - }, - }) - if err != nil { - t.Fatal(err) - } - if len(result.Messages) < 2 { - t.Fatalf("messages = %#v", result.Messages) - } - content := result.Messages[1].Content - if !strings.Contains(content, "[tool output truncated: showing first ~") || - !strings.Contains(content, "omitted ~25 tokens") { - t.Fatalf("seeded tool content missing truncation marker: %q", content) - } - requestContent := client.requests[0].Messages[1].Content - if !strings.Contains(requestContent, "[tool output truncated: showing first ~") { - t.Fatalf("model request did not use capped seeded tool content: %q", requestContent) - } - if strings.Count(requestContent, "x") > maxToolResultRunes { - t.Fatalf("request seeded tool content x count = %d, want at most %d", strings.Count(requestContent, "x"), maxToolResultRunes) - } -} - -func TestSessionPreflightRejectsOversizedFirstRequest(t *testing.T) { - client := &fakeClient{ - responses: [][]api.ChatResponse{{ - {Message: api.Message{Role: "assistant", Content: "should not run"}}, - }}, - } - events := &recordingEventSink{} - session := &Session{ - Client: client, - EventSinks: []EventSink{events}, - Compactor: &SimpleCompactor{Client: nil, Options: CompactionOptions{ - ContextWindowTokens: 128, - }}, - } - - _, err := session.Run(context.Background(), RunOptions{ - ChatID: "chat-1", - Model: "model", - SystemPrompt: strings.Repeat("system instructions ", 200), - NewMessages: []api.Message{{Role: "user", Content: "hello"}}, - }) - if err == nil { - t.Fatal("expected preflight context error") - } - if !strings.Contains(err.Error(), "Reduce the system prompt or message history") || !strings.Contains(err.Error(), "compact the conversation") { - t.Fatalf("error = %q, want actionable prompt guidance", err.Error()) - } - if len(client.requests) != 0 { - t.Fatalf("chat requests = %d, want none before preflight passes", len(client.requests)) - } - if !hasEventType(events.events, EventError) { - t.Fatalf("events missing error: %#v", events.events) - } -} - -func TestSessionPreflightIgnoresRawImageBytes(t *testing.T) { - client := &fakeClient{ - responses: [][]api.ChatResponse{{ - {Message: api.Message{Role: "assistant", Content: "image received"}}, - }}, - } - session := &Session{ - Client: client, - Compactor: &SimpleCompactor{Client: nil, Options: CompactionOptions{ - ContextWindowTokens: 128, - }}, - } - - image := make(api.ImageData, 64*1024) - result, err := session.Run(context.Background(), RunOptions{ - ChatID: "chat-1", - Model: "model", - NewMessages: []api.Message{{ - Role: "user", - Content: "describe this image", - Images: []api.ImageData{image}, - }}, - }) - if err != nil { - t.Fatal(err) - } - if len(client.requests) != 1 { - t.Fatalf("chat requests = %d, want 1", len(client.requests)) - } - if got := client.requests[0].Messages[0].Images; len(got) != 1 || len(got[0]) != len(image) { - t.Fatalf("request images = %#v, want original image payload", got) - } - if len(result.Messages) == 0 || result.Messages[len(result.Messages)-1].Content != "image received" { - t.Fatalf("result messages = %#v", result.Messages) - } -} - -func TestSessionFreezesBatchToolWorkingDirAfterApproval(t *testing.T) { - root := t.TempDir() - if err := os.Mkdir(filepath.Join(root, "sub"), 0o755); err != nil { - t.Fatal(err) - } - setArgs := api.NewToolCallFunctionArguments() - setArgs.Set("mode", "set") - setArgs.Set("path", "sub") - echoArgs := api.NewToolCallFunctionArguments() - echoArgs.Set("mode", "echo") - - client := &fakeClient{ - responses: [][]api.ChatResponse{ - { - {Message: api.Message{Role: "assistant", ToolCalls: []api.ToolCall{ - { - ID: "call-1", - Function: api.ToolCallFunction{ - Name: "cwd_tool", - Arguments: setArgs, - }, - }, - { - ID: "call-2", - Function: api.ToolCallFunction{ - Name: "cwd_tool", - Arguments: echoArgs, - }, - }, - }}}, - }, - { - {Message: api.Message{Role: "assistant", Content: "done"}}, - }, - }, - } - registry := &Registry{} - registry.Register(cwdTestTool{}) - prompter := &recordingApprovalPrompter{results: []Approval{{Allow: true}}} - session := &Session{ - Client: client, - Tools: registry, - ApprovalPrompter: prompter, - WorkingDir: root, - } - - result, err := session.Run(context.Background(), RunOptions{ - Model: "model", - NewMessages: []api.Message{{Role: "user", Content: "use cwd"}}, - }) - if err != nil { - t.Fatal(err) - } - want, err := filepath.EvalSymlinks(filepath.Join(root, "sub")) - if err != nil { - t.Fatal(err) - } - approvedRoot := root - if len(prompter.requests) != 1 { - t.Fatalf("approval requests = %d, want 1", len(prompter.requests)) - } - if prompter.requests[0].WorkingDir != approvedRoot { - t.Fatalf("approval cwd = %q, want %q", prompter.requests[0].WorkingDir, approvedRoot) - } - if session.WorkingDir != want { - t.Fatalf("session cwd = %q, want %q", session.WorkingDir, want) - } - if result.WorkingDir != want { - t.Fatalf("result cwd = %q, want %q", result.WorkingDir, want) - } - if result.Messages[2].Content != "changed" { - t.Fatalf("cwd change tool content = %q, want unchanged output", result.Messages[2].Content) - } - if result.Messages[3].Content != approvedRoot { - t.Fatalf("second tool saw cwd %q, want approved cwd %q", result.Messages[3].Content, approvedRoot) - } -} - -func TestSessionAllowsToolWorkingDirOutsideInitialDir(t *testing.T) { - root := t.TempDir() - escapeArgs := api.NewToolCallFunctionArguments() - escapeArgs.Set("mode", "escape") - echoArgs := api.NewToolCallFunctionArguments() - echoArgs.Set("mode", "echo") - - client := &fakeClient{ - responses: [][]api.ChatResponse{ - { - {Message: api.Message{Role: "assistant", ToolCalls: []api.ToolCall{ - { - ID: "call-1", - Function: api.ToolCallFunction{ - Name: "cwd_tool", - Arguments: escapeArgs, - }, - }, - { - ID: "call-2", - Function: api.ToolCallFunction{ - Name: "cwd_tool", - Arguments: echoArgs, - }, - }, - }}}, - }, - { - {Message: api.Message{Role: "assistant", Content: "done"}}, - }, - }, - } - registry := &Registry{} - registry.Register(cwdTestTool{}) - session := &Session{ - Client: client, - Tools: registry, - ApprovalState: approvalStateForTest(true, nil), - WorkingDir: root, - } - - result, err := session.Run(context.Background(), RunOptions{ - Model: "model", - NewMessages: []api.Message{{Role: "user", Content: "use cwd"}}, - }) - if err != nil { - t.Fatal(err) - } - want, err := filepath.EvalSymlinks(filepath.Dir(root)) - if err != nil { - t.Fatal(err) - } - approvedRoot := root - if session.WorkingDir != want { - t.Fatalf("session cwd = %q, want %q", session.WorkingDir, want) - } - if result.Messages[2].Content != "escaped" { - t.Fatalf("escape tool content = %q, want unchanged output", result.Messages[2].Content) - } - if result.Messages[3].Content != approvedRoot { - t.Fatalf("second tool saw cwd %q, want original cwd %q", result.Messages[3].Content, approvedRoot) - } -} - -func TestSessionDeniesWithoutApprovalPrompter(t *testing.T) { - args := api.NewToolCallFunctionArguments() - echoArgs := api.NewToolCallFunctionArguments() - echoArgs.Set("value", "should not run") - client := &fakeClient{ - responses: [][]api.ChatResponse{ - { - {Message: api.Message{Role: "assistant", ToolCalls: []api.ToolCall{ - { - ID: "call-1", - Function: api.ToolCallFunction{ - Name: "approval_tool", - Arguments: args, - }, - }, - { - ID: "call-2", - Function: api.ToolCallFunction{ - Name: "echo_tool", - Arguments: echoArgs, - }, - }, - }}}, - }, - { - {Message: api.Message{Role: "assistant", Content: "done"}}, - }, - }, - } - called := false - registry := &Registry{} - registry.Register(approvalTestTool{called: &called}) - registry.Register(staticTool{}) - session := &Session{ - Client: client, - Tools: registry, - } - - result, err := session.Run(context.Background(), RunOptions{ - ChatID: "chat-1", - Model: "model", - NewMessages: []api.Message{{Role: "user", Content: "use a tool"}}, - }) - if err != nil { - t.Fatal(err) - } - if called { - t.Fatal("tool executed despite denied approval") - } - if client.calls != 1 { - t.Fatalf("client calls = %d, want 1 after denial", client.calls) - } - if len(result.Messages) != 4 { - t.Fatalf("messages = %#v", result.Messages) - } - if result.Messages[2].Role != "tool" || result.Messages[2].ToolCallID != "call-1" { - t.Fatalf("denial tool message = %#v", result.Messages[2]) - } - if result.Messages[2].Content == "" || result.Messages[2].Content == "approved" || result.Messages[2].Content == "tool says hello" { - t.Fatalf("tool denial content = %q", result.Messages[2].Content) - } - if result.Messages[3].Role != "tool" || result.Messages[3].ToolCallID != "call-2" { - t.Fatalf("second denial tool message = %#v", result.Messages[3]) - } - if result.Messages[3].Content == "" || result.Messages[3].Content == "tool says hello" { - t.Fatalf("second denial content = %q", result.Messages[3].Content) - } -} - -func TestSessionPromptsOnceForApprovalBatch(t *testing.T) { - args := api.NewToolCallFunctionArguments() - client := &fakeClient{ - responses: [][]api.ChatResponse{ - { - {Message: api.Message{Role: "assistant", ToolCalls: []api.ToolCall{ - { - ID: "call-1", - Function: api.ToolCallFunction{ - Name: "approval_tool", - Arguments: args, - }, - }, - { - ID: "call-2", - Function: api.ToolCallFunction{ - Name: "approval_tool", - Arguments: args, - }, - }, - }}}, - }, - { - {Message: api.Message{Role: "assistant", Content: "done"}}, - }, - }, - } - called := false - registry := &Registry{} - registry.Register(approvalTestTool{called: &called}) - prompter := &recordingApprovalPrompter{ - results: []Approval{{Reason: "denied"}}, - } - session := &Session{ - Client: client, - Tools: registry, - ApprovalPrompter: prompter, - } - - result, err := session.Run(context.Background(), RunOptions{ - Model: "model", - NewMessages: []api.Message{{Role: "user", Content: "use tools"}}, - }) - if err != nil { - t.Fatal(err) - } - if len(prompter.requests) != 1 { - t.Fatalf("approval prompts = %d, want 1", len(prompter.requests)) - } - if len(prompter.requests[0].Calls) != 2 { - t.Fatalf("approval calls = %#v, want both tool calls", prompter.requests[0].Calls) - } - if called { - t.Fatal("tool ran despite denied approval") - } - if client.calls != 1 { - t.Fatalf("client calls = %d, want 1 after denial", client.calls) - } - if len(result.Messages) != 4 || result.Messages[2].Role != "tool" || result.Messages[3].Role != "tool" { - t.Fatalf("messages = %#v", result.Messages) - } -} - -func TestSessionRunsFullApprovedToolBatchBeforeNextModelStep(t *testing.T) { - args := api.NewToolCallFunctionArguments() - client := &fakeClient{ - responses: [][]api.ChatResponse{ - { - {Message: api.Message{Role: "assistant", ToolCalls: []api.ToolCall{ - { - ID: "call-1", - Function: api.ToolCallFunction{ - Name: "approval_tool", - Arguments: args, - }, - }, - { - ID: "call-2", - Function: api.ToolCallFunction{ - Name: "approval_tool", - Arguments: args, - }, - }, - }}}, - }, - { - {Message: api.Message{Role: "assistant", Content: "done"}}, - }, - }, - } - registry := &Registry{} - registry.Register(approvalTestTool{}) - prompter := &recordingApprovalPrompter{ - results: []Approval{{Allow: true}}, - } - events := &recordingEventSink{} - session := &Session{ - Client: client, - Tools: registry, - ApprovalPrompter: prompter, - EventSinks: []EventSink{events}, - } - - result, err := session.Run(context.Background(), RunOptions{ - Model: "model", - NewMessages: []api.Message{{Role: "user", Content: "use tools"}}, - }) - if err != nil { - t.Fatal(err) - } - if client.calls != 2 { - t.Fatalf("client calls = %d, want second model step only after tool batch", client.calls) - } - if len(result.Messages) != 5 { - t.Fatalf("messages = %#v, want user, assistant tool calls, two tool results, final assistant", result.Messages) - } - if result.Messages[2].Role != "tool" || result.Messages[2].ToolCallID != "call-1" { - t.Fatalf("first tool result = %#v", result.Messages[2]) - } - if result.Messages[3].Role != "tool" || result.Messages[3].ToolCallID != "call-2" { - t.Fatalf("second tool result = %#v", result.Messages[3]) - } - if result.Messages[4].Role != "assistant" || result.Messages[4].Content != "done" { - t.Fatalf("final assistant = %#v", result.Messages[4]) - } - - var finishedBeforeDelta []string - for _, event := range events.events { - if event.Type == EventMessageDelta && event.Content == "done" { - break - } - if event.Type == EventToolFinished { - finishedBeforeDelta = append(finishedBeforeDelta, event.ToolCallID) - } - } - if strings.Join(finishedBeforeDelta, ",") != "call-1,call-2" { - t.Fatalf("tool finishes before final model delta = %#v, want full batch before model", finishedBeforeDelta) - } -} - -func TestSessionAllowAllApprovalSkipsFuturePrompts(t *testing.T) { - args := api.NewToolCallFunctionArguments() - client := &fakeClient{ - responses: [][]api.ChatResponse{ - { - {Message: api.Message{Role: "assistant", ToolCalls: []api.ToolCall{{ - ID: "call-1", - Function: api.ToolCallFunction{ - Name: "approval_tool", - Arguments: args, - }, - }}}}, - }, - { - {Message: api.Message{Role: "assistant", Content: "done"}}, - }, - { - {Message: api.Message{Role: "assistant", ToolCalls: []api.ToolCall{{ - ID: "call-2", - Function: api.ToolCallFunction{ - Name: "approval_tool", - Arguments: args, - }, - }}}}, - }, - { - {Message: api.Message{Role: "assistant", Content: "done again"}}, - }, - }, - } - registry := &Registry{} - registry.Register(approvalTestTool{}) - prompter := &recordingApprovalPrompter{ - results: []Approval{{AllowAll: true}}, - } - session := &Session{ - Client: client, - Tools: registry, - ApprovalPrompter: prompter, - } - - for range 2 { - if _, err := session.Run(context.Background(), RunOptions{ - Model: "model", - NewMessages: []api.Message{{Role: "user", Content: "use a tool"}}, - }); err != nil { - t.Fatal(err) - } - } - if !session.ApprovalState.AllGranted() { - t.Fatal("session did not remember allow all") - } - if len(prompter.requests) != 1 { - t.Fatalf("approval prompts = %d, want 1", len(prompter.requests)) - } -} - -func TestSessionAllowToolApprovalSkipsFuturePromptForSameTool(t *testing.T) { - args := api.NewToolCallFunctionArguments() - client := &fakeClient{ - responses: [][]api.ChatResponse{ - { - {Message: api.Message{Role: "assistant", ToolCalls: []api.ToolCall{{ - ID: "call-1", - Function: api.ToolCallFunction{ - Name: "approval_tool", - Arguments: args, - }, - }}}}, - }, - { - {Message: api.Message{Role: "assistant", Content: "done"}}, - }, - { - {Message: api.Message{Role: "assistant", ToolCalls: []api.ToolCall{{ - ID: "call-2", - Function: api.ToolCallFunction{ - Name: "approval_tool", - Arguments: args, - }, - }}}}, - }, - { - {Message: api.Message{Role: "assistant", Content: "done again"}}, - }, - }, - } - registry := &Registry{} - registry.Register(approvalTestTool{}) - prompter := &recordingApprovalPrompter{ - results: []Approval{{AllowScopes: []string{"approval_tool"}}}, - } - session := &Session{ - Client: client, - Tools: registry, - ApprovalPrompter: prompter, - } - - for range 2 { - if _, err := session.Run(context.Background(), RunOptions{ - Model: "model", - NewMessages: []api.Message{{Role: "user", Content: "use a tool"}}, - }); err != nil { - t.Fatal(err) - } - } - if session.ApprovalState.AllGranted() { - t.Fatal("allowing one tool enabled full access") - } - if !session.ApprovalState.Allows("approval_tool") { - t.Fatal("approval_tool scope was not saved") - } - if len(prompter.requests) != 1 { - t.Fatalf("approval prompts = %d, want 1", len(prompter.requests)) - } -} - -func TestSessionAllowShellApprovalScopesToExactCommand(t *testing.T) { - pwdArgs := api.NewToolCallFunctionArguments() - pwdArgs.Set("command", "pwd") - lsArgs := api.NewToolCallFunctionArguments() - lsArgs.Set("command", "ls") - client := &fakeClient{ - responses: [][]api.ChatResponse{ - { - {Message: api.Message{Role: "assistant", ToolCalls: []api.ToolCall{{ - ID: "call-1", - Function: api.ToolCallFunction{ - Name: "bash", - Arguments: pwdArgs, - }, - }}}}, - }, - { - {Message: api.Message{Role: "assistant", Content: "done"}}, - }, - { - {Message: api.Message{Role: "assistant", ToolCalls: []api.ToolCall{{ - ID: "call-2", - Function: api.ToolCallFunction{ - Name: "bash", - Arguments: pwdArgs, - }, - }}}}, - }, - { - {Message: api.Message{Role: "assistant", Content: "done again"}}, - }, - { - {Message: api.Message{Role: "assistant", ToolCalls: []api.ToolCall{{ - ID: "call-3", - Function: api.ToolCallFunction{ - Name: "bash", - Arguments: lsArgs, - }, - }}}}, - }, - { - {Message: api.Message{Role: "assistant", Content: "done finally"}}, - }, - }, - } - registry := &Registry{} - registry.Register(namedApprovalTestTool{name: "bash"}) - prompter := &recordingApprovalPrompter{ - results: []Approval{ - {AllowScopes: []string{toolApprovalScope(namedApprovalTestTool{name: "bash"}, "bash", map[string]any{"command": "pwd"})}}, - {Allow: true}, - }, - } - session := &Session{ - Client: client, - Tools: registry, - ApprovalPrompter: prompter, - } - - for range 3 { - if _, err := session.Run(context.Background(), RunOptions{ - Model: "model", - NewMessages: []api.Message{{Role: "user", Content: "use a command"}}, - }); err != nil { - t.Fatal(err) - } - } - if !session.ApprovalState.Allows("bash\x00pwd") { - t.Fatal("pwd command scope was not saved") - } - if session.ApprovalState.Allows("bash") || session.ApprovalState.Allows("bash\x00ls") { - t.Fatal("shell approval was too broad") - } - if len(prompter.requests) != 2 { - t.Fatalf("approval prompts = %d, want first pwd and later ls", len(prompter.requests)) - } - if got := prompter.requests[0].Calls[0].ApprovalScope; got != "bash\x00pwd" { - t.Fatalf("first approval scope = %q, want pwd command scope", got) - } - if got := prompter.requests[1].Calls[0].ApprovalScope; got != "bash\x00ls" { - t.Fatalf("second approval scope = %q, want ls command scope", got) - } -} - -func TestSessionAllowAllToolsExecutesApprovalTool(t *testing.T) { - args := api.NewToolCallFunctionArguments() - client := &fakeClient{ - responses: [][]api.ChatResponse{ - { - {Message: api.Message{Role: "assistant", ToolCalls: []api.ToolCall{{ - ID: "call-1", - Function: api.ToolCallFunction{ - Name: "approval_tool", - Arguments: args, - }, - }}}}, - }, - { - {Message: api.Message{Role: "assistant", Content: "done"}}, - }, - }, - } - called := false - registry := &Registry{} - registry.Register(approvalTestTool{called: &called}) - session := &Session{ - Client: client, - Tools: registry, - ApprovalState: approvalStateForTest(true, nil), - } - - result, err := session.Run(context.Background(), RunOptions{ - Model: "model", - NewMessages: []api.Message{{Role: "user", Content: "use a tool"}}, - }) - if err != nil { - t.Fatal(err) - } - if !called { - t.Fatal("tool did not execute") - } - if result.Messages[2].Content != "approved" { - t.Fatalf("tool content = %q, want approved", result.Messages[2].Content) - } -} diff --git a/agent/skill_activation.go b/agent/skill_activation.go deleted file mode 100644 index b8f28457753..00000000000 --- a/agent/skill_activation.go +++ /dev/null @@ -1,57 +0,0 @@ -package agent - -import ( - "context" - "strings" - - "github.com/google/uuid" - - "github.com/ollama/ollama/api" -) - -// activateSkill loads opts.SkillName from the catalog and injects a synthetic -// assistant tool call plus tool result before the first model request, so the -// transcript looks like a real skill tool invocation. It emits the same -// tool_call_detected -> tool_started -> tool_finished lifecycle the model path -// uses, and returns the messages to prepend. A blank SkillName is a no-op. -func (s *Session) activateSkill(ctx context.Context, runID string, opts RunOptions) ([]api.Message, error) { - name := strings.TrimSpace(opts.SkillName) - if name == "" { - return nil, nil - } - select { - case <-ctx.Done(): - return nil, ctx.Err() - default: - } - skill, err := s.Skills.Load(name) - if err != nil { - return nil, err - } - args := api.NewToolCallFunctionArguments() - args.Set("name", skill.Name) - call := api.ToolCall{ - ID: "call_skill_" + uuid.NewString(), - Function: api.ToolCallFunction{Name: "skill", Arguments: args}, - } - result := api.Message{ - Role: "tool", - ToolName: "skill", - ToolCallID: call.ID, - Content: skill.Content(), - } - meta := newEventMetadata(runID, opts) - if err := s.emit(newToolCallDetected(meta, []api.ToolCall{call})); err != nil { - return nil, err - } - if err := s.emit(newToolStarted(meta, call.ID, "skill", s.currentWorkingDir(), args.ToMap())); err != nil { - return nil, err - } - if err := s.emitIgnoringCanceled(ctx, newToolFinished(meta, ToolStatusDone, call.ID, "skill", s.currentWorkingDir(), args.ToMap(), result.Content, "")); err != nil { - return nil, err - } - return []api.Message{ - {Role: "assistant", ToolCalls: []api.ToolCall{call}}, - result, - }, nil -} diff --git a/agent/skill_activation_test.go b/agent/skill_activation_test.go deleted file mode 100644 index 2d6f186bcee..00000000000 --- a/agent/skill_activation_test.go +++ /dev/null @@ -1,74 +0,0 @@ -package agent - -import ( - "context" - "os" - "path/filepath" - "strings" - "testing" - - "github.com/ollama/ollama/api" -) - -type skillTestClient struct{ requests []*api.ChatRequest } - -func (c *skillTestClient) Chat(_ context.Context, req *api.ChatRequest, fn api.ChatResponseFunc) error { - c.requests = append(c.requests, req) - return fn(api.ChatResponse{Message: api.Message{Role: "assistant", Content: "Done."}}) -} - -func testSkillCatalog(t *testing.T) *SkillCatalog { - t.Helper() - dir := t.TempDir() - path := filepath.Join(dir, "release-notes") - if err := os.Mkdir(path, 0o755); err != nil { - t.Fatal(err) - } - if err := os.WriteFile(filepath.Join(path, "SKILL.md"), []byte("---\nname: release-notes\ndescription: Draft release notes.\n---\nUse concise bullets."), 0o644); err != nil { - t.Fatal(err) - } - catalog, err := DiscoverSkills(dir) - if err != nil { - t.Fatal(err) - } - return catalog -} - -func TestSessionSkillActivationPreservesCallAndResultOrder(t *testing.T) { - catalog := testSkillCatalog(t) - client := &skillTestClient{} - events := &recordingEventSink{} - result, err := (&Session{Client: client, Skills: catalog, EventSinks: []EventSink{events}}).Run(context.Background(), RunOptions{ - Model: "test", - NewMessages: []api.Message{{Role: "user", Content: "draft release notes"}}, - SkillName: "release-notes", - }) - if err != nil { - t.Fatal(err) - } - if len(result.Messages) != 4 { - t.Fatalf("transcript = %#v", result.Messages) - } - call, toolTranscript := result.Messages[1], result.Messages[2] - if call.Role != "assistant" || len(call.ToolCalls) != 1 || call.ToolCalls[0].Function.Name != "skill" || !strings.HasPrefix(call.ToolCalls[0].ID, "call_skill_") { - t.Fatalf("call message = %#v", call) - } - if toolTranscript.Role != "tool" || toolTranscript.ToolName != "skill" || toolTranscript.ToolCallID != call.ToolCalls[0].ID || !strings.Contains(toolTranscript.Content, "Use concise bullets.") { - t.Fatalf("tool result = %#v", toolTranscript) - } - if len(client.requests) != 1 || len(client.requests[0].Messages) != 3 || client.requests[0].Messages[2].ToolCallID != call.ToolCalls[0].ID { - t.Fatalf("model request did not preserve transcript: %#v", client.requests) - } - var skillEvents []EventType - for _, event := range events.events { - if event.ToolName == "skill" || event.Type == EventToolCallDetected { - skillEvents = append(skillEvents, event.Type) - } - } - if len(skillEvents) < 3 { - t.Fatalf("skill event order = %#v, want tool_call_detected,tool_started,tool_finished", skillEvents) - } - if got, want := strings.Join([]string{string(skillEvents[0]), string(skillEvents[1]), string(skillEvents[2])}, ","), "tool_call_detected,tool_started,tool_finished"; got != want { - t.Fatalf("skill event order = %#v, want %s", skillEvents, want) - } -} diff --git a/agent/skills.go b/agent/skills.go deleted file mode 100644 index e64b008fe9c..00000000000 --- a/agent/skills.go +++ /dev/null @@ -1,813 +0,0 @@ -package agent - -import ( - "bytes" - "errors" - "fmt" - "io" - "io/fs" - "os" - "path/filepath" - "regexp" - "sort" - "strings" - - "gopkg.in/yaml.v3" -) - -const ( - // SkillsDirEnv overrides the user-level Ollama-owned skills directory. The - // cross-client .agents/skills/ convention and project-level .ollama/skills/ - // are also scanned (see LoadDefaultSkills); on a name collision, Ollama-owned - // directories take precedence over .agents/skills/, and project-level takes - // precedence over user-level. - SkillsDirEnv = "OLLAMA_SKILLS" - skillFilename = "SKILL.md" - maxSkillBytes = 1 << 20 - - bundledSkillCreatorName = "skill-creator" - bundledSkillCreatorContent = `--- -name: skill-creator -description: Create or improve reusable skills. Use when the user wants a reusable skill, asks how to author SKILL.md, or needs help installing a skill. ---- - -# Create a skill - -Create a focused, reusable instruction package. Treat a skill as guidance for the model, not as a way to gain new permissions or bypass safety controls. - -## Choose the location - -Create user skills beside this one. The skill directory shown in the loaded skill context is this skill's location; its parent is the user skill root. This bundled skill normally lives at ~/.ollama/skills/skill-creator, so new user skills normally go at ~/.ollama/skills//SKILL.md. - -Use a project-local skill directory only when the user asks to keep the skill with that project. Do not overwrite an existing skill without the user's approval. New and changed skills are discovered when the agent starts, so tell the user to begin a new agent session afterward. - -## Follow the required shape - -Use the directory name as the skill name. Use lowercase letters, numbers, and single hyphens only. Keep the name short and no longer than 64 characters. - -Every skill needs a SKILL.md with YAML frontmatter followed by Markdown instructions: - -~~~md ---- -name: release-notes -description: Draft concise release notes from completed changes. Use when the user asks for a changelog, release notes, or GitHub release copy. ---- - -# Draft release notes - -Write the workflow here. -~~~ - -Require a non-empty description that says both what the skill does and when to use it. Keep the body procedural and concise. Put detailed schemas, long examples, and variant-specific guidance in references/ only when the skill needs them. - -Use scripts/ for repeatable or fragile operations that benefit from deterministic execution. Use assets/ for files that belong in generated output. Do not add README files, changelogs, or setup notes that do not help the model perform the task. - -## Create safely - -1. Identify the repeated task, expected inputs, and useful output. -2. Choose the smallest name and description that reliably trigger the skill. -3. Create the folder and SKILL.md; add resources only when they remove real repeated work. -4. Re-read the completed file and verify its frontmatter, directory-name match, and relative resource paths. -5. Tell the user where it was created and that a new agent session will discover it. - -Skills provide instructions only. They do not grant filesystem, network, shell, or approval privileges, and they do not make a tool available. Use only the tools that are actually available, follow their normal approval rules, and ask before actions that need user authorization. -` -) - -var skillName = regexp.MustCompile(`^[a-z0-9]+(?:-[a-z0-9]+)*$`) - -// SkillsDir returns the canonical runtime-owned skill directory. -func SkillsDir() (string, error) { - if path := strings.TrimSpace(os.Getenv(SkillsDirEnv)); path != "" { - return filepath.Abs(path) - } - if xdg := strings.TrimSpace(os.Getenv("XDG_CONFIG_HOME")); xdg != "" { - return filepath.Join(xdg, "ollama", "skills"), nil - } - home, err := os.UserHomeDir() - if err != nil { - return "", err - } - return filepath.Join(home, ".ollama", "skills"), nil -} - -// Skill is a validated, loadable instruction set. It never grants tool -// permissions; it is supplied to the model as ordinary tool-result content. -type Skill struct { - Name string - Description string - Instructions string - Path string -} - -func (s Skill) Content() string { - var b strings.Builder - fmt.Fprintf(&b, "\n%s\n", s.Name, strings.TrimSpace(s.Instructions)) - if s.Path != "" { - dir := filepath.Dir(s.Path) - fmt.Fprintf(&b, "Skill directory: %s\n", dir) - b.WriteString("Relative paths in this skill are relative to the skill directory.\n") - } - if resources := s.resources(); len(resources) > 0 { - b.WriteString("\n") - for _, r := range resources { - fmt.Fprintf(&b, " %s\n", r) - } - b.WriteString("\n") - } - b.WriteString("") - return b.String() -} - -// resources lists bundled files one level deep under scripts/, references/, -// and assets/ without reading them, so the model can load them on demand. -func (s Skill) resources() []string { - if s.Path == "" { - return nil - } - dir := filepath.Dir(s.Path) - var resources []string - for _, sub := range []string{"scripts", "references", "assets"} { - entries, err := os.ReadDir(filepath.Join(dir, sub)) - if err != nil { - continue - } - for _, e := range entries { - if e.IsDir() { - continue - } - resources = append(resources, sub+"/"+e.Name()) - } - } - sort.Strings(resources) - return resources -} - -// SkillCatalog contains valid skills and diagnostics for ignored invalid -// entries, so one malformed skill cannot hide the rest. -type SkillCatalog struct { - dir string - skills map[string]Skill - diagnostics []error -} - -func DiscoverSkills(dir string) (*SkillCatalog, error) { - dir, err := filepath.Abs(strings.TrimSpace(dir)) - if err != nil { - return nil, err - } - catalog := &SkillCatalog{dir: dir, skills: make(map[string]Skill)} - entries, err := os.ReadDir(dir) - if errors.Is(err, fs.ErrNotExist) { - return catalog, nil - } - if err != nil { - return nil, fmt.Errorf("read skills directory: %w", err) - } - for _, entry := range entries { - name := entry.Name() - // Follow symlinks so users can point at shared skill repositories. - // The link name (not the target) is the canonical skill name. - info, err := os.Stat(filepath.Join(dir, name)) - if err != nil { - if errors.Is(err, fs.ErrNotExist) { - continue - } - catalog.diagnostics = append(catalog.diagnostics, fmt.Errorf("skill %q: %w", name, err)) - continue - } - if !info.IsDir() { - continue - } - if !skillName.MatchString(name) { - catalog.diagnostics = append(catalog.diagnostics, fmt.Errorf("invalid skill directory %q", name)) - continue - } - skill, err := parseSkill(filepath.Join(dir, name, skillFilename), name) - if errors.Is(err, fs.ErrNotExist) { - continue - } - if err != nil { - catalog.diagnostics = append(catalog.diagnostics, err) - continue - } - catalog.skills[skill.Name] = skill - } - return catalog, nil -} - -// LoadDefaultSkills discovers skills from the spec's scopes, merged with -// deterministic precedence. Roots are scanned lowest-precedence first so later -// roots override earlier ones on name collisions (recording a diagnostic): -// -// 1. ~/.agents/skills/ (user, cross-client) -// 2. user Ollama skills dir (user, Ollama-owned; SkillsDir) -// 3. /.agents/skills/ (project, cross-client) -// 4. /.ollama/skills/ (project, Ollama-owned) -// -// Project-level overrides user-level, and within a scope Ollama-owned -// directories override .agents/skills/. projectDir is the agent's working -// directory at startup (discovery is a session-start snapshot per the spec). -func LoadDefaultSkills(projectDir string) (*SkillCatalog, error) { - roots, err := defaultSkillRoots(projectDir) - if err != nil { - return nil, err - } - catalog := &SkillCatalog{skills: make(map[string]Skill)} - bundled, err := bundledSkillCreator() - if err != nil { - return nil, err - } - catalog.skills[bundled.Name] = bundled - if err := installBundledSkillCreator(); err != nil { - catalog.diagnostics = append(catalog.diagnostics, err) - } - for _, root := range roots { - sub, err := DiscoverSkills(root.path) - if err != nil { - catalog.diagnostics = append(catalog.diagnostics, fmt.Errorf("discover skills in %s: %w", root.path, err)) - continue - } - catalog.diagnostics = append(catalog.diagnostics, sub.diagnostics...) - for _, skill := range sub.skills { - // Name collisions across roots are expected precedence resolution, - // not errors: later (higher-precedence) roots legitimately override - // earlier ones. The skill is still loaded; no diagnostic needed. - catalog.skills[skill.Name] = skill - } - } - return catalog, nil -} - -func bundledSkillCreator() (Skill, error) { - skill, err := parseSkillContent("", bundledSkillCreatorName, bundledSkillCreatorContent) - if err != nil { - return Skill{}, fmt.Errorf("load bundled %s skill: %w", bundledSkillCreatorName, err) - } - return skill, nil -} - -func installBundledSkillCreator() error { - dir, err := SkillsDir() - if err != nil { - return fmt.Errorf("resolve bundled skill directory: %w", err) - } - path := filepath.Join(dir, bundledSkillCreatorName, skillFilename) - if err := os.MkdirAll(filepath.Dir(path), 0o755); err != nil { - return fmt.Errorf("create bundled skill directory: %w", err) - } - contents, err := os.ReadFile(path) - if err == nil && string(contents) == bundledSkillCreatorContent { - return nil - } - if err != nil && !errors.Is(err, fs.ErrNotExist) { - return fmt.Errorf("read bundled skill: %w", err) - } - if err := os.WriteFile(path, []byte(bundledSkillCreatorContent), 0o644); err != nil { - return fmt.Errorf("write bundled skill: %w", err) - } - return nil -} - -type skillRoot struct { - path string -} - -// SkillImportResult describes one import attempt. Failed skills do not prevent -// other valid skills in the same source root from being imported. -type SkillImportResult struct { - Source string - SourceDir string - Destination string - Imported []string - Existing []string - Failures []SkillImportFailure -} - -// SkillImportFailure identifies a source skill that was deliberately skipped. -// The destination is never changed for a failed skill. -type SkillImportFailure struct { - Name string - Err error -} - -// ImportSkills imports skills from a conventional coding-agent source into the -// canonical Ollama skills directory. Supported sources are codex, claude, and -// pi. Existing skills are left untouched: an identical directory is reported -// as existing, and a differing one is reported as a conflict. -func ImportSkills(source string) (SkillImportResult, error) { - home, err := os.UserHomeDir() - if err != nil { - return SkillImportResult{}, fmt.Errorf("resolve home directory: %w", err) - } - - destination, err := SkillsDir() - if err != nil { - return SkillImportResult{}, fmt.Errorf("resolve Ollama skills directory: %w", err) - } - return importSkillsFromRoots(source, conventionalSkillImportRoots(home), destination) -} - -func conventionalSkillImportRoots(home string) map[string]string { - return map[string]string{ - "codex": filepath.Join(home, ".codex", "skills"), - "claude": filepath.Join(home, ".claude", "skills"), - "pi": filepath.Join(home, ".pi", "agent", "skills"), - } -} - -func importSkillsFromRoots(source string, roots map[string]string, destination string) (SkillImportResult, error) { - source = strings.ToLower(strings.TrimSpace(source)) - sourceDir, ok := roots[source] - if !ok { - return SkillImportResult{}, fmt.Errorf("unknown skill source %q", source) - } - return importSkillsFromDir(source, sourceDir, destination) -} - -func importSkillsFromDir(source, sourceDir, destination string) (SkillImportResult, error) { - result := SkillImportResult{Source: source, SourceDir: sourceDir, Destination: destination} - info, err := os.Lstat(sourceDir) - if errors.Is(err, fs.ErrNotExist) { - return result, nil - } - if err != nil { - return result, fmt.Errorf("inspect %s skills directory: %w", source, err) - } - if info.Mode()&os.ModeSymlink != 0 { - return result, fmt.Errorf("inspect %s skills directory: symlinks are not supported", source) - } - if !info.IsDir() { - return result, fmt.Errorf("inspect %s skills directory: not a directory", source) - } - - entries, err := os.ReadDir(sourceDir) - if err != nil { - return result, fmt.Errorf("read %s skills directory: %w", source, err) - } - for _, entry := range entries { - name := entry.Name() - path := filepath.Join(sourceDir, name) - if entry.Type()&os.ModeSymlink != 0 { - result.Failures = append(result.Failures, SkillImportFailure{Name: name, Err: errors.New("symlinked skill directories are not supported")}) - continue - } - info, err := entry.Info() - if err != nil { - result.Failures = append(result.Failures, SkillImportFailure{Name: name, Err: fmt.Errorf("inspect source: %w", err)}) - continue - } - if !info.IsDir() { - continue - } - if !skillName.MatchString(name) { - result.Failures = append(result.Failures, SkillImportFailure{Name: name, Err: errors.New("invalid skill directory name")}) - continue - } - if err := validateImportSkill(path, name); err != nil { - result.Failures = append(result.Failures, SkillImportFailure{Name: name, Err: err}) - continue - } - - state, err := importSkillDirectory(path, filepath.Join(destination, name)) - if err != nil { - result.Failures = append(result.Failures, SkillImportFailure{Name: name, Err: err}) - continue - } - if state == skillImportExisting { - result.Existing = append(result.Existing, name) - } else { - result.Imported = append(result.Imported, name) - } - } - return result, nil -} - -func validateImportSkill(dir, name string) error { - manifest := filepath.Join(dir, skillFilename) - info, err := os.Lstat(manifest) - if err != nil { - return fmt.Errorf("inspect %s: %w", skillFilename, err) - } - if info.Mode()&os.ModeSymlink != 0 || !info.Mode().IsRegular() { - return fmt.Errorf("%s must be a regular, non-symlinked file", skillFilename) - } - if _, err := parseSkill(manifest, name); err != nil { - return err - } - return walkImportTree(dir, func(path string, entry fs.DirEntry, info fs.FileInfo) error { - if info.IsDir() || path == dir { - return nil - } - if !info.Mode().IsRegular() { - return fmt.Errorf("only regular files may be imported: %s", path) - } - file, err := os.Open(path) - if err != nil { - return fmt.Errorf("read %s: %w", path, err) - } - return file.Close() - }) -} - -func walkImportTree(root string, visit func(string, fs.DirEntry, fs.FileInfo) error) error { - return filepath.WalkDir(root, func(path string, entry fs.DirEntry, err error) error { - if err != nil { - return err - } - rel, err := filepath.Rel(root, path) - if err != nil || rel == ".." || strings.HasPrefix(rel, ".."+string(filepath.Separator)) { - return fmt.Errorf("unsafe skill path %q", path) - } - if entry.Type()&os.ModeSymlink != 0 { - return fmt.Errorf("symlinks may not be imported: %s", path) - } - info, err := entry.Info() - if err != nil { - return err - } - return visit(path, entry, info) - }) -} - -type skillImportState int - -const ( - skillImportCopied skillImportState = iota - skillImportExisting -) - -func importSkillDirectory(source, destination string) (skillImportState, error) { - if info, err := os.Lstat(destination); err == nil { - if info.Mode()&os.ModeSymlink != 0 || !info.IsDir() { - return 0, errors.New("destination exists but is not a regular directory") - } - same, err := sameImportTree(source, destination) - if err != nil { - return 0, fmt.Errorf("inspect existing destination: %w", err) - } - if same { - return skillImportExisting, nil - } - return 0, errors.New("destination skill already exists with different contents") - } else if !errors.Is(err, fs.ErrNotExist) { - return 0, fmt.Errorf("inspect destination: %w", err) - } - - if err := ensureImportDestination(filepath.Dir(destination)); err != nil { - return 0, err - } - stage, err := os.MkdirTemp(filepath.Dir(destination), "."+filepath.Base(destination)+".import-") - if err != nil { - return 0, fmt.Errorf("create import staging directory: %w", err) - } - defer os.RemoveAll(stage) - if err := copyImportTree(source, stage); err != nil { - return 0, err - } - if _, err := os.Lstat(destination); err == nil { - return 0, errors.New("destination skill was created during import") - } else if !errors.Is(err, fs.ErrNotExist) { - return 0, fmt.Errorf("inspect destination before install: %w", err) - } - if err := os.Rename(stage, destination); err != nil { - return 0, fmt.Errorf("install imported skill: %w", err) - } - return skillImportCopied, nil -} - -func ensureImportDestination(dir string) error { - if err := os.MkdirAll(dir, 0o755); err != nil { - return fmt.Errorf("create Ollama skills directory: %w", err) - } - info, err := os.Lstat(dir) - if err != nil { - return fmt.Errorf("inspect Ollama skills directory: %w", err) - } - if info.Mode()&os.ModeSymlink != 0 || !info.IsDir() { - return errors.New("Ollama skills directory must be a regular, non-symlinked directory") - } - return nil -} - -func copyImportTree(source, destination string) error { - return walkImportTree(source, func(path string, entry fs.DirEntry, info fs.FileInfo) error { - rel, err := filepath.Rel(source, path) - if err != nil { - return err - } - target := destination - if rel != "." { - target = filepath.Join(destination, rel) - } - if info.IsDir() { - if rel == "." { - return nil - } - return os.Mkdir(target, info.Mode().Perm()) - } - if !info.Mode().IsRegular() { - return fmt.Errorf("only regular files may be imported: %s", path) - } - return copyImportFile(path, target, info.Mode().Perm()) - }) -} - -func copyImportFile(source, destination string, mode fs.FileMode) error { - in, err := os.Open(source) - if err != nil { - return fmt.Errorf("read %s: %w", source, err) - } - defer in.Close() - out, err := os.OpenFile(destination, os.O_WRONLY|os.O_CREATE|os.O_EXCL, mode) - if err != nil { - return fmt.Errorf("create %s: %w", destination, err) - } - _, copyErr := io.Copy(out, in) - closeErr := out.Close() - if copyErr != nil { - return fmt.Errorf("copy %s: %w", source, copyErr) - } - if closeErr != nil { - return fmt.Errorf("write %s: %w", destination, closeErr) - } - return nil -} - -func sameImportTree(source, destination string) (bool, error) { - seen := make(map[string]struct{}) - same := true - err := walkImportTree(source, func(path string, entry fs.DirEntry, info fs.FileInfo) error { - rel, err := filepath.Rel(source, path) - if err != nil { - return err - } - seen[rel] = struct{}{} - other := destination - if rel != "." { - other = filepath.Join(destination, rel) - } - otherInfo, err := os.Lstat(other) - if errors.Is(err, fs.ErrNotExist) { - same = false - return nil - } - if err != nil { - return err - } - if otherInfo.Mode()&os.ModeSymlink != 0 || otherInfo.IsDir() != info.IsDir() || (!info.IsDir() && !otherInfo.Mode().IsRegular()) { - same = false - return nil - } - if info.Mode().IsRegular() { - equal, err := sameImportFile(path, other) - if err != nil { - return err - } - if !equal { - same = false - } - } - return nil - }) - if err != nil || !same { - return same, err - } - err = walkImportTree(destination, func(path string, entry fs.DirEntry, info fs.FileInfo) error { - rel, err := filepath.Rel(destination, path) - if err != nil { - return err - } - if _, ok := seen[rel]; !ok { - same = false - } - return nil - }) - return same, err -} - -func sameImportFile(first, second string) (bool, error) { - a, err := os.Open(first) - if err != nil { - return false, err - } - defer a.Close() - b, err := os.Open(second) - if err != nil { - return false, err - } - defer b.Close() - - left := make([]byte, 32*1024) - right := make([]byte, len(left)) - for { - n, errA := a.Read(left) - m, errB := b.Read(right) - if n != m || !bytes.Equal(left[:n], right[:m]) { - return false, nil - } - if errA == io.EOF && errB == io.EOF { - return true, nil - } - if errA != nil && errA != io.EOF { - return false, errA - } - if errB != nil && errB != io.EOF { - return false, errB - } - if errA == io.EOF || errB == io.EOF { - return false, nil - } - } -} - -// defaultSkillRoots returns skill directories ordered lowest- to -// highest-precedence. Non-existent directories are scanned harmlessly -// (DiscoverSkills skips them). -func defaultSkillRoots(projectDir string) ([]skillRoot, error) { - var roots []skillRoot - - if home, err := os.UserHomeDir(); err == nil && home != "" { - roots = append(roots, skillRoot{path: filepath.Join(home, ".agents", "skills")}) - } - - userOllama, err := SkillsDir() - if err != nil { - return nil, err - } - roots = append(roots, skillRoot{path: userOllama}) - - projectDir = strings.TrimSpace(projectDir) - if projectDir != "" { - if abs, err := filepath.Abs(projectDir); err == nil { - roots = append(roots, - skillRoot{path: filepath.Join(abs, ".agents", "skills")}, - skillRoot{path: filepath.Join(abs, ".ollama", "skills")}, - ) - } - } - return roots, nil -} - -func (c *SkillCatalog) Dir() string { - if c == nil { - return "" - } - return c.dir -} - -func (c *SkillCatalog) List() []Skill { - if c == nil { - return nil - } - list := make([]Skill, 0, len(c.skills)) - for _, skill := range c.skills { - list = append(list, skill) - } - sort.Slice(list, func(i, j int) bool { return list[i].Name < list[j].Name }) - return list -} - -func (c *SkillCatalog) Diagnostics() []error { - if c == nil { - return nil - } - return append([]error(nil), c.diagnostics...) -} - -// ExcludeNames removes skills whose names are reserved by a caller. It returns -// the excluded names in sorted order. -func (c *SkillCatalog) ExcludeNames(names []string) []string { - if c == nil { - return nil - } - reserved := make(map[string]struct{}, len(names)) - for _, name := range names { - name = strings.TrimPrefix(strings.ToLower(strings.TrimSpace(name)), "/") - if name != "" { - reserved[name] = struct{}{} - } - } - var excluded []string - for name := range c.skills { - if _, ok := reserved[name]; !ok { - continue - } - delete(c.skills, name) - excluded = append(excluded, name) - } - sort.Strings(excluded) - return excluded -} - -func (c *SkillCatalog) Load(name string) (Skill, error) { - name = strings.TrimSpace(name) - if !skillName.MatchString(name) { - return Skill{}, fmt.Errorf("invalid skill name %q", name) - } - if c == nil { - return Skill{}, errors.New("skills are unavailable") - } - skill, ok := c.skills[name] - if !ok { - return Skill{}, fmt.Errorf("skill %q not found in %s", name, c.dir) - } - return skill, nil -} - -// SystemContext advertises the catalog without expanding full instructions in -// every request. The skill call is the explicit loading boundary. -func (c *SkillCatalog) SystemContext() string { - list := c.List() - if len(list) == 0 { - return "" - } - lines := []string{""} - for _, skill := range list { - description := skill.Description - if description == "" { - description = "No description provided." - } - lines = append(lines, fmt.Sprintf("- %s: %s", skill.Name, description)) - } - lines = append(lines, "", "Load a matching skill with the skill tool before following its instructions. Skills only provide instructions; use ordinary tools for filesystem or network access, with their normal approval rules.") - return strings.Join(lines, "\n") -} - -func parseSkill(path, directoryName string) (Skill, error) { - // Stat (not Lstat) so a symlinked SKILL.md resolves to its target file. - info, err := os.Stat(path) - if err != nil { - return Skill{}, err - } - if !info.Mode().IsRegular() { - return Skill{}, fmt.Errorf("skill %q: %s is not a regular file", directoryName, skillFilename) - } - if info.Size() > maxSkillBytes { - return Skill{}, fmt.Errorf("skill %q: %s exceeds %d bytes", directoryName, skillFilename, maxSkillBytes) - } - data, err := os.ReadFile(path) - if err != nil { - return Skill{}, fmt.Errorf("read skill %q: %w", directoryName, err) - } - return parseSkillContent(path, directoryName, string(data)) -} - -func parseSkillContent(path, directoryName, input string) (Skill, error) { - instructions := strings.TrimSpace(input) - if instructions == "" { - return Skill{}, fmt.Errorf("skill %q: %s is empty", directoryName, skillFilename) - } - if !strings.HasPrefix(instructions, "---\n") && !strings.HasPrefix(instructions, "---\r\n") { - return Skill{}, fmt.Errorf("skill %q: missing YAML front matter", directoryName) - } - metadata, body, err := skillFrontMatter(instructions) - if err != nil { - return Skill{}, fmt.Errorf("skill %q: %w", directoryName, err) - } - if metadata.Name == "" { - return Skill{}, fmt.Errorf("skill %q: front matter requires name", directoryName) - } - if metadata.Description == "" { - return Skill{}, fmt.Errorf("skill %q: front matter requires description", directoryName) - } - if !skillName.MatchString(metadata.Name) { - return Skill{}, fmt.Errorf("skill %q: invalid front matter name %q", directoryName, metadata.Name) - } - if metadata.Name != directoryName { - return Skill{}, fmt.Errorf("skill %q: front matter name %q must match directory name", directoryName, metadata.Name) - } - skill := Skill{Name: metadata.Name, Description: metadata.Description, Path: path} - instructions = body - if strings.TrimSpace(instructions) == "" { - return Skill{}, fmt.Errorf("skill %q: instructions are empty", directoryName) - } - skill.Instructions = strings.TrimSpace(instructions) - return skill, nil -} - -type skillFrontMatterMetadata struct { - Name string `yaml:"name"` - Description string `yaml:"description"` - Metadata map[string]any `yaml:"metadata"` -} - -func skillFrontMatter(input string) (skillFrontMatterMetadata, string, error) { - input = strings.ReplaceAll(input, "\r\n", "\n") - lines := strings.Split(input, "\n") - if len(lines) < 3 || lines[0] != "---" { - return skillFrontMatterMetadata{}, "", errors.New("invalid front matter") - } - for i := 1; i < len(lines); i++ { - if lines[i] == "---" { - var metadata skillFrontMatterMetadata - if err := yaml.Unmarshal([]byte(strings.Join(lines[1:i], "\n")), &metadata); err != nil { - return skillFrontMatterMetadata{}, "", fmt.Errorf("parse YAML front matter: %w", err) - } - metadata.Name = strings.TrimSpace(metadata.Name) - metadata.Description = strings.TrimSpace(metadata.Description) - return metadata, strings.Join(lines[i+1:], "\n"), nil - } - } - return skillFrontMatterMetadata{}, "", errors.New("front matter is not closed") -} diff --git a/agent/skills_test.go b/agent/skills_test.go deleted file mode 100644 index 99e93feb6e9..00000000000 --- a/agent/skills_test.go +++ /dev/null @@ -1,516 +0,0 @@ -package agent - -import ( - "os" - "path/filepath" - "strings" - "testing" -) - -func writeCatalogSkill(t *testing.T, dir, name, content string) { - t.Helper() - path := filepath.Join(dir, name) - if err := os.MkdirAll(path, 0o755); err != nil { - t.Fatal(err) - } - if !strings.HasPrefix(content, "---") { - content = "---\nname: " + name + "\ndescription: Test skill.\n---\n" + content - } - if err := os.WriteFile(filepath.Join(path, skillFilename), []byte(content), 0o644); err != nil { - t.Fatal(err) - } -} - -func writeImportFixtureSkill(t *testing.T, dir string) { - t.Helper() - contents, err := os.ReadFile(filepath.Join("testdata", "import", "release-notes", skillFilename)) - if err != nil { - t.Fatal(err) - } - path := filepath.Join(dir, "release-notes", skillFilename) - if err := os.MkdirAll(filepath.Dir(path), 0o755); err != nil { - t.Fatal(err) - } - if err := os.WriteFile(path, contents, 0o644); err != nil { - t.Fatal(err) - } -} - -func TestDiscoverAndLoadSkills(t *testing.T) { - dir := t.TempDir() - writeCatalogSkill(t, dir, "release-notes", "---\nname: release-notes\ndescription: Draft concise release notes.\nmetadata:\n author: Ollama\n labels:\n - release\n - docs\n---\n# Release notes\n\nUse short bullets.") - catalog, err := DiscoverSkills(dir) - if err != nil { - t.Fatal(err) - } - list := catalog.List() - if len(list) != 1 || list[0].Name != "release-notes" || list[0].Description != "Draft concise release notes." { - t.Fatalf("skills = %#v", list) - } - skill, err := catalog.Load("release-notes") - if err != nil { - t.Fatal(err) - } - if !strings.Contains(skill.Content(), ``) || !strings.Contains(skill.Content(), "Use short bullets.") { - t.Fatalf("skill content = %q", skill.Content()) - } - if context := catalog.SystemContext(); !strings.Contains(context, "release-notes: Draft concise release notes.") || !strings.Contains(context, "normal approval rules") { - t.Fatalf("system context = %q", context) - } -} - -func TestDiscoverSkillsSkipsMalformedEntries(t *testing.T) { - dir := t.TempDir() - writeCatalogSkill(t, dir, "valid", "do the useful thing") - writeCatalogSkill(t, dir, "mismatched", "---\nname: whatever\ndescription: wrong name\n---\nbody") - // Genuinely malformed front matter (a line without a key:value pair) is still rejected. - writeCatalogSkill(t, dir, "broken", "---\nname: broken\ndescription\n---\nnope") - writeCatalogSkill(t, dir, "missing-name", "---\ndescription: missing name\n---\nbody") - writeCatalogSkill(t, dir, "missing-description", "---\nname: missing-description\n---\nbody") - writeCatalogSkill(t, dir, "bad-name", "---\nname: bad_name\ndescription: invalid name\n---\nbody") - writeCatalogSkill(t, dir, "under_score", "---\nname: under_score\ndescription: invalid directory\n---\nbody") - if err := os.MkdirAll(filepath.Join(dir, "no-front-matter"), 0o755); err != nil { - t.Fatal(err) - } - if err := os.WriteFile(filepath.Join(dir, "no-front-matter", skillFilename), []byte("body"), 0o644); err != nil { - t.Fatal(err) - } - catalog, err := DiscoverSkills(dir) - if err != nil { - t.Fatal(err) - } - if got, want := len(catalog.List()), 1; got != want { - t.Fatalf("valid skills = %d, want %d", got, want) - } - if got, want := len(catalog.Diagnostics()), 7; got != want { - t.Fatalf("diagnostics = %d, want %d: %#v", got, want, catalog.Diagnostics()) - } - if _, err := catalog.Load("broken"); err == nil || !strings.Contains(err.Error(), "not found") { - t.Fatalf("load broken error = %v", err) - } - if _, err := catalog.Load("../valid"); err == nil || !strings.Contains(err.Error(), "invalid skill name") { - t.Fatalf("unsafe name error = %v", err) - } -} - -func TestDiscoverSkillsFollowsSymlinks(t *testing.T) { - dir := t.TempDir() - target := t.TempDir() - writeCatalogSkill(t, target, "shared", "---\nname: shared\ndescription: From a linked repo.\n---\nshared instructions") - if err := os.Symlink(filepath.Join(target, "shared"), filepath.Join(dir, "shared")); err != nil { - t.Skipf("symlink not supported: %v", err) - } - catalog, err := DiscoverSkills(dir) - if err != nil { - t.Fatal(err) - } - list := catalog.List() - if len(list) != 1 || list[0].Name != "shared" || list[0].Description != "From a linked repo." { - t.Fatalf("symlinked skills = %#v", list) - } - if !strings.Contains(list[0].Content(), "shared instructions") { - t.Fatalf("symlinked skill content = %q", list[0].Content()) - } -} - -func TestLoadDefaultSkillsContinuesAfterBadRoot(t *testing.T) { - home := t.TempDir() - t.Setenv("HOME", home) - t.Setenv("USERPROFILE", home) - project := t.TempDir() - writeCatalogSkill(t, filepath.Join(project, ".ollama", "skills"), "release-notes", "project instructions") - - badRoot := filepath.Join(t.TempDir(), "not-a-directory") - if err := os.WriteFile(badRoot, []byte("not a directory"), 0o644); err != nil { - t.Fatal(err) - } - t.Setenv(SkillsDirEnv, badRoot) - - catalog, err := LoadDefaultSkills(project) - if err != nil { - t.Fatal(err) - } - if _, err := catalog.Load("release-notes"); err != nil { - t.Fatalf("valid skill was hidden by bad root: %v", err) - } - if _, err := catalog.Load(bundledSkillCreatorName); err != nil { - t.Fatalf("bundled skill was hidden by bad root: %v", err) - } - var foundDiagnostic bool - for _, diagnostic := range catalog.Diagnostics() { - if strings.Contains(diagnostic.Error(), badRoot) { - foundDiagnostic = true - break - } - } - if !foundDiagnostic { - t.Fatalf("diagnostics = %#v, want bad root %q", catalog.Diagnostics(), badRoot) - } -} - -func TestLoadDefaultSkillsInstallsBundledSkillCreator(t *testing.T) { - dir := t.TempDir() - t.Setenv(SkillsDirEnv, dir) - - catalog, err := LoadDefaultSkills("") - if err != nil { - t.Fatal(err) - } - skill, err := catalog.Load(bundledSkillCreatorName) - if err != nil { - t.Fatal(err) - } - path := filepath.Join(dir, bundledSkillCreatorName, skillFilename) - contents, err := os.ReadFile(path) - if err != nil { - t.Fatal(err) - } - if string(contents) != bundledSkillCreatorContent { - t.Fatalf("installed skill = %q, want bundled contents", contents) - } - if skill.Path != path { - t.Fatalf("skill path = %q, want %q", skill.Path, path) - } - if !strings.Contains(skill.Content(), "Skill directory: "+filepath.Dir(path)) { - t.Fatalf("skill content does not identify its directory: %q", skill.Content()) - } -} - -func TestLoadDefaultSkillsUpdatesExistingSkillCreator(t *testing.T) { - dir := t.TempDir() - t.Setenv(SkillsDirEnv, dir) - writeCatalogSkill(t, dir, bundledSkillCreatorName, "custom instructions") - - if _, err := LoadDefaultSkills(""); err != nil { - t.Fatal(err) - } - contents, err := os.ReadFile(filepath.Join(dir, bundledSkillCreatorName, skillFilename)) - if err != nil { - t.Fatal(err) - } - if string(contents) != bundledSkillCreatorContent { - t.Fatalf("installed skill = %q, want bundled contents", contents) - } -} - -func TestSkillsDirUsesOverrideAndXDG(t *testing.T) { - base := t.TempDir() - - override := filepath.Join(base, "skills-override") - t.Setenv(SkillsDirEnv, override) - got, err := SkillsDir() - if err != nil { - t.Fatal(err) - } - want, err := filepath.Abs(override) - if err != nil { - t.Fatal(err) - } - if got != want { - t.Fatalf("SkillsDir override = %q, want %q", got, want) - } - - t.Setenv(SkillsDirEnv, "") - xdg := filepath.Join(base, "xdg") - t.Setenv("XDG_CONFIG_HOME", xdg) - if got, err := SkillsDir(); err != nil || got != filepath.Join(xdg, "ollama", "skills") { - t.Fatalf("SkillsDir xdg = %q, want %q, %v", got, filepath.Join(xdg, "ollama", "skills"), err) - } - - t.Setenv("XDG_CONFIG_HOME", "") - home := filepath.Join(base, "home") - t.Setenv("HOME", home) - t.Setenv("USERPROFILE", home) - if got, err := SkillsDir(); err != nil || got != filepath.Join(home, ".ollama", "skills") { - t.Fatalf("SkillsDir default = %q, want %q, %v", got, filepath.Join(home, ".ollama", "skills"), err) - } -} - -func TestLoadDefaultSkillsPrecedenceAndCollisions(t *testing.T) { - home := t.TempDir() - t.Setenv("HOME", home) - t.Setenv("USERPROFILE", home) // Windows: os.UserHomeDir uses %USERPROFILE% - - userOllama := t.TempDir() - t.Setenv(SkillsDirEnv, userOllama) - - userAgents := filepath.Join(home, ".agents", "skills") - project := t.TempDir() - projectAgents := filepath.Join(project, ".agents", "skills") - projectOllama := filepath.Join(project, ".ollama", "skills") - - // release-notes exists in all four roots; project ollama must win. - writeCatalogSkill(t, userAgents, "release-notes", "from user agents") - writeCatalogSkill(t, userOllama, "release-notes", "from user ollama") - writeCatalogSkill(t, projectOllama, "release-notes", "from project ollama") - // code-review exists in both project roots; project ollama beats project agents. - writeCatalogSkill(t, projectAgents, "code-review", "from project agents") - writeCatalogSkill(t, projectOllama, "code-review", "from project ollama") - // unique appears only in user ollama (via env override). - writeCatalogSkill(t, userOllama, "unique", "only here") - - catalog, err := LoadDefaultSkills(project) - if err != nil { - t.Fatal(err) - } - rn, err := catalog.Load("release-notes") - if err != nil || !strings.Contains(rn.Instructions, "from project ollama") || !strings.Contains(rn.Path, ".ollama") { - t.Fatalf("release-notes = %#v, want project ollama to win", rn) - } - cr, err := catalog.Load("code-review") - if err != nil || !strings.Contains(cr.Instructions, "from project ollama") { - t.Fatalf("code-review = %#v, want project ollama to win over project agents", cr) - } - if _, err := catalog.Load("unique"); err != nil { - t.Fatalf("unique should load from user ollama: %v", err) - } - // Collisions are resolved silently by precedence — no diagnostics. - for _, d := range catalog.Diagnostics() { - if strings.Contains(d.Error(), "shadows") { - t.Fatalf("unexpected shadow diagnostic: %v", d) - } - } -} - -func TestSkillCatalogExcludeNames(t *testing.T) { - dir := t.TempDir() - for _, name := range []string{"release-notes", "system", "exit"} { - writeCatalogSkill(t, dir, name, "instructions") - } - catalog, err := DiscoverSkills(dir) - if err != nil { - t.Fatal(err) - } - - if got, want := strings.Join(catalog.ExcludeNames([]string{"/system", "EXIT"}), ","), "exit,system"; got != want { - t.Fatalf("excluded skills = %q, want %q", got, want) - } - if _, err := catalog.Load("system"); err == nil { - t.Fatal("excluded system skill should not load") - } - if _, err := catalog.Load("exit"); err == nil { - t.Fatal("excluded exit skill should not load") - } - if _, err := catalog.Load("release-notes"); err != nil { - t.Fatalf("non-conflicting skill should remain available: %v", err) - } -} - -func TestSkillContentListsDirectoryAndResources(t *testing.T) { - root := t.TempDir() - skillDir := filepath.Join(root, "pdf-processing") - if err := os.MkdirAll(filepath.Join(skillDir, "scripts"), 0o755); err != nil { - t.Fatal(err) - } - if err := os.MkdirAll(filepath.Join(skillDir, "references"), 0o755); err != nil { - t.Fatal(err) - } - if err := os.WriteFile(filepath.Join(skillDir, "SKILL.md"), []byte("---\nname: pdf-processing\ndescription: Handle PDFs.\n---\nHandle PDFs."), 0o644); err != nil { - t.Fatal(err) - } - if err := os.WriteFile(filepath.Join(skillDir, "scripts", "extract.py"), []byte("#!/usr/bin/env python3"), 0o755); err != nil { - t.Fatal(err) - } - if err := os.WriteFile(filepath.Join(skillDir, "references", "ref.md"), []byte("ref"), 0o644); err != nil { - t.Fatal(err) - } - catalog, err := DiscoverSkills(root) - if err != nil { - t.Fatal(err) - } - skill, err := catalog.Load("pdf-processing") - if err != nil { - t.Fatal(err) - } - content := skill.Content() - if !strings.Contains(content, "Skill directory:") || !strings.Contains(content, skillDir) { - t.Fatalf("content missing skill directory: %q", content) - } - if !strings.Contains(content, "scripts/extract.py") || !strings.Contains(content, "references/ref.md") { - t.Fatalf("content missing resource listing: %q", content) - } -} - -func TestImportSkillsCopiesFixtureAndIsIdempotent(t *testing.T) { - source := t.TempDir() - destination := t.TempDir() - writeImportFixtureSkill(t, source) - writeCatalogSkill(t, source, "broken", "---\nname: another-skill\ndescription: Deliberately invalid.\n---\nIgnore this.") - if err := os.MkdirAll(filepath.Join(source, "release-notes", "references"), 0o755); err != nil { - t.Fatal(err) - } - if err := os.WriteFile(filepath.Join(source, "release-notes", "references", "style.txt"), []byte("Keep it short.\n"), 0o644); err != nil { - t.Fatal(err) - } - if err := os.MkdirAll(filepath.Join(source, "release-notes", "scripts"), 0o755); err != nil { - t.Fatal(err) - } - if err := os.WriteFile(filepath.Join(source, "release-notes", "scripts", "prepare.sh"), []byte("#!/bin/sh\n"), 0o755); err != nil { - t.Fatal(err) - } - if err := os.WriteFile(filepath.Join(source, "ignored.md"), []byte("Ignored root file.\n"), 0o644); err != nil { - t.Fatal(err) - } - - result, err := importSkillsFromDir("codex", source, destination) - if err != nil { - t.Fatal(err) - } - if got, want := strings.Join(result.Imported, ","), "release-notes"; got != want { - t.Fatalf("imported = %q, want %q", got, want) - } - catalog, err := DiscoverSkills(destination) - if err != nil { - t.Fatal(err) - } - skill, err := catalog.Load("release-notes") - if err != nil || skill.Description != "Draft concise release notes." { - t.Fatalf("imported skill = %#v, %v", skill, err) - } - if got := len(result.Failures); got != 1 || result.Failures[0].Name != "broken" { - t.Fatalf("failures = %#v, want broken fixture failure", result.Failures) - } - for _, file := range []string{skillFilename, filepath.Join("references", "style.txt"), filepath.Join("scripts", "prepare.sh")} { - if _, err := os.Stat(filepath.Join(destination, "release-notes", file)); err != nil { - t.Fatalf("imported fixture file %q: %v", file, err) - } - } - - result, err = importSkillsFromDir("codex", source, destination) - if err != nil { - t.Fatal(err) - } - if got, want := strings.Join(result.Existing, ","), "release-notes"; got != want { - t.Fatalf("existing = %q, want %q", got, want) - } - if len(result.Imported) != 0 { - t.Fatalf("repeated import copied skills: %#v", result.Imported) - } -} - -func TestImportSkillsLeavesConflictsAndUnsafeSourcesUntouched(t *testing.T) { - source := t.TempDir() - destination := t.TempDir() - writeCatalogSkill(t, source, "release-notes", "source instructions") - writeCatalogSkill(t, destination, "release-notes", "existing instructions") - writeCatalogSkill(t, source, "nested-link", "safe manifest") - if err := os.Symlink(filepath.Join(source, "release-notes", skillFilename), filepath.Join(source, "nested-link", "reference")); err != nil { - t.Skipf("symlink not supported: %v", err) - } - if err := os.Symlink(filepath.Join(source, "release-notes"), filepath.Join(source, "linked-skill")); err != nil { - t.Skipf("symlink not supported: %v", err) - } - - result, err := importSkillsFromDir("codex", source, destination) - if err != nil { - t.Fatal(err) - } - if len(result.Imported) != 0 || len(result.Existing) != 0 { - t.Fatalf("unexpected successful import: %#v", result) - } - if got, err := os.ReadFile(filepath.Join(destination, "release-notes", skillFilename)); err != nil || !strings.Contains(string(got), "existing instructions") { - t.Fatalf("conflicting destination changed: %q, %v", got, err) - } - failed := make(map[string]bool) - for _, failure := range result.Failures { - failed[failure.Name] = true - } - for _, name := range []string{"release-notes", "nested-link", "linked-skill"} { - if !failed[name] { - t.Fatalf("missing failure for %q: %#v", name, result.Failures) - } - } -} - -func TestImportSkillsRejectsSymlinkedRoot(t *testing.T) { - root := t.TempDir() - source := filepath.Join(t.TempDir(), "codex-skills") - if err := os.Symlink(root, source); err != nil { - t.Skipf("symlink not supported: %v", err) - } - result, err := importSkillsFromDir("codex", source, t.TempDir()) - if err == nil || !strings.Contains(err.Error(), "symlinks are not supported") { - t.Fatalf("symlinked root error = %v", err) - } - if len(result.Imported) != 0 || len(result.Existing) != 0 || len(result.Failures) != 0 { - t.Fatalf("symlinked root result = %#v", result) - } -} - -func TestImportSkillsMissingRootAndConfiguredRoots(t *testing.T) { - result, err := importSkillsFromDir("codex", filepath.Join(t.TempDir(), "missing"), t.TempDir()) - if err != nil { - t.Fatal(err) - } - if len(result.Imported) != 0 || len(result.Existing) != 0 || len(result.Failures) != 0 { - t.Fatalf("missing root result = %#v", result) - } - - destination := t.TempDir() - rootBase := t.TempDir() - roots := map[string]string{ - "codex": filepath.Join(rootBase, "codex"), - "claude": filepath.Join(rootBase, "claude"), - "pi": filepath.Join(rootBase, "pi"), - } - for _, test := range []struct { - source string - root string - name string - }{ - {source: "codex", root: roots["codex"], name: "from-codex"}, - {source: "claude", root: roots["claude"], name: "from-claude"}, - {source: "pi", root: roots["pi"], name: "from-pi"}, - } { - t.Run(test.source, func(t *testing.T) { - writeCatalogSkill(t, test.root, test.name, "from "+test.source) - result, err = importSkillsFromRoots(test.source, roots, destination) - if err != nil { - t.Fatal(err) - } - if result.SourceDir != test.root { - t.Fatalf("source dir = %q, want %q", result.SourceDir, test.root) - } - if _, err := os.Stat(filepath.Join(destination, test.name, skillFilename)); err != nil { - t.Fatalf("conventional source was not imported: %v", err) - } - }) - } - if _, err := importSkillsFromRoots("unknown", roots, destination); err == nil || !strings.Contains(err.Error(), "unknown skill source") { - t.Fatalf("unknown source error = %v", err) - } -} - -func TestConventionalSkillImportRoots(t *testing.T) { - home := t.TempDir() - roots := conventionalSkillImportRoots(home) - for source, want := range map[string]string{ - "codex": filepath.Join(home, ".codex", "skills"), - "claude": filepath.Join(home, ".claude", "skills"), - "pi": filepath.Join(home, ".pi", "agent", "skills"), - } { - if got := roots[source]; got != want { - t.Fatalf("%s root = %q, want %q", source, got, want) - } - } -} - -func TestImportSkillsRejectsUnreadableManifest(t *testing.T) { - source := t.TempDir() - writeCatalogSkill(t, source, "private", "do not read") - manifest := filepath.Join(source, "private", skillFilename) - if err := os.Chmod(manifest, 0); err != nil { - t.Fatal(err) - } - t.Cleanup(func() { _ = os.Chmod(manifest, 0o644) }) - if _, err := os.ReadFile(manifest); err == nil { - t.Skip("test user can read a mode-000 file") - } - result, err := importSkillsFromDir("codex", source, t.TempDir()) - if err != nil { - t.Fatal(err) - } - if len(result.Failures) != 1 || result.Failures[0].Name != "private" { - t.Fatalf("failures = %#v", result.Failures) - } -} diff --git a/agent/testdata/import/release-notes/SKILL.md b/agent/testdata/import/release-notes/SKILL.md deleted file mode 100644 index dae33c01526..00000000000 --- a/agent/testdata/import/release-notes/SKILL.md +++ /dev/null @@ -1,8 +0,0 @@ ---- -name: release-notes -description: Draft concise release notes. ---- - -# Release notes - -Use short bullets. diff --git a/agent/tools/bash.go b/agent/tools/bash.go deleted file mode 100644 index 733629b515f..00000000000 --- a/agent/tools/bash.go +++ /dev/null @@ -1,450 +0,0 @@ -package tools - -import ( - "context" - "errors" - "fmt" - "os" - "os/exec" - "path/filepath" - "runtime" - "strings" - "time" - "unicode/utf8" - - "github.com/ollama/ollama/agent" - "github.com/ollama/ollama/api" -) - -const ( - bashTimeout = 3 * time.Minute - bashWaitDelay = 1 * time.Second - maxBashOutputBytes = 60_000 -) - -type Bash struct{} - -func (b *Bash) Name() string { - return shellToolName() -} - -func (b *Bash) Description() string { - return shellToolDescription() -} - -func (b *Bash) Schema() api.ToolFunction { - props := api.NewToolPropertiesMap() - props.Set("command", api.ToolProperty{ - Type: api.PropertyType{"string"}, - Description: shellCommandDescription(), - }) - return api.ToolFunction{ - Name: b.Name(), - Description: b.Description(), - Parameters: api.ToolFunctionParameters{ - Type: "object", - Properties: props, - Required: []string{"command"}, - }, - } -} - -func (b *Bash) RequiresApproval(map[string]any) bool { - return true -} - -// ApprovalScope scopes shell approval to the exact, trimmed command string -// using a NUL separator: "\x00". "Always allow this command" -// matches ONLY that precise string — any whitespace, quoting, or casing -// variant re-prompts. The NUL separator is safe because a shell command -// string cannot contain a literal NUL. -func (b *Bash) ApprovalScope(args map[string]any) string { - name := b.Name() - if command, ok := args["command"].(string); ok { - command = strings.TrimSpace(command) - if command != "" { - return name + "\x00" + command - } - } - return name -} - -func (b *Bash) Execute(ctx context.Context, toolCtx agent.ToolContext, args map[string]any) (agent.ToolResult, error) { - // TODO: use shared agent.RequiredStringArg for the "command" parameter (see agent package cleanup plan). - command, ok := args["command"].(string) - if !ok || strings.TrimSpace(command) == "" { - return agent.ToolResult{}, fmt.Errorf("command parameter is required") - } - if err := rejectUnsafeShellCommand(command); err != nil { - return agent.ToolResult{}, err - } - - ctx, cancel := context.WithTimeout(ctx, bashTimeout) - defer cancel() - - cwdFile, err := os.CreateTemp("", "ollama-agent-cwd-*") - if err != nil { - return agent.ToolResult{}, err - } - cwdPath := cwdFile.Name() - _ = cwdFile.Close() - defer os.Remove(cwdPath) - - cmd := newBashCommand(ctx, command, cwdPath) - cmd.WaitDelay = bashWaitDelay - cmd.Cancel = func() error { - return killBashCommand(cmd) - } - if toolCtx.WorkingDir != "" { - cmd.Dir = toolCtx.WorkingDir - } - - var stdout, stderr boundedOutput - stdout.Limit = maxBashOutputBytes - stderr.Limit = maxBashOutputBytes - cmd.Stdout = &stdout - cmd.Stderr = &stderr - - err = runBashCommand(cmd) - finalWorkingDir := readFinalWorkingDir(cwdPath) - - var sb strings.Builder - if stdout.Len() > 0 { - sb.WriteString(stdout.String("stdout")) - } - if stderr.Len() > 0 { - if sb.Len() > 0 { - sb.WriteString("\n") - } - sb.WriteString("stderr:\n") - sb.WriteString(stderr.String("stderr")) - } - - if err != nil { - if ctx.Err() == context.DeadlineExceeded { - return agent.ToolResult{Content: bashContentWithError(sb.String(), "Error: command timed out after "+bashTimeout.String()), WorkingDir: finalWorkingDir}, nil - } - if ctx.Err() == context.Canceled { - return agent.ToolResult{Content: bashContentWithError(sb.String(), "Error: command was canceled"), WorkingDir: finalWorkingDir}, nil - } - if errors.Is(err, exec.ErrWaitDelay) { - _ = killBashCommand(cmd) - return agent.ToolResult{Content: bashContentWithError(sb.String(), "Error: command output pipes did not close after "+bashWaitDelay.String()), WorkingDir: finalWorkingDir}, nil - } - if exitErr, ok := err.(*exec.ExitError); ok { - return agent.ToolResult{Content: bashContentWithError(sb.String(), fmt.Sprintf("Exit code: %d", exitErr.ExitCode())), WorkingDir: finalWorkingDir}, nil - } - return agent.ToolResult{Content: sb.String(), WorkingDir: finalWorkingDir}, fmt.Errorf("executing command: %w", err) - } - - if sb.Len() == 0 { - return agent.ToolResult{Content: "(no output)", WorkingDir: finalWorkingDir}, nil - } - return agent.ToolResult{Content: sb.String(), WorkingDir: finalWorkingDir}, nil -} - -func bashContentWithError(content, msg string) string { - if content == "" { - return msg - } - return content + "\n\n" + msg -} - -// rejectUnsafeShellCommand applies a best-effort blocklist for obviously -// destructive or credential-exfiltrating commands. It is defense-in-depth -// ONLY: the interactive approval prompt is the real security control, and -// this check must not be relied upon as a sandbox. Sophisticated or novel -// dangerous commands (e.g. find / -delete, dd, fork bombs, custom binaries) -// are NOT caught here and will simply be routed through approval like any -// other command. Keep the approval prompt as the gate. -func rejectUnsafeShellCommand(command string) error { - switch { - case hasUnsafeRecursiveDelete(command): - return fmt.Errorf("refusing to run unsafe command: recursive delete target is too broad") - case readsCredentialPath(command): - return fmt.Errorf("refusing to run unsafe command: credential file reads are not allowed") - default: - return nil - } -} - -func hasUnsafeRecursiveDelete(command string) bool { - // Check each command segment independently. shellSafetyText flattens - // separators (; & | newlines) to spaces, which would otherwise let the - // rm target scan bleed across command boundaries — e.g. - // "rm -rf build && echo ~/.ssh/config" flattened to one token stream - // would treat the unrelated ~/.ssh/config (a ~/-prefixed "unsafe - // target") as an rm argument. Splitting on separators first restores - // command boundaries while still catching multi-target single commands - // like "rm -rf build /etc". - for _, segment := range shellSegments(command) { - fields := shellSafetyFields(segment) - for i, field := range fields { - if isRMCommand(field) && rmCommandDeletesUnsafeTarget(fields[i+1:]) { - return true - } - if isPowerShellDeleteCommand(field) && powerShellDeleteCommandDeletesUnsafeTarget(fields[i+1:]) { - return true - } - } - } - return false -} - -// shellSegments splits a command on shell control operators (;, &, |, &&, -// ||) and newlines, returning the individual command segments. It operates on -// the lowercased raw command before quote/separator normalization so that -// command boundaries are preserved for per-segment checks. Subshell parens are -// intentionally NOT treated as separators: splitting on them would fragment -// command substitutions like "rm -rf $(echo /)" into "rm -rf $" and "echo /", -// hiding the destructive "/" target from the per-segment scan. Empty segments -// are dropped. -func shellSegments(command string) []string { - command = strings.ToLower(command) - var segments []string - for _, segment := range strings.FieldsFunc(command, func(r rune) bool { - switch r { - case ';', '&', '|', '\n', '\r': - return true - } - return false - }) { - if segment = strings.TrimSpace(segment); segment != "" { - segments = append(segments, segment) - } - } - return segments -} - -func rmCommandDeletesUnsafeTarget(fields []string) bool { - var flags string - for _, field := range fields { - if field == "--" { - continue - } - if strings.HasPrefix(field, "-") { - flags += field - continue - } - if strings.Contains(flags, "r") && strings.Contains(flags, "f") && isUnsafeDeleteTarget(field) { - return true - } - } - return false -} - -func powerShellDeleteCommandDeletesUnsafeTarget(fields []string) bool { - var recurse, force bool - var targets []string - for _, field := range fields { - switch field { - case "-r", "-recurse", "-recursive": - recurse = true - case "-f", "-force": - force = true - default: - if !strings.HasPrefix(field, "-") { - targets = append(targets, field) - } - } - } - if !recurse || !force { - return false - } - for _, target := range targets { - if isUnsafeDeleteTarget(target) { - return true - } - } - return false -} - -func readsCredentialPath(command string) bool { - fields := shellSafetyFields(command) - if !hasCredentialReadVerb(fields) { - return false - } - normalized := shellSafetyText(command) - for _, fragment := range []string{ - "/.ssh/id_rsa", - "/.ssh/id_dsa", - "/.ssh/id_ecdsa", - "/.ssh/id_ed25519", - "/.ssh/config", - "/.ssh/known_hosts", - "/.aws/credentials", - "/.aws/config", - "/.config/gcloud/application_default_credentials.json", - "/.kube/config", - "/.netrc", - "/.npmrc", - "/.docker/config.json", - "/.config/gh/hosts.yml", - "/.gnupg/", - "/etc/shadow", - } { - if strings.Contains(normalized, fragment) { - return true - } - } - return false -} - -func hasCredentialReadVerb(fields []string) bool { - for _, field := range fields { - switch field { - case "cat", "less", "more", "head", "tail", "type", "get-content", "gc", "select-string", "grep", "rg", "sed", "awk": - return true - case "env", "printenv": - return true - } - } - return false -} - -func isRMCommand(field string) bool { - return field == "rm" || strings.HasSuffix(field, "/rm") -} - -func isPowerShellDeleteCommand(field string) bool { - switch field { - case "remove-item", "del", "erase", "rd", "rmdir": - return true - default: - return false - } -} - -func isUnsafeDeleteTarget(target string) bool { - if target == "." || target == "./" || target == "*" { - return true - } - if target == "/*" { - return true - } - target = strings.TrimSuffix(target, "/*") - for _, prefix := range []string{"~/", "$home/", "${home}/", "$env:home/", "$env:userprofile/", "%userprofile%/"} { - if strings.HasPrefix(target, prefix) { - return true - } - } - for _, prefix := range []string{"/etc/", "/bin/", "/sbin/", "/usr/", "/var/", "/lib/", "/library/", "/system/", "/applications/", "c:/windows/", "c:/program files/"} { - if strings.HasPrefix(target, prefix) { - return true - } - } - for _, exact := range []string{"/", "~", "$home", "${home}", "$env:home", "$env:userprofile", "%userprofile%", "c:", "c:/", "/etc", "/bin", "/sbin", "/usr", "/var", "/lib", "/library", "/system", "/applications", "c:/windows", "c:/program files"} { - if target == exact { - return true - } - } - return false -} - -func shellSafetyFields(command string) []string { - return strings.Fields(shellSafetyText(command)) -} - -func shellSafetyText(command string) string { - command = strings.ToLower(command) - return strings.NewReplacer( - "\\", "/", - "\n", " ", - "\t", " ", - ";", " ", - "&", " ", - "|", " ", - "(", " ", - ")", " ", - "\"", "", - "'", "", - "`", "", - ).Replace(command) -} - -func readFinalWorkingDir(path string) string { - content, err := os.ReadFile(path) - if err != nil { - return "" - } - workingDir := strings.TrimPrefix(string(content), "\ufeff") - workingDir = strings.TrimSpace(workingDir) - if workingDir == "" { - return "" - } - workingDir = normalizeBashWorkingDir(workingDir) - info, err := os.Stat(workingDir) - if err != nil || !info.IsDir() { - return "" - } - return workingDir -} - -func normalizeBashWorkingDir(workingDir string) string { - if runtime.GOOS == "windows" && len(workingDir) >= 3 && workingDir[0] == '/' && workingDir[2] == '/' && isASCIIAlpha(workingDir[1]) { - workingDir = strings.ToUpper(string(workingDir[1])) + ":" + workingDir[2:] - } - workingDir = filepath.Clean(filepath.FromSlash(workingDir)) - if runtime.GOOS == "windows" && len(workingDir) >= 2 && workingDir[1] == ':' && isASCIIAlpha(workingDir[0]) { - workingDir = strings.ToUpper(string(workingDir[0])) + workingDir[1:] - } - return workingDir -} - -func isASCIIAlpha(b byte) bool { - return (b >= 'a' && b <= 'z') || (b >= 'A' && b <= 'Z') -} - -type boundedOutput struct { - Limit int - buf []byte - omitted int -} - -func (b *boundedOutput) Write(p []byte) (int, error) { - if b.Limit <= 0 { - b.omitted += len(p) - return len(p), nil - } - remaining := b.Limit - len(b.buf) - if remaining <= 0 { - b.omitted += len(p) - return len(p), nil - } - if len(p) <= remaining { - b.buf = append(b.buf, p...) - return len(p), nil - } - writeLen := utf8SafePrefixLen(p[:remaining]) - b.buf = append(b.buf, p[:writeLen]...) - b.omitted += len(p) - writeLen - return len(p), nil -} - -func (b *boundedOutput) Len() int { - return len(b.buf) + b.omitted -} - -func (b *boundedOutput) String(label string) string { - safeLen := utf8SafePrefixLen(b.buf) - content := string(b.buf[:safeLen]) - omitted := b.omitted + len(b.buf) - safeLen - if omitted == 0 { - return content - } - return content + agent.TruncMarker(label, safeLen, 0, omitted, false, "") -} - -func utf8SafePrefixLen(p []byte) int { - if len(p) == 0 { - return 0 - } - for i := 0; i < len(p); { - r, size := utf8.DecodeRune(p[i:]) - if r == utf8.RuneError && size == 1 { - return i - } - i += size - } - return len(p) -} diff --git a/agent/tools/bash_test.go b/agent/tools/bash_test.go deleted file mode 100644 index 537ac98a961..00000000000 --- a/agent/tools/bash_test.go +++ /dev/null @@ -1,258 +0,0 @@ -package tools - -import ( - "context" - "os" - "path/filepath" - "runtime" - "strings" - "testing" - "unicode/utf8" - - "github.com/ollama/ollama/agent" -) - -func TestBashReportsFinalWorkingDir(t *testing.T) { - root := t.TempDir() - subdir := filepath.Join(root, "sub") - if err := os.Mkdir(subdir, 0o755); err != nil { - t.Fatal(err) - } - - result, err := (&Bash{}).Execute(context.Background(), agent.ToolContext{WorkingDir: root}, map[string]any{ - "command": shellTestCommand("cd sub && pwd", "Set-Location sub; Get-Location"), - }) - if err != nil { - t.Fatal(err) - } - wantDir, err := filepath.EvalSymlinks(subdir) - if err != nil { - t.Fatal(err) - } - if result.WorkingDir != wantDir { - t.Fatalf("working dir = %q, want %q", result.WorkingDir, wantDir) - } - if !strings.Contains(result.Content, "sub") { - t.Fatalf("content = %q, want pwd output", result.Content) - } -} - -func TestBashBoundsOutputWhileRunning(t *testing.T) { - result, err := (&Bash{}).Execute(context.Background(), agent.ToolContext{WorkingDir: t.TempDir()}, map[string]any{ - "command": shellTestCommand("yes x | head -c 70000", "[Console]::Out.Write(('x' * 70000))"), - }) - if err != nil { - t.Fatal(err) - } - if !strings.Contains(result.Content, "[stdout truncated: showing first ~") || !strings.Contains(result.Content, "omitted ~") || !strings.Contains(result.Content, " tokens.]") { - t.Fatalf("content = %q, want stdout truncation marker", result.Content) - } - if count, want := strings.Count(result.Content, "x"), shellTestCapturedXCount(); count != want { - t.Fatalf("captured x count = %d, want %d", count, want) - } - if len(result.Content) > maxBashOutputBytes+200 { - t.Fatalf("content length = %d, want bounded output", len(result.Content)) - } -} - -func TestBoundedOutputTruncatesAtUTF8Boundary(t *testing.T) { - var out boundedOutput - out.Limit = len([]byte("abc")) + 1 - - if _, err := out.Write([]byte("abcédef")); err != nil { - t.Fatal(err) - } - content := out.String("stdout") - if !utf8.ValidString(content) { - t.Fatalf("content is not valid UTF-8: %q", content) - } - if strings.ContainsRune(content, utf8.RuneError) { - t.Fatalf("content contains replacement rune: %q", content) - } - if !strings.HasPrefix(content, "abc\n\n[stdout truncated:") { - t.Fatalf("content = %q, want complete ASCII prefix and truncation marker", content) - } -} - -func TestBoundedOutputKeepsCompleteUTF8AtBoundary(t *testing.T) { - var out boundedOutput - out.Limit = len([]byte("abcé")) - - if _, err := out.Write([]byte("abcédef")); err != nil { - t.Fatal(err) - } - if content := out.String("stdout"); !strings.HasPrefix(content, "abcé\n\n[stdout truncated:") { - t.Fatalf("content = %q, want complete UTF-8 prefix", content) - } -} - -func TestBoundedOutputTrimsTrailingPartialUTF8(t *testing.T) { - var out boundedOutput - out.Limit = 4 - - if _, err := out.Write([]byte{'a', 'b', 'c', 0xc3}); err != nil { - t.Fatal(err) - } - if _, err := out.Write([]byte{0xa9}); err != nil { - t.Fatal(err) - } - if content := out.String("stdout"); !utf8.ValidString(content) || !strings.HasPrefix(content, "abc\n\n[stdout truncated:") { - t.Fatalf("content = %q, want valid UTF-8 with partial suffix trimmed", content) - } -} - -func TestUTF8SafePrefixRejectsMalformedLeadByte(t *testing.T) { - input := []byte{'a', 0xc0, 0x80, 'b'} - if got := utf8SafePrefixLen(input); got != 1 { - t.Fatalf("safe prefix length = %d, want 1", got) - } -} - -func TestBoundedOutputDropsMalformedUTF8(t *testing.T) { - var out boundedOutput - out.Limit = 4 - - if _, err := out.Write([]byte{'a', 0xc0, 0x80, 'b'}); err != nil { - t.Fatal(err) - } - content := out.String("stdout") - if !utf8.ValidString(content) { - t.Fatalf("content is not valid UTF-8: %q", content) - } - if strings.ContainsRune(content, utf8.RuneError) { - t.Fatalf("content contains replacement rune: %q", content) - } - if !strings.HasPrefix(content, "a\n\n[stdout truncated:") { - t.Fatalf("content = %q, want valid prefix and truncation marker", content) - } -} - -func TestBashReportsCanceledCommand(t *testing.T) { - ctx, cancel := context.WithCancel(context.Background()) - cancel() - - result, err := (&Bash{}).Execute(ctx, agent.ToolContext{WorkingDir: t.TempDir()}, map[string]any{ - "command": shellTestCommand("sleep 10", "Start-Sleep -Seconds 10"), - }) - if err != nil { - t.Fatal(err) - } - if !strings.Contains(result.Content, "Error: command was canceled") { - t.Fatalf("content = %q, want canceled message", result.Content) - } - if strings.Contains(result.Content, "Exit code: -1") { - t.Fatalf("content = %q, should not mask cancellation as exit code", result.Content) - } -} - -func TestRejectUnsafeShellCommand(t *testing.T) { - tests := []struct { - name string - command string - wantErr bool - }{ - {name: "rm root", command: "rm -rf /", wantErr: true}, - {name: "sudo rm root", command: "sudo rm -rf -- /", wantErr: true}, - {name: "rm home", command: "rm -fr $HOME", wantErr: true}, - {name: "rm root wildcard", command: "rm -rf /*", wantErr: true}, - {name: "rm system subdir", command: "rm -rf /etc/ssh", wantErr: true}, - {name: "rm cwd", command: "rm -rf .", wantErr: true}, - {name: "powershell remove root", command: `Remove-Item -Recurse -Force C:\`, wantErr: true}, - {name: "powershell remove system subdir", command: `Remove-Item -Recurse -Force C:\Windows\Temp`, wantErr: true}, - {name: "ssh private key", command: "cat ~/.ssh/id_rsa", wantErr: true}, - {name: "aws credentials", command: "Get-Content $HOME/.aws/credentials", wantErr: true}, - {name: "shadow", command: "head /etc/shadow", wantErr: true}, - {name: "netrc", command: "cat ~/.netrc", wantErr: true}, - {name: "docker config", command: "cat ~/.docker/config.json", wantErr: true}, - {name: "gnupg dir", command: "cat ~/.gnupg/private-keys-v1.d/key", wantErr: true}, - {name: "gh hosts", command: "cat ~/.config/gh/hosts.yml", wantErr: true}, - {name: "ssh config", command: "cat ~/.ssh/config", wantErr: true}, - {name: "printenv dump", command: "printenv", wantErr: false}, - {name: "delete build dir", command: "rm -rf build", wantErr: false}, - {name: "read project file", command: "cat README.md", wantErr: false}, - {name: "mention key text", command: "rg id_rsa docs", wantErr: false}, - {name: "env example", command: "cat .env.example", wantErr: false}, - {name: "rm build then unrelated tilde path", command: "rm -rf build && echo ~/.ssh/config", wantErr: false}, - {name: "rm build then unrelated slash path", command: "rm -rf build; cat /etc/passwd", wantErr: false}, - {name: "rm build then unrelated star glob", command: "rm -rf build && ls *.go", wantErr: false}, - {name: "rm multiple targets one unsafe", command: "rm -rf build /etc", wantErr: true}, - {name: "rm unsafe then safe piped", command: "rm -rf / | tee log", wantErr: true}, - {name: "rm unsafe via command substitution", command: "rm -rf $(echo /)", wantErr: true}, - } - - for _, tt := range tests { - t.Run(tt.name, func(t *testing.T) { - err := rejectUnsafeShellCommand(tt.command) - if tt.wantErr && err == nil { - t.Fatal("expected unsafe command to be rejected") - } - if !tt.wantErr && err != nil { - t.Fatalf("command rejected: %v", err) - } - }) - } -} - -func TestBashRejectsUnsafeCommandBeforeExecution(t *testing.T) { - _, err := (&Bash{}).Execute(context.Background(), agent.ToolContext{WorkingDir: t.TempDir()}, map[string]any{ - "command": "rm -rf /", - }) - if err == nil || !strings.Contains(err.Error(), "refusing to run unsafe command") { - t.Fatalf("err = %v, want unsafe command rejection", err) - } -} - -func shellTestCommand(unix, windows string) string { - if runtime.GOOS == "windows" { - return windows - } - return unix -} - -func shellTestCapturedXCount() int { - if runtime.GOOS == "windows" { - return maxBashOutputBytes - } - return maxBashOutputBytes / 2 -} - -func TestReadFinalWorkingDirRejectsInvalidPaths(t *testing.T) { - dir := t.TempDir() - cwdFile := filepath.Join(dir, "cwd") - notDir := filepath.Join(dir, "file.txt") - if err := os.WriteFile(notDir, []byte("not a dir"), 0o644); err != nil { - t.Fatal(err) - } - - if err := os.WriteFile(cwdFile, []byte(notDir+"\n"), 0o644); err != nil { - t.Fatal(err) - } - if got := readFinalWorkingDir(cwdFile); got != "" { - t.Fatalf("regular file cwd = %q, want empty", got) - } - - if err := os.WriteFile(cwdFile, []byte(filepath.Join(dir, "missing")+"\n"), 0o644); err != nil { - t.Fatal(err) - } - if got := readFinalWorkingDir(cwdFile); got != "" { - t.Fatalf("missing cwd = %q, want empty", got) - } - - if err := os.WriteFile(cwdFile, []byte(dir+"\n"), 0o644); err != nil { - t.Fatal(err) - } - if got := readFinalWorkingDir(cwdFile); got != dir { - t.Fatalf("directory cwd = %q, want %q", got, dir) - } -} - -func TestNormalizeBashWorkingDirWindowsDriveLetter(t *testing.T) { - if runtime.GOOS != "windows" { - t.Skip("windows path normalization") - } - got := normalizeBashWorkingDir("/c/Users/jdoe/project") - want := filepath.Clean(`C:\Users\jdoe\project`) - if got != want { - t.Fatalf("working dir = %q, want %q", got, want) - } -} diff --git a/agent/tools/bash_unix.go b/agent/tools/bash_unix.go deleted file mode 100644 index 10683e8d00e..00000000000 --- a/agent/tools/bash_unix.go +++ /dev/null @@ -1,49 +0,0 @@ -//go:build !windows - -package tools - -import ( - "context" - "os/exec" - "strings" - "syscall" -) - -func shellToolName() string { - return "bash" -} - -func shellToolDescription() string { - return "Execute a bash command on the system. Use this to inspect files, run tests, and perform development tasks." -} - -func shellCommandDescription() string { - return "The bash command to execute." -} - -func newBashCommand(ctx context.Context, command, cwdPath string) *exec.Cmd { - script := command + "\n__ollama_status=$?\npwd -P > " + shellQuote(cwdPath) + "\nexit $__ollama_status" - cmd := exec.CommandContext(ctx, "bash", "-c", script) - configureBashCommand(cmd) - return cmd -} - -func shellQuote(value string) string { - return "'" + strings.ReplaceAll(value, "'", "'\\''") + "'" -} - -func configureBashCommand(cmd *exec.Cmd) { - cmd.SysProcAttr = &syscall.SysProcAttr{Setpgid: true} -} - -func runBashCommand(cmd *exec.Cmd) error { - return cmd.Run() -} - -func killBashCommand(cmd *exec.Cmd) error { - if cmd == nil || cmd.Process == nil { - return nil - } - _ = syscall.Kill(-cmd.Process.Pid, syscall.SIGKILL) - return nil -} diff --git a/agent/tools/bash_unix_test.go b/agent/tools/bash_unix_test.go deleted file mode 100644 index 5c7992adb85..00000000000 --- a/agent/tools/bash_unix_test.go +++ /dev/null @@ -1,40 +0,0 @@ -//go:build !windows - -package tools - -import ( - "context" - "os/exec" - "strings" - "testing" - "time" - - "github.com/ollama/ollama/agent" -) - -func TestConfigureBashCommandSetsProcessGroup(t *testing.T) { - cmd := exec.Command("bash", "-c", "true") - configureBashCommand(cmd) - if cmd.SysProcAttr == nil || !cmd.SysProcAttr.Setpgid { - t.Fatalf("configureBashCommand should start bash in a new process group") - } -} - -func TestBashWaitDelayBoundsBackgroundOutputPipe(t *testing.T) { - start := time.Now() - result, err := (&Bash{}).Execute(context.Background(), agent.ToolContext{WorkingDir: t.TempDir()}, map[string]any{ - "command": "sleep 5 & echo done", - }) - if err != nil { - t.Fatal(err) - } - if elapsed := time.Since(start); elapsed > bashWaitDelay+2*time.Second { - t.Fatalf("command elapsed = %s, want bounded near %s", elapsed, bashWaitDelay) - } - if !strings.Contains(result.Content, "done") { - t.Fatalf("content = %q, want command output", result.Content) - } - if !strings.Contains(result.Content, "output pipes did not close") { - t.Fatalf("content = %q, want wait delay message", result.Content) - } -} diff --git a/agent/tools/bash_windows.go b/agent/tools/bash_windows.go deleted file mode 100644 index 840ff24fca2..00000000000 --- a/agent/tools/bash_windows.go +++ /dev/null @@ -1,134 +0,0 @@ -//go:build windows - -package tools - -import ( - "context" - "os/exec" - "strings" - "sync" - "unsafe" - - "golang.org/x/sys/windows" -) - -var bashJobHandles sync.Map - -func shellToolName() string { - return "powershell" -} - -func shellToolDescription() string { - return "Execute a PowerShell command on the system. Use this to inspect files, run tests, and perform development tasks." -} - -func shellCommandDescription() string { - return "The PowerShell command to execute." -} - -func newBashCommand(ctx context.Context, command, cwdPath string) *exec.Cmd { - return exec.CommandContext( - ctx, - "powershell.exe", - "-NoLogo", - "-NoProfile", - "-NonInteractive", - "-ExecutionPolicy", - "Bypass", - "-Command", - powerShellCommandScript(command, cwdPath), - ) -} - -func powerShellCommandScript(command, cwdPath string) string { - cwdPath = powerShellSingleQuote(cwdPath) - return strings.Join([]string{ - "$__ollama_status = 0", - ". {", - "try {", - command, - " $__ollama_success = $?", - " $__ollama_last_exit = $global:LASTEXITCODE", - " if ($__ollama_success) {", - " $__ollama_status = 0", - " } elseif ($__ollama_last_exit -is [int] -and $__ollama_last_exit -ne 0) {", - " $__ollama_status = $__ollama_last_exit", - " } else {", - " $__ollama_status = 1", - " }", - "} catch {", - " Write-Error $_", - " $__ollama_status = 1", - "} finally {", - " try { [System.IO.File]::WriteAllText(" + cwdPath + ", (Get-Location).ProviderPath, [System.Text.Encoding]::UTF8) } catch {}", - "}", - "} | Out-String -Stream -Width 4096", - "exit $__ollama_status", - }, "\n") -} - -func powerShellSingleQuote(value string) string { - return "'" + strings.ReplaceAll(value, "'", "''") + "'" -} - -func runBashCommand(cmd *exec.Cmd) error { - if err := cmd.Start(); err != nil { - return err - } - if job, err := createBashJob(cmd.Process.Pid); err == nil { - bashJobHandles.Store(cmd.Process.Pid, job) - defer releaseBashJob(cmd.Process.Pid) - } - return cmd.Wait() -} - -func killBashCommand(cmd *exec.Cmd) error { - if cmd == nil || cmd.Process == nil { - return nil - } - releaseBashJob(cmd.Process.Pid) - _ = cmd.Process.Kill() - return nil -} - -func createBashJob(pid int) (windows.Handle, error) { - job, err := windows.CreateJobObject(nil, nil) - if err != nil { - return 0, err - } - - info := windows.JOBOBJECT_EXTENDED_LIMIT_INFORMATION{} - info.BasicLimitInformation.LimitFlags = windows.JOB_OBJECT_LIMIT_KILL_ON_JOB_CLOSE - if _, err := windows.SetInformationJobObject( - job, - windows.JobObjectExtendedLimitInformation, - uintptr(unsafe.Pointer(&info)), - uint32(unsafe.Sizeof(info)), - ); err != nil { - _ = windows.CloseHandle(job) - return 0, err - } - - process, err := windows.OpenProcess(windows.PROCESS_SET_QUOTA|windows.PROCESS_TERMINATE, false, uint32(pid)) - if err != nil { - _ = windows.CloseHandle(job) - return 0, err - } - defer windows.CloseHandle(process) - - if err := windows.AssignProcessToJobObject(job, process); err != nil { - _ = windows.CloseHandle(job) - return 0, err - } - return job, nil -} - -func releaseBashJob(pid int) { - value, ok := bashJobHandles.LoadAndDelete(pid) - if !ok { - return - } - if job, ok := value.(windows.Handle); ok { - _ = windows.CloseHandle(job) - } -} diff --git a/agent/tools/bash_windows_test.go b/agent/tools/bash_windows_test.go deleted file mode 100644 index 9a9be4c23e3..00000000000 --- a/agent/tools/bash_windows_test.go +++ /dev/null @@ -1,15 +0,0 @@ -//go:build windows - -package tools - -import ( - "strings" - "testing" -) - -func TestPowerShellCommandScriptUsesWideOutString(t *testing.T) { - script := powerShellCommandScript("Get-ChildItem", `C:\cwd.txt`) - if !strings.Contains(script, "Out-String -Stream -Width 4096") { - t.Fatalf("script = %q, want explicit Out-String width", script) - } -} diff --git a/agent/tools/file.go b/agent/tools/file.go deleted file mode 100644 index b12a1c0e00a..00000000000 --- a/agent/tools/file.go +++ /dev/null @@ -1,711 +0,0 @@ -package tools - -import ( - "bufio" - "cmp" - "context" - "encoding/json" - "fmt" - "io" - "os" - "path/filepath" - "slices" - "strconv" - "strings" - - "github.com/ollama/ollama/agent" - "github.com/ollama/ollama/api" -) - -const ( - maxReadBytes = 200000 -) - -type Read struct{} - -func (r *Read) Name() string { - return "read" -} - -func (r *Read) Description() string { - return "Read a text file from the current working directory." -} - -func (r *Read) Schema() api.ToolFunction { - props := api.NewToolPropertiesMap() - props.Set("path", api.ToolProperty{ - Type: api.PropertyType{"string"}, - Description: "Path to the file to read, relative to the working directory.", - }) - props.Set("start", api.ToolProperty{ - Type: api.PropertyType{"integer"}, - Description: "Optional 1-based line to start reading from.", - }) - props.Set("end", api.ToolProperty{ - Type: api.PropertyType{"integer"}, - Description: "Optional 1-based inclusive line to stop reading at.", - }) - return api.ToolFunction{ - Name: r.Name(), - Description: r.Description(), - Parameters: api.ToolFunctionParameters{ - Type: "object", - Properties: props, - Required: []string{"path"}, - }, - } -} - -func (r *Read) RequiresApproval(map[string]any) bool { - return true -} - -func (r *Read) Execute(ctx context.Context, toolCtx agent.ToolContext, args map[string]any) (agent.ToolResult, error) { - // TODO: use shared agent.RequiredStringArg / agent.OptionalIntArg for args (see agent package cleanup plan). - path, ok := args["path"].(string) - if !ok || strings.TrimSpace(path) == "" { - return agent.ToolResult{}, fmt.Errorf("path parameter is required") - } - - file, info, err := openRegularFile(toolCtx.WorkingDir, path, true) - if err != nil { - return agent.ToolResult{}, err - } - defer file.Close() - - selection, err := readSelectionFromArgs(args) - if err != nil { - return agent.ToolResult{}, err - } - if !selection.enabled && info.Size() > maxReadBytes { - return agent.ToolResult{}, fmt.Errorf("%s is too large to read (%d bytes)", path, info.Size()) - } - - select { - case <-ctx.Done(): - return agent.ToolResult{}, ctx.Err() - default: - } - - var content string - if selection.enabled { - content, err = readLineSelection(file, selection) - } else { - var contentBytes []byte - contentBytes, err = readAllWithinLimit(file, maxReadBytes) - content = string(contentBytes) - } - if err != nil { - return agent.ToolResult{}, err - } - return agent.ToolResult{Content: content}, nil -} - -type Edit struct{} - -func (e *Edit) Name() string { - return "edit" -} - -func (e *Edit) Description() string { - return "Edit a text file in the current working directory by replacing exact text. Pass multiple edits to change separate parts of the file in one call." -} - -func (e *Edit) Schema() api.ToolFunction { - editProps := api.NewToolPropertiesMap() - editProps.Set("old_text", api.ToolProperty{ - Type: api.PropertyType{"string"}, - Description: "Exact text for one targeted replacement. Must match the original file exactly once and must not overlap with any other edit's old_text.", - }) - editProps.Set("new_text", api.ToolProperty{ - Type: api.PropertyType{"string"}, - Description: "Replacement text for this targeted edit.", - }) - - props := api.NewToolPropertiesMap() - props.Set("path", api.ToolProperty{ - Type: api.PropertyType{"string"}, - Description: "Path to the file to edit, relative to the working directory.", - }) - props.Set("edits", api.ToolProperty{ - Type: api.PropertyType{"array"}, - Items: api.ToolProperty{ - Type: api.PropertyType{"object"}, - Properties: editProps, - Required: []string{"old_text", "new_text"}, - }, - Description: "One or more exact-text replacements. Each is matched against the original file, not against the output of earlier edits. Keep old_text as small as possible while still unique in the file; merge changes to the same or adjacent lines into a single edit.", - }) - props.Set("replace_all", api.ToolProperty{ - Type: api.PropertyType{"boolean"}, - Description: "Replace every occurrence. Defaults to false; only applies when a single edit is provided.", - }) - return api.ToolFunction{ - Name: e.Name(), - Description: e.Description(), - Parameters: api.ToolFunctionParameters{ - Type: "object", - Properties: props, - Required: []string{"path", "edits"}, - }, - } -} - -func (e *Edit) RequiresApproval(map[string]any) bool { - return true -} - -func (e *Edit) Execute(ctx context.Context, toolCtx agent.ToolContext, args map[string]any) (agent.ToolResult, error) { - // TODO: use shared agent.RequiredStringArg / agent.OptionalBoolArg for args (see agent package cleanup plan). - path, ok := args["path"].(string) - if !ok || strings.TrimSpace(path) == "" { - return agent.ToolResult{}, fmt.Errorf("path parameter is required") - } - - edits, replaceAll, err := parseEditArgs(args) - if err != nil { - return agent.ToolResult{}, err - } - - if err := rejectFinalSymlink(toolCtx.WorkingDir, path); err != nil { - return agent.ToolResult{}, err - } - - file, info, err := openRegularFile(toolCtx.WorkingDir, path, false) - if err != nil { - return agent.ToolResult{}, err - } - if info.Size() > maxReadBytes { - file.Close() - return agent.ToolResult{}, fmt.Errorf("%s is too large to edit (%d bytes)", path, info.Size()) - } - - select { - case <-ctx.Done(): - file.Close() - return agent.ToolResult{}, ctx.Err() - default: - } - - contentBytes, err := readAllWithinLimit(file, maxReadBytes) - if closeErr := file.Close(); err == nil && closeErr != nil { - err = closeErr - } - if err != nil { - return agent.ToolResult{}, err - } - content := string(contentBytes) - - var updated string - replacements := 0 - if replaceAll { - matches := strings.Count(content, edits[0].OldText) - if matches == 0 { - return agent.ToolResult{}, fmt.Errorf("old_text was not found in %s", path) - } - updated = strings.ReplaceAll(content, edits[0].OldText, edits[0].NewText) - replacements = matches - } else { - // Every edit is matched against the original file content rather - // than the output of earlier edits, so each edit must match exactly - // once and edits must target disjoint regions. - matched := make([]editMatch, 0, len(edits)) - for i, edit := range edits { - count := strings.Count(content, edit.OldText) - if count == 0 { - return agent.ToolResult{}, editNotFoundError(path, i, len(edits)) - } - if count > 1 { - return agent.ToolResult{}, editAmbiguousError(path, i, len(edits), count) - } - matched = append(matched, editMatch{ - editIndex: i, - offset: strings.Index(content, edit.OldText), - length: len(edit.OldText), - newText: edit.NewText, - }) - replacements++ - } - - slices.SortFunc(matched, func(a, b editMatch) int { return cmp.Compare(a.offset, b.offset) }) - for i := 1; i < len(matched); i++ { - prev, cur := matched[i-1], matched[i] - if prev.offset+prev.length > cur.offset { - return agent.ToolResult{}, fmt.Errorf("edits[%d] and edits[%d] overlap in %s; merge them into one edit or target disjoint text", prev.editIndex, cur.editIndex, path) - } - } - - // Apply from the end of the file backwards so earlier offsets stay valid. - updated = content - for i := len(matched) - 1; i >= 0; i-- { - m := matched[i] - updated = updated[:m.offset] + m.newText + updated[m.offset+m.length:] - } - } - - if updated == content { - return agent.ToolResult{}, fmt.Errorf("edit produced no changes in %s; replacement text is identical to the original", path) - } - if len(updated) > maxReadBytes { - return agent.ToolResult{}, fmt.Errorf("edited content is too large (%d bytes)", len(updated)) - } - - if err := writeFileAtomic(toolCtx.WorkingDir, path, []byte(updated), info.Mode().Perm()); err != nil { - return agent.ToolResult{}, err - } - - return agent.ToolResult{Content: fmt.Sprintf("Updated %s (%d edit%s, %d replacement%s).", path, len(edits), plural(len(edits)), replacements, plural(replacements))}, nil -} - -// editReplacement is one targeted replacement within an edit call. -type editReplacement struct { - OldText string - NewText string -} - -// editMatch locates one editReplacement within the original file content. -type editMatch struct { - editIndex int - offset int - length int - newText string -} - -// parseEditArgs normalizes edit arguments from a tool call into a list of -// replacements. It accepts the `edits` array form and tolerates legacy -// top-level old_text/new_text args as well as stringified JSON, mirroring -// the pi coding agent's argument handling. -func parseEditArgs(args map[string]any) ([]editReplacement, bool, error) { - replaceAll, _ := args["replace_all"].(bool) - - var edits []editReplacement - if raw, ok := args["edits"]; ok { - parsed, err := parseEditArray(raw) - if err != nil { - return nil, false, err - } - edits = parsed - } - - // Fold a legacy top-level old_text/new_text pair into edits. - if oldText, ok := args["old_text"].(string); ok { - newText, ok := args["new_text"].(string) - if !ok { - return nil, false, fmt.Errorf("new_text parameter is required") - } - edits = append(edits, editReplacement{OldText: oldText, NewText: newText}) - } - - if len(edits) == 0 { - return nil, false, fmt.Errorf("edits parameter is required") - } - for i, edit := range edits { - if edit.OldText == "" { - if len(edits) == 1 { - return nil, false, fmt.Errorf("old_text parameter is required") - } - return nil, false, fmt.Errorf("edits[%d].old_text must not be empty", i) - } - } - if replaceAll && len(edits) != 1 { - return nil, false, fmt.Errorf("replace_all only applies to a single edit") - } - return edits, replaceAll, nil -} - -func parseEditArray(raw any) ([]editReplacement, error) { - if s, ok := raw.(string); ok { - // Some models serialize array arguments as a JSON string. - if err := json.Unmarshal([]byte(s), &raw); err != nil { - return nil, fmt.Errorf("edits must be an array of {old_text, new_text} objects") - } - } - items, ok := raw.([]any) - if !ok { - return nil, fmt.Errorf("edits must be an array of {old_text, new_text} objects") - } - - edits := make([]editReplacement, 0, len(items)) - for i, item := range items { - entry, ok := item.(map[string]any) - if !ok { - return nil, fmt.Errorf("edits[%d] must be an object with old_text and new_text", i) - } - oldText, oldOK := editTextArg(entry, "old_text", "oldText") - newText, newOK := editTextArg(entry, "new_text", "newText") - if !oldOK || !newOK { - return nil, fmt.Errorf("edits[%d] must be an object with old_text and new_text", i) - } - edits = append(edits, editReplacement{OldText: oldText, NewText: newText}) - } - return edits, nil -} - -// editTextArg reads the first present string key, tolerating both snake_case -// and camelCase spellings that models emit. -func editTextArg(entry map[string]any, keys ...string) (string, bool) { - for _, key := range keys { - if value, ok := entry[key].(string); ok { - return value, true - } - } - return "", false -} - -func editNotFoundError(path string, editIndex, totalEdits int) error { - if totalEdits == 1 { - return fmt.Errorf("old_text was not found in %s", path) - } - return fmt.Errorf("edits[%d].old_text was not found in %s", editIndex, path) -} - -func editAmbiguousError(path string, editIndex, totalEdits, occurrences int) error { - if totalEdits == 1 { - return fmt.Errorf("old_text matched %d times in %s; set replace_all to true to replace every match", occurrences, path) - } - return fmt.Errorf("edits[%d].old_text matched %d times in %s; each edit must match exactly once, so provide more surrounding context", editIndex, occurrences, path) -} - -func cleanRelativePath(path string) (string, error) { - path = strings.TrimSpace(path) - if path == "" { - return "", fmt.Errorf("path parameter is required") - } - if filepath.IsAbs(path) { - return "", fmt.Errorf("absolute paths are not allowed") - } - cleaned := filepath.Clean(path) - if cleaned == "." || cleaned == ".." || strings.HasPrefix(cleaned, ".."+string(os.PathSeparator)) { - return "", fmt.Errorf("path escapes working directory") - } - return cleaned, nil -} - -func openRegularFile(workingDir, path string, allowAbsolute bool) (*os.File, os.FileInfo, error) { - path = strings.TrimSpace(path) - if path == "" { - return nil, nil, fmt.Errorf("path parameter is required") - } - if allowAbsolute && filepath.IsAbs(path) { - cleaned := filepath.Clean(path) - info, err := os.Lstat(cleaned) - if err != nil { - return nil, nil, err - } - if info.Mode()&os.ModeSymlink != 0 { - return nil, nil, fmt.Errorf("%s is a symlink; read the target file directly", path) - } - if err := rejectNonRegularFile(path, info); err != nil { - return nil, nil, err - } - file, err := os.Open(cleaned) - if err != nil { - return nil, nil, err - } - info, err = file.Stat() - if err != nil { - file.Close() - return nil, nil, err - } - if err := rejectNonRegularFile(path, info); err != nil { - file.Close() - return nil, nil, err - } - return file, info, nil - } - - rel, err := cleanRelativePath(path) - if err != nil { - return nil, nil, err - } - root, err := openWorkingRoot(workingDir) - if err != nil { - return nil, nil, err - } - defer root.Close() - - if _, err := regularRootFileInfo(root, rel, path); err != nil { - return nil, nil, err - } - file, err := root.Open(rel) - if err != nil { - return nil, nil, rootPathError(err) - } - info, err := file.Stat() - if err != nil { - file.Close() - return nil, nil, err - } - if err := rejectNonRegularFile(path, info); err != nil { - file.Close() - return nil, nil, err - } - return file, info, nil -} - -func regularRootFileInfo(root *os.Root, rel, path string) (os.FileInfo, error) { - info, err := root.Lstat(rel) - if err != nil { - return nil, rootPathError(err) - } - // Reject symlinks outright. os.Root.Open follows symlinks via openat - // without O_NOFOLLOW, so a symlink inside the working root that points - // outside it (e.g. ./notes -> ~/.ssh/id_rsa) would otherwise be read - // transparently, bypassing the working-directory confinement that the - // bash denylist enforces for direct credential reads. The caller must - // operate on the real target file instead. - if info.Mode()&os.ModeSymlink != 0 { - return nil, fmt.Errorf("%s is a symlink; read the target file directly", path) - } - if err := rejectNonRegularFile(path, info); err != nil { - return nil, err - } - return info, nil -} - -func rejectNonRegularFile(path string, info os.FileInfo) error { - if info.IsDir() { - return fmt.Errorf("%s is a directory", path) - } - if !info.Mode().IsRegular() { - return fmt.Errorf("%s is not a regular file", path) - } - return nil -} - -func writeFileAtomic(workingDir, path string, data []byte, perm os.FileMode) error { - rel, err := cleanRelativePath(path) - if err != nil { - return err - } - root, err := openWorkingRoot(workingDir) - if err != nil { - return err - } - defer root.Close() - if err := rejectRootFinalSymlink(root, rel, path); err != nil { - return err - } - - parent, name := filepath.Split(rel) - tmpBase := fmt.Sprintf(".%s.ollama-tmp-%d", name, os.Getpid()) - for i := 0; ; i++ { - candidateName := tmpBase - if i > 0 { - candidateName = fmt.Sprintf("%s-%d", tmpBase, i) - } - candidate := filepath.Join(parent, candidateName) - file, err := root.OpenFile(candidate, os.O_WRONLY|os.O_CREATE|os.O_EXCL, perm) - if os.IsExist(err) { - continue - } - if err != nil { - return rootPathError(err) - } - if err := file.Chmod(perm); err != nil { - closeErr := file.Close() - _ = root.Remove(candidate) - if closeErr != nil { - return closeErr - } - return err - } - writeErr := writeAllAndSync(file, data) - closeErr := file.Close() - if writeErr != nil || closeErr != nil { - _ = root.Remove(candidate) - if writeErr != nil { - return writeErr - } - return closeErr - } - if err := root.Rename(candidate, rel); err != nil { - _ = root.Remove(candidate) - return rootPathError(err) - } - return nil - } -} - -func rejectFinalSymlink(workingDir, path string) error { - rel, err := cleanRelativePath(path) - if err != nil { - return err - } - root, err := openWorkingRoot(workingDir) - if err != nil { - return err - } - defer root.Close() - return rejectRootFinalSymlink(root, rel, path) -} - -func rejectRootFinalSymlink(root *os.Root, rel, path string) error { - info, err := root.Lstat(rel) - if err != nil { - return rootPathError(err) - } - if info.Mode()&os.ModeSymlink != 0 { - return fmt.Errorf("%s is a symlink; edit the target file directly", path) - } - return nil -} - -func rootPathError(err error) error { - if err != nil && strings.Contains(err.Error(), "path escapes") { - return fmt.Errorf("path escapes working directory") - } - return err -} - -func openWorkingRoot(workingDir string) (*os.Root, error) { - base, err := workingDirAbs(workingDir) - if err != nil { - return nil, err - } - return os.OpenRoot(base) -} - -func writeAllAndSync(file *os.File, data []byte) error { - if _, err := file.Write(data); err != nil { - return err - } - return file.Sync() -} - -func readAllWithinLimit(reader io.Reader, limit int) ([]byte, error) { - if limit < 0 { - limit = 0 - } - content, err := io.ReadAll(io.LimitReader(reader, int64(limit)+1)) - if err != nil { - return nil, err - } - if len(content) > limit { - return nil, fmt.Errorf("content is too large (%d byte limit)", limit) - } - return content, nil -} - -func workingDirAbs(workingDir string) (string, error) { - base := workingDir - if base == "" { - var err error - base, err = os.Getwd() - if err != nil { - return "", err - } - } - return canonicalPath(base) -} - -func canonicalPath(path string) (string, error) { - abs, err := filepath.Abs(path) - if err != nil { - return "", err - } - resolved, err := filepath.EvalSymlinks(abs) - if err == nil { - return resolved, nil - } - return abs, nil -} - -type readSelection struct { - enabled bool - start int - end int -} - -func readSelectionFromArgs(args map[string]any) (readSelection, error) { - selection := readSelection{start: 1} - - if start, ok, err := intReadArg(args, "start"); err != nil { - return readSelection{}, err - } else if ok { - selection.enabled = true - selection.start = start - } - if end, ok, err := intReadArg(args, "end"); err != nil { - return readSelection{}, err - } else if ok { - selection.enabled = true - selection.end = end - } - - if !selection.enabled { - return selection, nil - } - if selection.start < 1 { - return readSelection{}, fmt.Errorf("start must be greater than 0") - } - if selection.end > 0 && selection.end < selection.start { - return readSelection{}, fmt.Errorf("end must be greater than or equal to start") - } - return selection, nil -} - -func readLineSelection(file *os.File, selection readSelection) (string, error) { - reader := bufio.NewReader(file) - var b strings.Builder - for lineNo := 1; ; { - line, err := reader.ReadSlice('\n') - if lineNo >= selection.start && (selection.end == 0 || lineNo <= selection.end) { - if b.Len()+len(line) > maxReadBytes { - return "", fmt.Errorf("selected content is too large (%d byte limit)", maxReadBytes) - } - b.Write(line) - } - if err != nil { - if err == bufio.ErrBufferFull { - continue - } - if err == io.EOF { - break - } - return "", err - } - if selection.end > 0 && lineNo >= selection.end { - break - } - lineNo++ - } - return b.String(), nil -} - -func intReadArg(args map[string]any, key string) (int, bool, error) { - value, ok := args[key] - if !ok { - return 0, false, nil - } - switch v := value.(type) { - case int: - return v, true, nil - case int64: - return int(v), true, nil - case float64: - if v != float64(int(v)) { - return 0, true, fmt.Errorf("%s must be a whole number", key) - } - return int(v), true, nil - case string: - v = strings.TrimSpace(v) - if v == "" { - return 0, false, nil - } - n, err := strconv.Atoi(v) - if err != nil { - return 0, true, fmt.Errorf("%s must be a whole number", key) - } - return n, true, nil - default: - return 0, true, fmt.Errorf("%s must be a whole number", key) - } -} - -func plural(n int) string { - if n == 1 { - return "" - } - return "s" -} diff --git a/agent/tools/file_test.go b/agent/tools/file_test.go deleted file mode 100644 index 15f28ab2792..00000000000 --- a/agent/tools/file_test.go +++ /dev/null @@ -1,571 +0,0 @@ -package tools - -import ( - "context" - "io" - "os" - "path/filepath" - "strings" - "testing" - - "github.com/ollama/ollama/agent" -) - -func TestEditReplacesUniqueText(t *testing.T) { - dir := t.TempDir() - path := filepath.Join(dir, "note.txt") - if err := os.WriteFile(path, []byte("hello world\n"), 0o644); err != nil { - t.Fatal(err) - } - - result, err := (&Edit{}).Execute(context.Background(), agent.ToolContext{WorkingDir: dir}, map[string]any{ - "path": "note.txt", - "old_text": "hello", - "new_text": "hi", - }) - if err != nil { - t.Fatal(err) - } - if !strings.Contains(result.Content, "Updated note.txt") { - t.Fatalf("result = %q", result.Content) - } - - content, err := os.ReadFile(path) - if err != nil { - t.Fatal(err) - } - if string(content) != "hi world\n" { - t.Fatalf("content = %q", content) - } -} - -func TestEditRequiresUniqueMatchByDefault(t *testing.T) { - dir := t.TempDir() - path := filepath.Join(dir, "note.txt") - if err := os.WriteFile(path, []byte("same same\n"), 0o644); err != nil { - t.Fatal(err) - } - - _, err := (&Edit{}).Execute(context.Background(), agent.ToolContext{WorkingDir: dir}, map[string]any{ - "path": "note.txt", - "old_text": "same", - "new_text": "other", - }) - if err == nil { - t.Fatal("expected ambiguous edit to fail") - } - if !strings.Contains(err.Error(), "matched 2 times") { - t.Fatalf("err = %v", err) - } -} - -func TestEditAppliesMultipleEdits(t *testing.T) { - dir := t.TempDir() - path := filepath.Join(dir, "note.txt") - if err := os.WriteFile(path, []byte("alpha beta gamma delta\n"), 0o644); err != nil { - t.Fatal(err) - } - - result, err := (&Edit{}).Execute(context.Background(), agent.ToolContext{WorkingDir: dir}, map[string]any{ - "path": "note.txt", - "edits": []any{ - map[string]any{"old_text": "beta", "new_text": "BETA"}, - map[string]any{"old_text": "delta", "new_text": "DELTA"}, - }, - }) - if err != nil { - t.Fatal(err) - } - if !strings.Contains(result.Content, "2 edits, 2 replacements") { - t.Fatalf("result = %q", result.Content) - } - - content, err := os.ReadFile(path) - if err != nil { - t.Fatal(err) - } - if string(content) != "alpha BETA gamma DELTA\n" { - t.Fatalf("content = %q", content) - } -} - -func TestEditMatchesEditsAgainstOriginalContent(t *testing.T) { - dir := t.TempDir() - path := filepath.Join(dir, "note.txt") - if err := os.WriteFile(path, []byte("abc def\n"), 0o644); err != nil { - t.Fatal(err) - } - - // edits[1] must target the original "def", not the one introduced by edits[0]. - _, err := (&Edit{}).Execute(context.Background(), agent.ToolContext{WorkingDir: dir}, map[string]any{ - "path": "note.txt", - "edits": []any{ - map[string]any{"old_text": "abc", "new_text": "def"}, - map[string]any{"old_text": "def", "new_text": "ghi"}, - }, - }) - if err != nil { - t.Fatal(err) - } - - content, err := os.ReadFile(path) - if err != nil { - t.Fatal(err) - } - if string(content) != "def ghi\n" { - t.Fatalf("content = %q", content) - } -} - -func TestEditRejectsOverlappingEdits(t *testing.T) { - dir := t.TempDir() - path := filepath.Join(dir, "note.txt") - if err := os.WriteFile(path, []byte("abc\n"), 0o644); err != nil { - t.Fatal(err) - } - - _, err := (&Edit{}).Execute(context.Background(), agent.ToolContext{WorkingDir: dir}, map[string]any{ - "path": "note.txt", - "edits": []any{ - map[string]any{"old_text": "ab", "new_text": "x"}, - map[string]any{"old_text": "bc", "new_text": "y"}, - }, - }) - if err == nil { - t.Fatal("expected overlapping edits to fail") - } - if !strings.Contains(err.Error(), "overlap") { - t.Fatalf("err = %v", err) - } -} - -func TestEditMultipleEditsNotFoundIndexed(t *testing.T) { - dir := t.TempDir() - path := filepath.Join(dir, "note.txt") - if err := os.WriteFile(path, []byte("hello world\n"), 0o644); err != nil { - t.Fatal(err) - } - - _, err := (&Edit{}).Execute(context.Background(), agent.ToolContext{WorkingDir: dir}, map[string]any{ - "path": "note.txt", - "edits": []any{ - map[string]any{"old_text": "hello", "new_text": "hi"}, - map[string]any{"old_text": "missing", "new_text": "x"}, - }, - }) - if err == nil { - t.Fatal("expected missing edit to fail") - } - if !strings.Contains(err.Error(), "edits[1]") { - t.Fatalf("err = %v", err) - } -} - -func TestEditMultipleEditsAmbiguousIndexed(t *testing.T) { - dir := t.TempDir() - path := filepath.Join(dir, "note.txt") - if err := os.WriteFile(path, []byte("hello same same\n"), 0o644); err != nil { - t.Fatal(err) - } - - _, err := (&Edit{}).Execute(context.Background(), agent.ToolContext{WorkingDir: dir}, map[string]any{ - "path": "note.txt", - "edits": []any{ - map[string]any{"old_text": "hello", "new_text": "hi"}, - map[string]any{"old_text": "same", "new_text": "x"}, - }, - }) - if err == nil { - t.Fatal("expected ambiguous edit to fail") - } - if !strings.Contains(err.Error(), "edits[1]") || !strings.Contains(err.Error(), "matched 2 times") { - t.Fatalf("err = %v", err) - } -} - -func TestEditRejectsEmptyEdits(t *testing.T) { - dir := t.TempDir() - path := filepath.Join(dir, "note.txt") - if err := os.WriteFile(path, []byte("hello\n"), 0o644); err != nil { - t.Fatal(err) - } - - for name, args := range map[string]map[string]any{ - "missing edits": {"path": "note.txt"}, - "empty edits": {"path": "note.txt", "edits": []any{}}, - } { - if _, err := (&Edit{}).Execute(context.Background(), agent.ToolContext{WorkingDir: dir}, args); err == nil { - t.Fatalf("%s: expected error", name) - } else if !strings.Contains(err.Error(), "edits parameter is required") { - t.Fatalf("%s: err = %v", name, err) - } - } -} - -func TestEditRejectsEmptyOldTextInArray(t *testing.T) { - dir := t.TempDir() - path := filepath.Join(dir, "note.txt") - if err := os.WriteFile(path, []byte("hello\n"), 0o644); err != nil { - t.Fatal(err) - } - - _, err := (&Edit{}).Execute(context.Background(), agent.ToolContext{WorkingDir: dir}, map[string]any{ - "path": "note.txt", - "edits": []any{ - map[string]any{"old_text": "hello", "new_text": "hi"}, - map[string]any{"old_text": "", "new_text": "x"}, - }, - }) - if err == nil { - t.Fatal("expected empty old_text to fail") - } - if !strings.Contains(err.Error(), "edits[1].old_text must not be empty") { - t.Fatalf("err = %v", err) - } -} - -func TestEditAcceptsJSONStringEdits(t *testing.T) { - dir := t.TempDir() - path := filepath.Join(dir, "note.txt") - if err := os.WriteFile(path, []byte("hello world\n"), 0o644); err != nil { - t.Fatal(err) - } - - // Some models serialize array arguments as a JSON string. - _, err := (&Edit{}).Execute(context.Background(), agent.ToolContext{WorkingDir: dir}, map[string]any{ - "path": "note.txt", - "edits": `[{"oldText": "hello", "newText": "hi"}, {"oldText": "world", "newText": "earth"}]`, - }) - if err != nil { - t.Fatal(err) - } - - content, err := os.ReadFile(path) - if err != nil { - t.Fatal(err) - } - if string(content) != "hi earth\n" { - t.Fatalf("content = %q", content) - } -} - -func TestEditRejectsReplaceAllWithMultipleEdits(t *testing.T) { - dir := t.TempDir() - path := filepath.Join(dir, "note.txt") - if err := os.WriteFile(path, []byte("a b c\n"), 0o644); err != nil { - t.Fatal(err) - } - - _, err := (&Edit{}).Execute(context.Background(), agent.ToolContext{WorkingDir: dir}, map[string]any{ - "path": "note.txt", - "replace_all": true, - "edits": []any{ - map[string]any{"old_text": "a", "new_text": "x"}, - map[string]any{"old_text": "b", "new_text": "y"}, - }, - }) - if err == nil { - t.Fatal("expected replace_all with multiple edits to fail") - } - if !strings.Contains(err.Error(), "replace_all") { - t.Fatalf("err = %v", err) - } -} - -func TestEditRejectsNoChange(t *testing.T) { - dir := t.TempDir() - path := filepath.Join(dir, "note.txt") - if err := os.WriteFile(path, []byte("hello\n"), 0o644); err != nil { - t.Fatal(err) - } - - _, err := (&Edit{}).Execute(context.Background(), agent.ToolContext{WorkingDir: dir}, map[string]any{ - "path": "note.txt", - "old_text": "hello", - "new_text": "hello", - }) - if err == nil { - t.Fatal("expected no-change edit to fail") - } - if !strings.Contains(err.Error(), "no changes") { - t.Fatalf("err = %v", err) - } -} - -func TestEditRejectsEscapingPath(t *testing.T) { - dir := t.TempDir() - _, err := (&Edit{}).Execute(context.Background(), agent.ToolContext{WorkingDir: dir}, map[string]any{ - "path": "../outside.txt", - "old_text": "old", - "new_text": "new", - }) - if err == nil { - t.Fatal("expected escaping path to fail") - } - if !strings.Contains(err.Error(), "path escapes working directory") { - t.Fatalf("err = %v", err) - } -} - -func TestEditRejectsSymlinkEscape(t *testing.T) { - dir := t.TempDir() - outside := t.TempDir() - if err := os.WriteFile(filepath.Join(outside, "note.txt"), []byte("old\n"), 0o644); err != nil { - t.Fatal(err) - } - if err := os.Symlink(outside, filepath.Join(dir, "link")); err != nil { - t.Skipf("symlinks unavailable: %v", err) - } - - _, err := (&Edit{}).Execute(context.Background(), agent.ToolContext{WorkingDir: dir}, map[string]any{ - "path": filepath.Join("link", "note.txt"), - "old_text": "old", - "new_text": "new", - }) - if err == nil { - t.Fatal("expected symlink escape to fail") - } - if !strings.Contains(err.Error(), "path escapes working directory") { - t.Fatalf("err = %v", err) - } - - content, err := os.ReadFile(filepath.Join(outside, "note.txt")) - if err != nil { - t.Fatal(err) - } - if string(content) != "old\n" { - t.Fatalf("outside content changed to %q", content) - } -} - -func TestEditRejectsFinalSymlink(t *testing.T) { - dir := t.TempDir() - target := filepath.Join(dir, "target.txt") - if err := os.WriteFile(target, []byte("old\n"), 0o644); err != nil { - t.Fatal(err) - } - link := filepath.Join(dir, "link.txt") - if err := os.Symlink("target.txt", link); err != nil { - t.Skipf("symlinks unavailable: %v", err) - } - - _, err := (&Edit{}).Execute(context.Background(), agent.ToolContext{WorkingDir: dir}, map[string]any{ - "path": "link.txt", - "old_text": "old", - "new_text": "new", - }) - if err == nil { - t.Fatal("expected final symlink edit to fail") - } - if !strings.Contains(err.Error(), "is a symlink") { - t.Fatalf("err = %v", err) - } - content, err := os.ReadFile(target) - if err != nil { - t.Fatal(err) - } - if string(content) != "old\n" { - t.Fatalf("target content changed to %q", content) - } - info, err := os.Lstat(link) - if err != nil { - t.Fatal(err) - } - if info.Mode()&os.ModeSymlink == 0 { - t.Fatalf("link mode = %v, want symlink", info.Mode()) - } -} - -func TestReadRejectsParentOutsideCurrentWorkingDir(t *testing.T) { - root := t.TempDir() - subdir := filepath.Join(root, "sub") - if err := os.Mkdir(subdir, 0o755); err != nil { - t.Fatal(err) - } - if err := os.WriteFile(filepath.Join(root, "note.txt"), []byte("hello"), 0o644); err != nil { - t.Fatal(err) - } - - _, err := (&Read{}).Execute(context.Background(), agent.ToolContext{WorkingDir: subdir}, map[string]any{ - "path": "../note.txt", - }) - if err == nil { - t.Fatal("expected parent path to fail") - } - if !strings.Contains(err.Error(), "path escapes working directory") { - t.Fatalf("err = %v", err) - } -} - -func TestReadRequiresApproval(t *testing.T) { - if !agent.ToolRequiresApproval((&Read{}), map[string]any{"path": "note.txt"}) { - t.Fatal("read should require approval") - } -} - -func TestReadDefaultsToEntireFile(t *testing.T) { - dir := t.TempDir() - content := "one\ntwo\nthree\n" - if err := os.WriteFile(filepath.Join(dir, "note.txt"), []byte(content), 0o644); err != nil { - t.Fatal(err) - } - - result, err := (&Read{}).Execute(context.Background(), agent.ToolContext{WorkingDir: dir}, map[string]any{ - "path": "note.txt", - }) - if err != nil { - t.Fatal(err) - } - if result.Content != content { - t.Fatalf("content = %q", result.Content) - } -} - -func TestReadAllowsAbsolutePath(t *testing.T) { - dir := t.TempDir() - path := filepath.Join(dir, "note.txt") - content := "one\ntwo\nthree\n" - if err := os.WriteFile(path, []byte(content), 0o644); err != nil { - t.Fatal(err) - } - - result, err := (&Read{}).Execute(context.Background(), agent.ToolContext{WorkingDir: t.TempDir()}, map[string]any{ - "path": path, - }) - if err != nil { - t.Fatal(err) - } - if result.Content != content { - t.Fatalf("content = %q", result.Content) - } -} - -func TestReadRejectsAbsoluteSymlink(t *testing.T) { - dir := t.TempDir() - target := filepath.Join(dir, "target.txt") - if err := os.WriteFile(target, []byte("hello\n"), 0o644); err != nil { - t.Fatal(err) - } - link := filepath.Join(dir, "alias") - if err := os.Symlink(target, link); err != nil { - t.Skipf("symlinks unavailable: %v", err) - } - - _, err := (&Read{}).Execute(context.Background(), agent.ToolContext{WorkingDir: t.TempDir()}, map[string]any{ - "path": link, - }) - if err == nil { - t.Fatal("expected absolute symlink to be rejected") - } - if !strings.Contains(err.Error(), "symlink") { - t.Fatalf("err = %v, want symlink rejection", err) - } -} - -func TestReadStartEnd(t *testing.T) { - dir := t.TempDir() - if err := os.WriteFile(filepath.Join(dir, "note.txt"), []byte("one\ntwo\nthree\nfour\n"), 0o644); err != nil { - t.Fatal(err) - } - - result, err := (&Read{}).Execute(context.Background(), agent.ToolContext{WorkingDir: dir}, map[string]any{ - "path": "note.txt", - "start": 2, - "end": 3, - }) - if err != nil { - t.Fatal(err) - } - if result.Content != "two\nthree\n" { - t.Fatalf("content = %q", result.Content) - } -} - -func TestReadStartOnly(t *testing.T) { - dir := t.TempDir() - if err := os.WriteFile(filepath.Join(dir, "note.txt"), []byte("one\ntwo\nthree\nfour\n"), 0o644); err != nil { - t.Fatal(err) - } - - result, err := (&Read{}).Execute(context.Background(), agent.ToolContext{WorkingDir: dir}, map[string]any{ - "path": "note.txt", - "start": 3, - }) - if err != nil { - t.Fatal(err) - } - if result.Content != "three\nfour\n" { - t.Fatalf("content = %q", result.Content) - } -} - -func TestReadEndOnly(t *testing.T) { - dir := t.TempDir() - if err := os.WriteFile(filepath.Join(dir, "note.txt"), []byte("one\ntwo\nthree\nfour\n"), 0o644); err != nil { - t.Fatal(err) - } - - result, err := (&Read{}).Execute(context.Background(), agent.ToolContext{WorkingDir: dir}, map[string]any{ - "path": "note.txt", - "end": 2, - }) - if err != nil { - t.Fatal(err) - } - if result.Content != "one\ntwo\n" { - t.Fatalf("content = %q", result.Content) - } -} - -func TestReadSelectionRejectsHugeSingleLine(t *testing.T) { - dir := t.TempDir() - if err := os.WriteFile(filepath.Join(dir, "note.txt"), []byte(strings.Repeat("x", maxReadBytes+1)), 0o644); err != nil { - t.Fatal(err) - } - - _, err := (&Read{}).Execute(context.Background(), agent.ToolContext{WorkingDir: dir}, map[string]any{ - "path": "note.txt", - "start": 1, - "end": 1, - }) - if err == nil { - t.Fatal("expected huge selected line to fail") - } - if !strings.Contains(err.Error(), "selected content is too large") { - t.Fatalf("err = %v", err) - } -} - -func TestReadAllWithinLimitRejectsGrowingRead(t *testing.T) { - reader := io.MultiReader( - strings.NewReader(strings.Repeat("x", maxReadBytes)), - strings.NewReader("x"), - ) - - _, err := readAllWithinLimit(reader, maxReadBytes) - if err == nil { - t.Fatal("expected over-limit read to fail") - } - if !strings.Contains(err.Error(), "content is too large") { - t.Fatalf("err = %v", err) - } -} - -func TestReadRejectsInvalidRange(t *testing.T) { - dir := t.TempDir() - if err := os.WriteFile(filepath.Join(dir, "note.txt"), []byte("one\ntwo\n"), 0o644); err != nil { - t.Fatal(err) - } - - _, err := (&Read{}).Execute(context.Background(), agent.ToolContext{WorkingDir: dir}, map[string]any{ - "path": "note.txt", - "start": 4, - "end": 2, - }) - if err == nil { - t.Fatal("expected invalid range to fail") - } - if !strings.Contains(err.Error(), "end must") { - t.Fatalf("err = %v", err) - } -} diff --git a/agent/tools/file_unix_test.go b/agent/tools/file_unix_test.go deleted file mode 100644 index a0973abc203..00000000000 --- a/agent/tools/file_unix_test.go +++ /dev/null @@ -1,121 +0,0 @@ -//go:build !windows - -package tools - -import ( - "context" - "os" - "path/filepath" - "strings" - "syscall" - "testing" - "time" - - "github.com/ollama/ollama/agent" -) - -func TestOpenRegularFileRejectsFIFO(t *testing.T) { - dir := t.TempDir() - path := filepath.Join(dir, "pipe") - if err := syscall.Mkfifo(path, 0o600); err != nil { - t.Skipf("mkfifo unavailable: %v", err) - } - - done := make(chan error, 1) - go func() { - file, _, err := openRegularFile(dir, "pipe", false) - if file != nil { - file.Close() - } - done <- err - }() - - select { - case err := <-done: - if err == nil { - t.Fatal("expected FIFO to be rejected") - } - if !strings.Contains(err.Error(), "not a regular file") { - t.Fatalf("err = %v", err) - } - case <-time.After(time.Second): - t.Fatal("openRegularFile blocked on FIFO") - } -} - -func TestEditPreservesModeDespiteUmask(t *testing.T) { - oldUmask := syscall.Umask(0o077) - defer syscall.Umask(oldUmask) - - dir := t.TempDir() - path := filepath.Join(dir, "note.txt") - if err := os.WriteFile(path, []byte("hello\n"), 0o666); err != nil { - t.Fatal(err) - } - if err := os.Chmod(path, 0o666); err != nil { - t.Fatal(err) - } - - _, err := (&Edit{}).Execute(context.Background(), agent.ToolContext{WorkingDir: dir}, map[string]any{ - "path": "note.txt", - "old_text": "hello", - "new_text": "hi", - }) - if err != nil { - t.Fatal(err) - } - - info, err := os.Stat(path) - if err != nil { - t.Fatal(err) - } - if got := info.Mode().Perm(); got != 0o666 { - t.Fatalf("mode = %#o, want 0666", got) - } -} - -func TestReadRejectsSymlinkEscapingWorkingDir(t *testing.T) { - root := t.TempDir() - secret := filepath.Join(t.TempDir(), "secret.txt") - if err := os.WriteFile(secret, []byte("top secret\n"), 0o600); err != nil { - t.Fatal(err) - } - link := filepath.Join(root, "notes") - if err := os.Symlink(secret, link); err != nil { - t.Fatal(err) - } - - _, err := (&Read{}).Execute(context.Background(), agent.ToolContext{WorkingDir: root}, map[string]any{ - "path": "notes", - }) - if err == nil { - t.Fatal("expected symlink escaping working dir to be rejected") - } - if !strings.Contains(err.Error(), "symlink") { - t.Fatalf("err = %v, want symlink rejection", err) - } -} - -func TestReadRejectsSymlinkInsideWorkingDirToOutside(t *testing.T) { - root := t.TempDir() - target := filepath.Join(root, "real.txt") - if err := os.WriteFile(target, []byte("hello\n"), 0o644); err != nil { - t.Fatal(err) - } - // A symlink to a sibling file still resolves inside the root; Read must - // reject it regardless, consistent with Edit's rejectFinalSymlink. - link := filepath.Join(root, "alias") - if err := os.Symlink(target, link); err != nil { - t.Fatal(err) - } - - _, err := (&Read{}).Execute(context.Background(), agent.ToolContext{WorkingDir: root}, map[string]any{ - "path": "alias", - }) - if err == nil { - t.Fatal("expected symlink to be rejected even when target is inside root") - } - if !strings.Contains(err.Error(), "symlink") { - t.Fatalf("err = %v, want symlink rejection", err) - } -} diff --git a/agent/tools/skill.go b/agent/tools/skill.go deleted file mode 100644 index 37620b6b239..00000000000 --- a/agent/tools/skill.go +++ /dev/null @@ -1,41 +0,0 @@ -package tools - -import ( - "context" - "errors" - - "github.com/ollama/ollama/agent" - "github.com/ollama/ollama/api" -) - -// Skill is the model-facing adapter for the core agent skill catalog. -// Model-initiated loads require approval because a skill's instructions can -// influence the rest of the run. Explicit user activation is handled by the -// session's synthetic skill call and bypasses this adapter. -type Skill struct{ Catalog *agent.SkillCatalog } - -func (t *Skill) Name() string { return "skill" } - -func (t *Skill) Description() string { - return "Load a named Ollama skill and return its instructions." -} - -func (t *Skill) Schema() api.ToolFunction { - props := api.NewToolPropertiesMap() - props.Set("name", api.ToolProperty{Type: api.PropertyType{"string"}, Description: "Name of the skill to load."}) - return api.ToolFunction{Name: t.Name(), Description: t.Description(), Parameters: api.ToolFunctionParameters{Type: "object", Properties: props, Required: []string{"name"}}} -} - -func (t *Skill) RequiresApproval(map[string]any) bool { return true } - -func (t *Skill) Execute(_ context.Context, _ agent.ToolContext, args map[string]any) (agent.ToolResult, error) { - name, ok := args["name"].(string) - if !ok { - return agent.ToolResult{}, errors.New("name parameter is required") - } - skill, err := t.Catalog.Load(name) - if err != nil { - return agent.ToolResult{}, err - } - return agent.ToolResult{Content: skill.Content()}, nil -} diff --git a/agent/tools/skill_test.go b/agent/tools/skill_test.go deleted file mode 100644 index 3835325abc0..00000000000 --- a/agent/tools/skill_test.go +++ /dev/null @@ -1,163 +0,0 @@ -package tools - -import ( - "context" - "os" - "path/filepath" - "strings" - "testing" - - "github.com/ollama/ollama/agent" - "github.com/ollama/ollama/api" -) - -func TestSkillLoadsCoreCatalogWithApproval(t *testing.T) { - catalog := testSkillCatalog(t) - tool := &Skill{Catalog: catalog} - if !agent.ToolRequiresApproval(tool, map[string]any{"name": "release-notes"}) { - t.Fatal("model-initiated skill loading should require approval") - } - result, err := tool.Execute(context.Background(), agent.ToolContext{}, map[string]any{"name": "release-notes"}) - if err != nil || !strings.Contains(result.Content, "Use concise bullets.") { - t.Fatalf("tool result = %#v, %v", result, err) - } -} - -func TestModelSkillLoadRequiresApproval(t *testing.T) { - for _, tt := range []struct { - name string - approval agent.Approval - prompt bool - wantCalls int - wantPrompts int - wantResult string - }{ - {name: "rejected", approval: agent.Approval{Reason: "Skill loading denied."}, prompt: true, wantCalls: 1, wantPrompts: 1, wantResult: "Skill loading denied."}, - {name: "approved", approval: agent.Approval{Allow: true}, prompt: true, wantCalls: 2, wantPrompts: 1, wantResult: "Use concise bullets."}, - {name: "headless denied", wantCalls: 1, wantResult: "Tool execution requires approval"}, - } { - t.Run(tt.name, func(t *testing.T) { - catalog := testSkillCatalog(t) - args := api.NewToolCallFunctionArguments() - args.Set("name", "release-notes") - client := &skillTestClient{responses: [][]api.ChatResponse{ - {{Message: api.Message{Role: "assistant", ToolCalls: []api.ToolCall{{ - ID: "call_skill_1", - Function: api.ToolCallFunction{Name: "skill", Arguments: args}, - }}}}}, - {{Message: api.Message{Role: "assistant", Content: "done"}}}, - }} - var prompter *skillApprovalPrompter - var approvalPrompter agent.ApprovalPrompter - if tt.prompt { - prompter = &skillApprovalPrompter{result: tt.approval} - approvalPrompter = prompter - } - registry := &agent.Registry{} - registry.Register(&Skill{Catalog: catalog}) - - result, err := (&agent.Session{ - Client: client, - Tools: registry, - ApprovalPrompter: approvalPrompter, - }).Run(context.Background(), agent.RunOptions{ - Model: "test", - NewMessages: []api.Message{{Role: "user", Content: "load the release-notes skill"}}, - }) - if err != nil { - t.Fatal(err) - } - if tt.prompt { - if got := len(prompter.requests); got != tt.wantPrompts { - t.Fatalf("approval prompts = %d, want %d", got, tt.wantPrompts) - } - request := prompter.requests[0] - if len(request.Calls) != 1 || request.Calls[0].ToolName != "skill" || request.Calls[0].ApprovalScope != "skill" || request.Calls[0].Args["name"] != "release-notes" { - t.Fatalf("approval request = %#v", request) - } - } - if got := client.calls; got != tt.wantCalls { - t.Fatalf("model calls = %d, want %d", got, tt.wantCalls) - } - var toolResult string - for _, message := range result.Messages { - if message.Role == "tool" && message.ToolCallID == "call_skill_1" { - toolResult = message.Content - break - } - } - if !strings.Contains(toolResult, tt.wantResult) { - t.Fatalf("skill tool result = %q, want it to contain %q", toolResult, tt.wantResult) - } - }) - } -} - -func TestExplicitSkillActivationBypassesApproval(t *testing.T) { - catalog := testSkillCatalog(t) - client := &skillTestClient{responses: [][]api.ChatResponse{{{Message: api.Message{Role: "assistant", Content: "done"}}}}} - prompter := &skillApprovalPrompter{result: agent.Approval{}} - result, err := (&agent.Session{ - Client: client, - Skills: catalog, - ApprovalPrompter: prompter, - }).Run(context.Background(), agent.RunOptions{ - Model: "test", - NewMessages: []api.Message{{Role: "user", Content: "draft release notes"}}, - SkillName: "release-notes", - }) - if err != nil { - t.Fatal(err) - } - if len(prompter.requests) != 0 { - t.Fatalf("explicit activation prompted for approval: %#v", prompter.requests) - } - if len(result.Messages) != 4 || result.Messages[2].ToolName != "skill" || !strings.Contains(result.Messages[2].Content, "Use concise bullets.") { - t.Fatalf("synthetic skill activation = %#v", result.Messages) - } -} - -func testSkillCatalog(t *testing.T) *agent.SkillCatalog { - t.Helper() - dir := t.TempDir() - path := filepath.Join(dir, "release-notes") - if err := os.Mkdir(path, 0o755); err != nil { - t.Fatal(err) - } - if err := os.WriteFile(filepath.Join(path, "SKILL.md"), []byte("---\nname: release-notes\ndescription: Draft release notes.\n---\nUse concise bullets."), 0o644); err != nil { - t.Fatal(err) - } - catalog, err := agent.DiscoverSkills(dir) - if err != nil { - t.Fatal(err) - } - return catalog -} - -type skillTestClient struct { - responses [][]api.ChatResponse - calls int -} - -func (c *skillTestClient) Chat(_ context.Context, _ *api.ChatRequest, fn api.ChatResponseFunc) error { - if c.calls >= len(c.responses) { - return nil - } - for _, response := range c.responses[c.calls] { - if err := fn(response); err != nil { - return err - } - } - c.calls++ - return nil -} - -type skillApprovalPrompter struct { - requests []agent.ApprovalRequest - result agent.Approval -} - -func (p *skillApprovalPrompter) PromptApproval(_ context.Context, request agent.ApprovalRequest) (agent.Approval, error) { - p.requests = append(p.requests, request) - return p.result, nil -} diff --git a/agent/tools/web.go b/agent/tools/web.go deleted file mode 100644 index fe6849ad61b..00000000000 --- a/agent/tools/web.go +++ /dev/null @@ -1,186 +0,0 @@ -package tools - -import ( - "context" - "errors" - "fmt" - "net/url" - "strings" - "time" - - "github.com/ollama/ollama/agent" - "github.com/ollama/ollama/api" - internalcloud "github.com/ollama/ollama/internal/cloud" -) - -const ( - maxWebFetchContentRunes = 60_000 - webSearchTimeout = 15 * time.Second - webFetchTimeout = 30 * time.Second -) - -var ErrWebAuthRequired = errors.New("Not authenticated. Run `ollama signin` and try again.") - -type WebSearch struct{} - -func (w *WebSearch) Name() string { - return "web_search" -} - -func (w *WebSearch) Description() string { - return "Search the web for current information that may not be in the model's training data." -} - -func (w *WebSearch) Schema() api.ToolFunction { - props := api.NewToolPropertiesMap() - props.Set("query", api.ToolProperty{ - Type: api.PropertyType{"string"}, - Description: "The search query to look up on the web.", - }) - return api.ToolFunction{ - Name: w.Name(), - Description: w.Description(), - Parameters: api.ToolFunctionParameters{ - Type: "object", - Properties: props, - Required: []string{"query"}, - }, - } -} - -func (w *WebSearch) RequiresApproval(map[string]any) bool { - return true -} - -func (w *WebSearch) Execute(ctx context.Context, _ agent.ToolContext, args map[string]any) (agent.ToolResult, error) { - // TODO: use shared agent.RequiredStringArg for the "query" parameter (see agent package cleanup plan). - if internalcloud.Disabled() { - return agent.ToolResult{}, errors.New(internalcloud.DisabledError("web search is unavailable")) - } - query, ok := args["query"].(string) - if !ok || strings.TrimSpace(query) == "" { - return agent.ToolResult{}, fmt.Errorf("query parameter is required") - } - - client, err := api.ClientFromEnvironment() - if err != nil { - return agent.ToolResult{}, err - } - - ctx, cancel := context.WithTimeout(ctx, webSearchTimeout) - defer cancel() - - searchResp, err := client.WebSearchExperimental(ctx, &api.WebSearchRequest{Query: query, MaxResults: 5}) - if err != nil { - var authErr api.AuthorizationError - if errors.As(err, &authErr) { - return agent.ToolResult{}, ErrWebAuthRequired - } - return agent.ToolResult{}, err - } - if len(searchResp.Results) == 0 { - return agent.ToolResult{Content: "No results found for query: " + query}, nil - } - - var sb strings.Builder - sb.WriteString(fmt.Sprintf("Search results for: %s\n\n", query)) - for i, result := range searchResp.Results { - sb.WriteString(fmt.Sprintf("%d. %s\n", i+1, result.Title)) - sb.WriteString(fmt.Sprintf(" URL: %s\n", result.URL)) - if result.Content != "" { - content := []rune(result.Content) - if len(content) > 300 { - content = append(content[:300], []rune("...")...) - } - sb.WriteString(fmt.Sprintf(" %s\n", string(content))) - } - sb.WriteByte('\n') - } - return agent.ToolResult{Content: sb.String()}, nil -} - -type WebFetch struct{} - -func (w *WebFetch) Name() string { - return "web_fetch" -} - -func (w *WebFetch) Description() string { - return "Fetch and extract text content from a web page." -} - -func (w *WebFetch) Schema() api.ToolFunction { - props := api.NewToolPropertiesMap() - props.Set("url", api.ToolProperty{ - Type: api.PropertyType{"string"}, - Description: "The URL to fetch and extract content from.", - }) - return api.ToolFunction{ - Name: w.Name(), - Description: w.Description(), - Parameters: api.ToolFunctionParameters{ - Type: "object", - Properties: props, - Required: []string{"url"}, - }, - } -} - -func (w *WebFetch) RequiresApproval(map[string]any) bool { - return true -} - -func (w *WebFetch) Execute(ctx context.Context, _ agent.ToolContext, args map[string]any) (agent.ToolResult, error) { - // TODO: use shared agent.RequiredStringArg for the "url" parameter (see agent package cleanup plan). - if internalcloud.Disabled() { - return agent.ToolResult{}, errors.New(internalcloud.DisabledError("web fetch is unavailable")) - } - urlStr, ok := args["url"].(string) - if !ok || strings.TrimSpace(urlStr) == "" { - return agent.ToolResult{}, fmt.Errorf("url parameter is required") - } - parsed, err := url.Parse(urlStr) - if err != nil { - return agent.ToolResult{}, fmt.Errorf("invalid URL: %w", err) - } - if scheme := strings.ToLower(parsed.Scheme); scheme != "http" && scheme != "https" { - return agent.ToolResult{}, fmt.Errorf("unsupported URL scheme %q: only http and https are allowed", parsed.Scheme) - } - - client, err := api.ClientFromEnvironment() - if err != nil { - return agent.ToolResult{}, err - } - - ctx, cancel := context.WithTimeout(ctx, webFetchTimeout) - defer cancel() - - fetchResp, err := client.WebFetchExperimental(ctx, &api.WebFetchRequest{URL: urlStr}) - if err != nil { - var authErr api.AuthorizationError - if errors.As(err, &authErr) { - return agent.ToolResult{}, ErrWebAuthRequired - } - return agent.ToolResult{}, err - } - - var sb strings.Builder - if fetchResp.Title != "" { - sb.WriteString(fmt.Sprintf("Title: %s\n\n", fetchResp.Title)) - } - if fetchResp.Content != "" { - sb.WriteString("Content:\n") - sb.WriteString(truncateWebFetchContent(fetchResp.Content)) - } else { - sb.WriteString("No content could be extracted from the page.") - } - return agent.ToolResult{Content: sb.String()}, nil -} - -func truncateWebFetchContent(content string) string { - return agent.Truncate(content, agent.TruncateConfig{ - MaxRunes: maxWebFetchContentRunes, - Label: "tool output", - Hint: "Use a narrower request or search query if more detail is needed.", - }) -} diff --git a/agent/tools/web_test.go b/agent/tools/web_test.go deleted file mode 100644 index a9327820fe3..00000000000 --- a/agent/tools/web_test.go +++ /dev/null @@ -1,214 +0,0 @@ -package tools - -import ( - "encoding/json" - "errors" - "net/http" - "net/http/httptest" - "os" - "path/filepath" - "strings" - "testing" - - coreagent "github.com/ollama/ollama/agent" - "github.com/ollama/ollama/api" - "github.com/ollama/ollama/envconfig" - internalcloud "github.com/ollama/ollama/internal/cloud" -) - -func TestWebToolsRequireApproval(t *testing.T) { - if !coreagent.ToolRequiresApproval((&WebSearch{}), map[string]any{"query": "ollama"}) { - t.Fatal("web search should require approval") - } - if !coreagent.ToolRequiresApproval((&WebFetch{}), map[string]any{"url": "https://ollama.com"}) { - t.Fatal("web fetch should require approval") - } -} - -var webToolCases = []struct { - name string - tool coreagent.Tool - args map[string]any - path string - operation string -}{ - {"search", &WebSearch{}, map[string]any{"query": "ollama"}, "/api/experimental/web_search", "web search is unavailable"}, - {"fetch", &WebFetch{}, map[string]any{"url": "https://ollama.com"}, "/api/experimental/web_fetch", "web fetch is unavailable"}, -} - -// enableWebToolsForTest isolates web tool tests from the runner's cloud -// policy. In particular, Windows can inherit both OLLAMA_NO_CLOUD and a -// server.json from USERPROFILE. -func enableWebToolsForTest(t *testing.T) { - t.Helper() - - // Register before t.Setenv so the cache is refreshed after t.Setenv has - // restored the runner's environment during cleanup. - t.Cleanup(envconfig.ReloadServerConfig) - - home := t.TempDir() - t.Setenv("HOME", home) - t.Setenv("USERPROFILE", home) - t.Setenv("OLLAMA_NO_CLOUD", "") - envconfig.ReloadServerConfig() -} - -// runWebTool executes tool against a stub server that responds to every -// request with status and body, returning the resulting error. -func runWebTool(t *testing.T, tool coreagent.Tool, args map[string]any, path string, status int, body string) error { - t.Helper() - ts := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { - if r.URL.Path != path { - t.Fatalf("path = %q, want %q", r.URL.Path, path) - } - w.Header().Set("Content-Type", "application/json") - w.WriteHeader(status) - _, _ = w.Write([]byte(body)) - })) - t.Cleanup(ts.Close) - t.Setenv("OLLAMA_HOST", ts.URL) - _, err := tool.Execute(t.Context(), coreagent.ToolContext{}, args) - return err -} - -func TestWebToolsReportAuthenticationError(t *testing.T) { - enableWebToolsForTest(t) - - for _, tt := range webToolCases { - t.Run(tt.name, func(t *testing.T) { - err := runWebTool(t, tt.tool, tt.args, tt.path, http.StatusUnauthorized, - `{"error":"unauthorized","signin_url":"https://ollama.com/signin"}`) - if !errors.Is(err, ErrWebAuthRequired) { - t.Fatalf("error = %v, want %v", err, ErrWebAuthRequired) - } - }) - } -} - -func TestWebToolsPreserveNonAuthenticationErrors(t *testing.T) { - enableWebToolsForTest(t) - - for _, tt := range webToolCases { - t.Run(tt.name, func(t *testing.T) { - err := runWebTool(t, tt.tool, tt.args, tt.path, http.StatusTooManyRequests, - `{"error":"web search quota exceeded"}`) - if err == nil { - t.Fatal("expected error") - } - if !strings.Contains(err.Error(), "web search quota exceeded") { - t.Fatalf("error = %q, want original error message", err) - } - }) - } -} - -func TestWebToolsIgnoreInheritedCloudPolicy(t *testing.T) { - // This cleanup is registered before the test environment, so it restores - // the server config cache after t.Setenv restores the runner's values. - t.Cleanup(envconfig.ReloadServerConfig) - - home := t.TempDir() - configPath := filepath.Join(home, ".ollama", "server.json") - if err := os.MkdirAll(filepath.Dir(configPath), 0o755); err != nil { - t.Fatal(err) - } - if err := os.WriteFile(configPath, []byte(`{"disable_ollama_cloud":true}`), 0o644); err != nil { - t.Fatal(err) - } - t.Setenv("HOME", home) - t.Setenv("USERPROFILE", home) - t.Setenv("OLLAMA_NO_CLOUD", "1") - envconfig.ReloadServerConfig() - - enableWebToolsForTest(t) - err := runWebTool(t, &WebSearch{}, map[string]any{"query": "ollama"}, "/api/experimental/web_search", http.StatusUnauthorized, - `{"error":"unauthorized","signin_url":"https://ollama.com/signin"}`) - if !errors.Is(err, ErrWebAuthRequired) { - t.Fatalf("error = %v, want %v", err, ErrWebAuthRequired) - } -} - -func TestWebFetchRejectsUnsupportedScheme(t *testing.T) { - enableWebToolsForTest(t) - - tests := []struct { - name string - url string - wantErr bool - }{ - {name: "file scheme", url: "file:///etc/passwd", wantErr: true}, - {name: "data scheme", url: "data:text/plain,secret", wantErr: true}, - {name: "ftp scheme", url: "ftp://example.com/secret", wantErr: true}, - {name: "http allowed", url: "http://example.com", wantErr: false}, - {name: "https allowed", url: "https://example.com", wantErr: false}, - } - for _, tt := range tests { - t.Run(tt.name, func(t *testing.T) { - _, err := (&WebFetch{}).Execute(t.Context(), coreagent.ToolContext{}, map[string]any{"url": tt.url}) - if tt.wantErr && err == nil { - t.Fatal("expected unsupported scheme to be rejected") - } - // For allowed schemes we expect an error only from the missing - // server/auth path, not from scheme validation. The http/https - // cases reach the client and may fail on connection/auth; we only - // assert that the error is NOT a scheme error. - if !tt.wantErr && err != nil && strings.Contains(err.Error(), "unsupported URL scheme") { - t.Fatalf("http/https rejected as unsupported: %v", err) - } - }) - } -} - -func TestWebFetchBoundsContentBeforeReturning(t *testing.T) { - enableWebToolsForTest(t) - - ts := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { - if r.URL.Path != "/api/experimental/web_fetch" { - t.Fatalf("path = %q, want /api/experimental/web_fetch", r.URL.Path) - } - var req api.WebFetchRequest - if err := json.NewDecoder(r.Body).Decode(&req); err != nil { - t.Fatal(err) - } - if req.URL != "https://ollama.com" { - t.Fatalf("request URL = %q, want https://ollama.com", req.URL) - } - if err := json.NewEncoder(w).Encode(api.WebFetchResponse{ - Title: "Ollama", - Content: strings.Repeat("x", maxWebFetchContentRunes+25), - }); err != nil { - t.Fatal(err) - } - })) - defer ts.Close() - t.Setenv("OLLAMA_HOST", ts.URL) - - result, err := (&WebFetch{}).Execute(t.Context(), coreagent.ToolContext{}, map[string]any{ - "url": "https://ollama.com", - }) - if err != nil { - t.Fatal(err) - } - if !strings.Contains(result.Content, "[tool output truncated: showing first ~") || - !strings.Contains(result.Content, "omitted ~7 tokens") || - !strings.Contains(result.Content, "Use a narrower request or search query") { - t.Fatalf("content missing truncation marker: %q", result.Content) - } - if count := strings.Count(result.Content, "x"); count != maxWebFetchContentRunes { - t.Fatalf("captured content count = %d, want %d", count, maxWebFetchContentRunes) - } -} - -func TestWebToolsRejectWhenCloudDisabled(t *testing.T) { - t.Setenv("OLLAMA_NO_CLOUD", "1") - - for _, tt := range webToolCases { - t.Run(tt.name, func(t *testing.T) { - _, err := tt.tool.Execute(t.Context(), coreagent.ToolContext{}, tt.args) - want := internalcloud.DisabledError(tt.operation) - if err == nil || err.Error() != want { - t.Fatalf("error = %v, want %q", err, want) - } - }) - } -} diff --git a/cmd/agent_tui.go b/cmd/agent_tui.go deleted file mode 100644 index 682fea6e20f..00000000000 --- a/cmd/agent_tui.go +++ /dev/null @@ -1,654 +0,0 @@ -package cmd - -import ( - "context" - "errors" - "fmt" - "net/http" - "os" - "runtime" - "slices" - "strconv" - "strings" - "time" - - "github.com/spf13/cobra" - - coreagent "github.com/ollama/ollama/agent" - agenttools "github.com/ollama/ollama/agent/tools" - "github.com/ollama/ollama/api" - "github.com/ollama/ollama/cmd/config" - "github.com/ollama/ollama/cmd/launch" - agentchat "github.com/ollama/ollama/cmd/tui/chat" - "github.com/ollama/ollama/format" - internalcloud "github.com/ollama/ollama/internal/cloud" - "github.com/ollama/ollama/internal/modelref" - "github.com/ollama/ollama/types/model" -) - -type agentTUIOptions struct { - Model string - System string - Format string - Options map[string]any - Think *api.ThinkValue - KeepAlive *api.Duration - ContextWindowTokens int - AllowAllTools bool - ToolsDisabled bool - MultiModal bool -} - -func saveLastAgentModel(model string) error { - model = strings.TrimSpace(model) - if model == "" { - return nil - } - return config.SetLastModel(model) -} - -func prepareAgentModel(cmd *cobra.Command, client *api.Client, opts *agentTUIOptions, thinkExplicit bool) (*api.ShowResponse, error) { - // Unlike `ollama run`, the bare `ollama` root command doesn't define - // --insecure, so GetBool would error; treat it as false. - insecure, _ := cmd.Flags().GetBool("insecure") - info, resolved, err := showOrPullModel(cmd, client, opts.Model, insecure, "run") - if err != nil { - return nil, err - } - // The model may have been resolved to a different name (e.g. its - // ":cloud" variant). - opts.Model = resolved - - ensureCloudStub(cmd.Context(), client, opts.Model) - opts.Think, err = inferThinkingOption(&info.Capabilities, &runOptions{Model: opts.Model, Think: opts.Think}, thinkExplicit) - if err != nil { - return nil, err - } - opts.MultiModal = showResponseSupportsMultimodal(info) - opts.ContextWindowTokens = showResponseContextWindow(info) - return info, nil -} - -func GenerateAgentTUI(cmd *cobra.Command, client *api.Client, opts agentTUIOptions) error { - cwd := agentWorkingDir() - contextWindowForModel := func(ctx context.Context, model string, fallback int) int { - return agentContextWindowForModel(ctx, client, model, fallback) - } - - var skillCatalog *coreagent.SkillCatalog - reloadSkills := func() (*coreagent.SkillCatalog, error) { - catalog, err := coreagent.LoadDefaultSkills(cwd) - if err != nil { - return nil, err - } - if ignored := catalog.ExcludeNames(agentchat.BuiltinSlashCommandNames()); len(ignored) > 0 { - fmt.Fprintf(os.Stderr, "\033[1mwarning:\033[0m ignoring agent skill(s): %s\n", strings.Join(ignored, ", ")) - } - for _, diagnostic := range catalog.Diagnostics() { - fmt.Fprintf(os.Stderr, "\033[1mwarning:\033[0m ignored invalid agent skill: %v\n", diagnostic) - } - skillCatalog = catalog - return catalog, nil - } - if _, err := reloadSkills(); err != nil { - return fmt.Errorf("load agent skills: %w", err) - } - var registry *coreagent.Registry - registryForModel := func(ctx context.Context, model string) *coreagent.Registry { - return agentToolsRegistry(ctx, client, model, skillCatalog) - } - if opts.Model != "" { - registry = agentToolsRegistry(cmd.Context(), client, opts.Model, skillCatalog) - } - systemPrompt := agentSystemPromptWithWorkingDir(opts.Model, opts.System, agentSkillSystemContext(skillCatalog, registry, opts.ToolsDisabled), cwd) - - _, err := agentchat.Run(cmd.Context(), agentchat.Options{ - Model: opts.Model, - Client: client, - Tools: registry, - ToolRegistryForModel: registryForModel, - ToolsDisabled: opts.ToolsDisabled, - MultiModalForModel: func(ctx context.Context, model string) bool { - return agentModelSupportsMultimodal(ctx, client, model) - }, - ModelOptions: func(ctx context.Context) ([]agentchat.ModelOption, error) { - return agentModelOptions(ctx, client) - }, - OnModelSelected: func(_ context.Context, model string) error { - return config.SetLastModel(model) - }, - SystemPromptForModel: func(ctx context.Context, model string, registry *coreagent.Registry, toolsDisabled bool) string { - return agentSystemPromptWithWorkingDir(model, agentSystemFromShow(ctx, client, model), agentSkillSystemContext(skillCatalog, registry, toolsDisabled), cwd) - }, - Skills: skillCatalog, - ImportSkills: coreagent.ImportSkills, - ReloadSkills: reloadSkills, - SystemPrompt: systemPrompt, - WorkingDir: cwd, - Format: opts.Format, - Options: opts.Options, - Think: opts.Think, - KeepAlive: opts.KeepAlive, - MultiModal: opts.MultiModal, - AllowAllTools: opts.AllowAllTools, - ContextWindowTokens: opts.ContextWindowTokens, - Compactor: &coreagent.SimpleCompactor{ - Client: client, - Options: coreagent.CompactionOptions{ContextWindowTokens: opts.ContextWindowTokens}, - }, - ContextWindowTokensForModel: func(ctx context.Context, model string, fallback int) int { - return contextWindowForModel(ctx, model, fallback) - }, - PreloadModel: func(ctx context.Context, model string, think *api.ThinkValue) (int, error) { - return preloadAgentModelIfLocal(ctx, client, opts, model, think) - }, - CheckCloudModel: func(ctx context.Context, model, requiredPlan string) error { - return ensureCloudModelAccess(ctx, client, model, requiredPlan) - }, - OpenBrowser: launch.OpenBrowser, - PollCloudAuth: func(ctx context.Context) (string, bool, error) { - user, err := client.Whoami(ctx) - if err != nil { - return "", false, err - } - if user == nil || user.Name == "" { - return "", false, nil - } - return user.Name, true, nil - }, - }) - return err -} - -func agentSkillSystemContext(catalog *coreagent.SkillCatalog, registry *coreagent.Registry, toolsDisabled bool) string { - if toolsDisabled || registry == nil { - return "" - } - if _, ok := registry.Get("skill"); !ok { - return "" - } - return catalog.SystemContext() -} - -func agentSelectionItems(models []agentchat.ModelOption) []launch.SelectionItem { - items := make([]launch.SelectionItem, 0, len(models)) - for _, model := range models { - items = append(items, launch.SelectionItem{ - Name: model.Name, - Description: strings.TrimSpace(model.Description), - Recommended: model.Recommended, - AvailabilityBadge: model.AvailabilityBadge, - }) - } - return items -} - -var agentGetwd = os.Getwd - -func agentWorkingDir() string { - cwd, err := agentGetwd() - if err != nil { - return "" - } - return cwd -} - -func agentSystemPromptWithWorkingDir(modelName string, modelSystem string, extra string, workingDir string) string { - return agentSystemPromptAtWithWorkingDir(time.Now(), modelName, modelSystem, extra, workingDir) -} - -func agentSystemPromptAtWithWorkingDir(now time.Time, modelName string, modelSystem string, extra string, workingDir string) string { - var parts []string - parts = append(parts, agentDefaultSystemPromptWithWorkingDir(now, modelName, workingDir)) - if strings.TrimSpace(modelSystem) != "" { - parts = append(parts, strings.TrimSpace(modelSystem)) - } - if strings.TrimSpace(extra) != "" { - parts = append(parts, strings.TrimSpace(extra)) - } - return strings.Join(parts, "\n\n") -} - -func agentDefaultSystemPromptWithWorkingDir(now time.Time, modelName string, workingDir string) string { - date := now.Format("Monday, January 2, 2006") - shellName := "bash" - if runtime.GOOS == "windows" { - shellName = "PowerShell" - } - parts := []string{ - "You are running in Ollama, in a harness to help the user accomplish tasks, and the model is " + modelName + ".", - "", - "Current date: " + date + ".", - "", - } - parts = append(parts, - "Be concise, practical, and action-oriented. Use tools when they materially help. Verify current or fast-changing facts with web tools when available; otherwise state uncertainty.", - "", - "Use "+shellName+" carefully. Prefer read-only inspection first. Stay within the current working directory unless explicitly asked. Surface intent before risky actions such as writes, deletes, moves, installs, git state changes, service changes, sudo, secrets access, network scripts, or commands outside the working directory. Request approval when required and do not work around denied approvals.", - "", - "Tell the user about meaningful changes, verification, failures, blockers, assumptions, and risks. Summarize routine tool output instead of dumping it.", - ) - if workingDir != "" { - parts = append(parts, "Current working directory: "+strconv.Quote(workingDir)+".") - } - return strings.Join(parts, "\n") -} - -func agentSystemFromShow(ctx context.Context, client *api.Client, modelName string) string { - if client == nil || strings.TrimSpace(modelName) == "" { - return "" - } - resp, err := client.Show(ctx, &api.ShowRequest{Model: modelName}) - if err != nil { - fmt.Fprintf(os.Stderr, "\033[1mwarning:\033[0m could not load model system prompt: %v\n", err) - return "" - } - return resp.System -} - -func agentToolsRegistry(ctx context.Context, client *api.Client, modelName string, skillCatalog *coreagent.SkillCatalog) *coreagent.Registry { - supportsTools, err := agentModelSupportsTools(ctx, client, modelName) - if err != nil { - fmt.Fprintf(os.Stderr, "\033[1mwarning:\033[0m could not check model capabilities: %v\n", err) - } - if !supportsTools { - return nil - } - - registry := &coreagent.Registry{} - if os.Getenv("OLLAMA_AGENT_DISABLE_SHELL") == "" { - registry.Register(&agenttools.Bash{}) - } - registry.Register(&agenttools.Read{}) - registry.Register(&agenttools.Edit{}) - if len(skillCatalog.List()) > 0 { - registry.Register(&agenttools.Skill{Catalog: skillCatalog}) - } - - if os.Getenv("OLLAMA_AGENT_DISABLE_WEBSEARCH") == "" { - if disabled, known := agentCloudStatusDisabled(ctx, client); !known || !disabled { - registry.Register(&agenttools.WebSearch{}) - registry.Register(&agenttools.WebFetch{}) - } else { - fmt.Fprintf(os.Stderr, "%s\n", internalcloud.DisabledError("web search is unavailable")) - } - } - return registry -} - -func agentModelSupportsTools(ctx context.Context, client *api.Client, modelName string) (bool, error) { - if client == nil || strings.TrimSpace(modelName) == "" { - return false, nil - } - resp, err := client.Show(ctx, &api.ShowRequest{Model: modelName}) - if err != nil { - return false, err - } - return slices.Contains(resp.Capabilities, model.CapabilityTools), nil -} - -func agentModelSupportsMultimodal(ctx context.Context, client *api.Client, modelName string) bool { - if client == nil || strings.TrimSpace(modelName) == "" { - return false - } - resp, err := client.Show(ctx, &api.ShowRequest{Model: modelName}) - if err != nil { - fmt.Fprintf(os.Stderr, "\033[1mwarning:\033[0m could not check model capabilities: %v\n", err) - return false - } - return showResponseSupportsMultimodal(resp) -} - -func showResponseSupportsMultimodal(resp *api.ShowResponse) bool { - if resp == nil { - return false - } - if slices.Contains(resp.Capabilities, model.CapabilityVision) || slices.Contains(resp.Capabilities, model.CapabilityAudio) { - return true - } - if len(resp.ProjectorInfo) != 0 { - return true - } - for key := range resp.ModelInfo { - if strings.Contains(key, ".vision.") { - return true - } - } - return false -} - -func agentContextWindowForModel(ctx context.Context, client *api.Client, modelName string, fallback int) int { - if client == nil || strings.TrimSpace(modelName) == "" { - return fallback - } - if tokens := launch.LoadedContextWindow(ctx, client, modelName); tokens > 0 { - return tokens - } - if modelref.HasExplicitCloudSource(modelName) { - if tokens := agentRecommendationContextWindowForModel(ctx, client, modelName); tokens > 0 { - return tokens - } - } - resp, err := client.Show(ctx, &api.ShowRequest{Model: modelName}) - if err != nil { - return fallback - } - if tokens := showResponseContextWindow(resp); tokens > 0 { - return tokens - } - return fallback -} - -func agentRecommendationContextWindowForModel(ctx context.Context, client *api.Client, modelName string) int { - if client == nil { - return 0 - } - recs, err := client.ModelRecommendationsExperimental(ctx) - if err != nil || recs == nil { - return 0 - } - return contextWindowFromRecommendations(modelName, recs.Recommendations) -} - -func contextWindowFromRecommendations(modelName string, recommendations []api.ModelRecommendation) int { - for _, rec := range recommendations { - if rec.ContextLength <= 0 { - continue - } - if launch.SameModelRef(modelName, rec.Model) { - return rec.ContextLength - } - } - return 0 -} - -func showResponseContextWindow(resp *api.ShowResponse) int { - if resp == nil { - return 0 - } - if resp.Details.ContextLength > 0 { - return resp.Details.ContextLength - } - if n, ok := numericModelInfo(resp.ModelInfo["general.context_length"]); ok { - return n - } - best := 0 - for key, value := range resp.ModelInfo { - if key != "context_length" && !strings.HasSuffix(key, ".context_length") { - continue - } - if n, ok := numericModelInfo(value); ok && n > best { - best = n - } - } - return best -} - -func numericModelInfo(value any) (int, bool) { - switch v := value.(type) { - case int: - return v, v > 0 - case int32: - return int(v), v > 0 - case int64: - return int(v), v > 0 - case uint: - return int(v), v > 0 - case uint32: - return int(v), v > 0 - case uint64: - return int(v), v > 0 - case float64: - return int(v), v > 0 - case string: - n, err := strconv.Atoi(strings.TrimSpace(v)) - return n, err == nil && n > 0 - default: - return 0, false - } -} - -func preloadAgentModelIfLocal(ctx context.Context, client *api.Client, opts agentTUIOptions, modelName string, think *api.ThinkValue) (int, error) { - modelName = strings.TrimSpace(modelName) - if client == nil || modelName == "" { - return 0, nil - } - if modelref.HasExplicitCloudSource(modelName) { - return 0, nil - } - info, err := client.Show(ctx, &api.ShowRequest{Model: modelName}) - if err != nil { - return 0, err - } - if info.RemoteHost != "" { - return 0, nil - } - if err := client.Generate(ctx, &api.GenerateRequest{ - Model: modelName, - KeepAlive: opts.KeepAlive, - Options: opts.Options, - Think: think, - }, func(api.GenerateResponse) error { - return nil - }); err != nil { - return 0, err - } - return launch.LoadedContextWindow(ctx, client, modelName), nil -} - -func agentModelOptions(ctx context.Context, client *api.Client) ([]agentchat.ModelOption, error) { - if client == nil { - return nil, errors.New("model picker requires an API client") - } - - list, err := client.List(ctx) - if err != nil { - return nil, err - } - - seen := make(map[string]struct{}) - var options []agentchat.ModelOption - add := func(name, description string, recommended bool, requiredPlan string, cloud bool) { - name = strings.TrimSpace(name) - if name == "" { - return - } - key := strings.ToLower(name) - if _, ok := seen[key]; ok { - return - } - seen[key] = struct{}{} - options = append(options, agentchat.ModelOption{ - Name: name, - Description: strings.TrimSpace(description), - Recommended: recommended, - RequiredPlan: requiredPlan, - Cloud: cloud, - }) - } - - if disabled, known := agentCloudStatusDisabled(ctx, client); !known || !disabled { - if recs, err := client.ModelRecommendationsExperimental(ctx); err == nil { - for _, rec := range recs.Recommendations { - name := strings.TrimSpace(rec.Model) - if !modelref.HasExplicitCloudSource(name) { - continue - } - add(name, agentRecommendationDescription(rec), true, strings.TrimSpace(rec.RequiredPlan), true) - } - } - } - - local := slices.Clone(list.Models) - slices.SortStableFunc(local, func(a, b api.ListModelResponse) int { - return strings.Compare(strings.ToLower(a.Name), strings.ToLower(b.Name)) - }) - for _, model := range local { - name := strings.TrimSpace(model.Name) - if name == "" { - name = strings.TrimSpace(model.Model) - } - name = strings.TrimSuffix(name, ":latest") - if modelref.HasExplicitCloudSource(name) { - add(name, agentCloudModelDescription(model), false, "", true) - continue - } - add(name, agentLocalModelDescription(model), false, "", false) - } - - badges, signInURLs := cloudAvailabilityBadges(ctx, client, options) - for i := range options { - options[i].AvailabilityBadge = badges[options[i].Name] - options[i].SignInURL = signInURLs[options[i].Name] - } - return options, nil -} - -func cloudAvailabilityBadges(ctx context.Context, client *api.Client, options []agentchat.ModelOption) (map[string]string, map[string]string) { - badges := make(map[string]string) - signInURLs := make(map[string]string) - hasCloud := false - for _, opt := range options { - if opt.Cloud { - hasCloud = true - break - } - } - if !hasCloud { - return badges, signInURLs - } - - if disabled, known := agentCloudStatusDisabled(ctx, client); known && disabled { - return badges, signInURLs - } - - whoamiCtx, cancel := context.WithTimeout(ctx, 3*time.Second) - defer cancel() - user, err := client.Whoami(whoamiCtx) - if err != nil { - var authErr api.AuthorizationError - signInURL := "" - if errors.As(err, &authErr) && (authErr.StatusCode == http.StatusUnauthorized || authErr.SigninURL != "") { - if authErr.SigninURL != "" { - signInURL = authErr.SigninURL - } - } else { - return badges, signInURLs - } - for _, opt := range options { - if opt.Cloud { - badges[opt.Name] = "Sign in required" - if signInURL != "" { - signInURLs[opt.Name] = signInURL - } - } - } - return badges, signInURLs - } - - signedIn := user != nil && user.Name != "" - for _, opt := range options { - if !opt.Cloud { - continue - } - if !signedIn { - badges[opt.Name] = "Sign in required" - } else if opt.RequiredPlan != "" && !launch.PlanSatisfies(user.Plan, opt.RequiredPlan) { - badges[opt.Name] = "Upgrade required" - } - } - return badges, signInURLs -} - -func agentRecommendationDescription(rec api.ModelRecommendation) string { - var parts []string - if description := strings.TrimSpace(rec.Description); description != "" { - parts = append(parts, description) - } else { - parts = append(parts, "cloud") - } - if rec.ContextLength > 0 { - parts = append(parts, format.HumanNumber(uint64(rec.ContextLength))+" ctx") - } - return strings.Join(parts, " - ") -} - -func agentLocalModelDescription(model api.ListModelResponse) string { - desc := agentModelArchDescription(model) - if desc == "" { - return "local" - } - return "local - " + desc -} - -func agentCloudModelDescription(model api.ListModelResponse) string { - return agentModelArchDescription(model) -} - -func agentModelArchDescription(model api.ListModelResponse) string { - var details []string - if model.Details.Family != "" { - details = append(details, model.Details.Family) - } - if ps := humanizedParameterSize(model.Details.ParameterSize); ps != "" { - details = append(details, ps) - } - if model.Details.QuantizationLevel != "" { - details = append(details, model.Details.QuantizationLevel) - } - var parts []string - if len(details) > 0 { - parts = append(parts, strings.Join(details, " ")) - } - if model.Details.ContextLength > 0 { - parts = append(parts, format.HumanNumber(uint64(model.Details.ContextLength))+" ctx") - } - return strings.Join(parts, " - ") -} - -func humanizedParameterSize(s string) string { - s = strings.TrimSpace(s) - if s == "" { - return "" - } - if f, err := strconv.ParseFloat(s, 64); err == nil { - return format.HumanNumber(uint64(f)) - } - return s -} - -func agentCloudStatusDisabled(ctx context.Context, client *api.Client) (disabled bool, known bool) { - if internalcloud.Disabled() { - return true, true - } - - status, err := client.CloudStatusExperimental(ctx) - if err != nil { - var statusErr api.StatusError - if errors.As(err, &statusErr) && statusErr.StatusCode == http.StatusNotFound { - return false, false - } - return false, false - } - return status.Cloud.Disabled, true -} - -func ensureCloudModelAccess(ctx context.Context, client *api.Client, modelName, requiredPlan string) error { - if client == nil { - return errors.New("no API client available") - } - if disabled, known := agentCloudStatusDisabled(ctx, client); known && disabled { - return errors.New("remote inference is unavailable") - } - user, err := client.Whoami(ctx) - if err != nil { - return err - } - if user != nil && user.Name != "" { - if requiredPlan != "" && !launch.PlanSatisfies(user.Plan, requiredPlan) { - return fmt.Errorf("plan upgrade required: %s needs plan %s, you have %s", modelName, requiredPlan, user.Plan) - } - return nil - } - return fmt.Errorf("%s requires sign in", modelName) -} diff --git a/cmd/agent_tui_test.go b/cmd/agent_tui_test.go deleted file mode 100644 index 299440b180d..00000000000 --- a/cmd/agent_tui_test.go +++ /dev/null @@ -1,189 +0,0 @@ -package cmd - -import ( - "errors" - "os" - "path/filepath" - "strconv" - "strings" - "testing" - "time" - - coreagent "github.com/ollama/ollama/agent" - agenttools "github.com/ollama/ollama/agent/tools" - "github.com/ollama/ollama/api" - "github.com/ollama/ollama/cmd/config" - agentchat "github.com/ollama/ollama/cmd/tui/chat" -) - -func TestAgentSystemPromptIncludesSessionWorkingDirOnce(t *testing.T) { - workingDir := t.TempDir() - prompt := agentSystemPromptAtWithWorkingDir( - time.Date(2026, time.July, 14, 0, 0, 0, 0, time.UTC), - "test-model", - "model instruction", - "caller instruction", - workingDir, - ) - - workingDirInstruction := "Current working directory: " + strconv.Quote(workingDir) + "." - if got := strings.Count(prompt, workingDirInstruction); got != 1 { - t.Fatalf("working directory instruction count = %d, want 1:\n%s", got, prompt) - } - for _, want := range []string{"model instruction", "caller instruction"} { - if !strings.Contains(prompt, want) { - t.Fatalf("prompt missing %q:\n%s", want, prompt) - } - } -} - -func TestAgentWorkingDirIgnoresGetwdFailure(t *testing.T) { - original := agentGetwd - agentGetwd = func() (string, error) { - return "", errors.New("getwd failed") - } - t.Cleanup(func() { - agentGetwd = original - }) - - if got := agentWorkingDir(); got != "" { - t.Fatalf("working directory = %q, want empty on getwd failure", got) - } -} - -func TestAgentSystemPromptIncludesSkillCatalog(t *testing.T) { - dir := t.TempDir() - if err := os.Mkdir(filepath.Join(dir, "release-notes"), 0o755); err != nil { - t.Fatal(err) - } - if err := os.WriteFile(filepath.Join(dir, "release-notes", "SKILL.md"), []byte("---\nname: release-notes\ndescription: Draft releases.\n---\nUse bullets."), 0o644); err != nil { - t.Fatal(err) - } - catalog, err := coreagent.DiscoverSkills(dir) - if err != nil { - t.Fatal(err) - } - got := agentSystemPromptAtWithWorkingDir(time.Date(2026, 7, 14, 0, 0, 0, 0, time.UTC), "model", "", catalog.SystemContext(), "") - if !strings.Contains(got, "release-notes: Draft releases.") || !strings.Contains(got, "normal approval rules") { - t.Fatalf("system prompt missing skill context: %q", got) - } -} - -func TestAgentSkillSystemContextRequiresAvailableEnabledSkillTool(t *testing.T) { - dir := t.TempDir() - if err := os.Mkdir(filepath.Join(dir, "release-notes"), 0o755); err != nil { - t.Fatal(err) - } - if err := os.WriteFile(filepath.Join(dir, "release-notes", "SKILL.md"), []byte("---\nname: release-notes\ndescription: Draft releases.\n---\nUse bullets."), 0o644); err != nil { - t.Fatal(err) - } - catalog, err := coreagent.DiscoverSkills(dir) - if err != nil { - t.Fatal(err) - } - registry := &coreagent.Registry{} - registry.Register(&agenttools.Skill{Catalog: catalog}) - - if got := agentSkillSystemContext(catalog, registry, false); !strings.Contains(got, "release-notes: Draft releases.") { - t.Fatalf("enabled skill context = %q", got) - } - if got := agentSkillSystemContext(catalog, registry, true); got != "" { - t.Fatalf("disabled tools should omit skill context, got %q", got) - } - if got := agentSkillSystemContext(catalog, &coreagent.Registry{}, false); got != "" { - t.Fatalf("unavailable skill tool should omit skill context, got %q", got) - } -} - -func TestAgentSkillCommandCollisionsAreIgnored(t *testing.T) { - dir := t.TempDir() - for _, name := range []string{"release-notes", "system", "exit"} { - if err := os.Mkdir(filepath.Join(dir, name), 0o755); err != nil { - t.Fatal(err) - } - content := "---\nname: " + name + "\ndescription: Test skill.\n---\nInstructions." - if err := os.WriteFile(filepath.Join(dir, name, "SKILL.md"), []byte(content), 0o644); err != nil { - t.Fatal(err) - } - } - catalog, err := coreagent.DiscoverSkills(dir) - if err != nil { - t.Fatal(err) - } - - ignored := catalog.ExcludeNames(agentchat.BuiltinSlashCommandNames()) - if got, want := strings.Join(ignored, ", "), "exit, system"; got != want { - t.Fatalf("ignored skills = %q, want %q", got, want) - } - if _, err := catalog.Load("release-notes"); err != nil { - t.Fatalf("non-conflicting skill should remain available: %v", err) - } - if context := catalog.SystemContext(); !strings.Contains(context, "release-notes: Test skill.") || strings.Contains(context, "system: Test skill.") || strings.Contains(context, "exit: Test skill.") { - t.Fatalf("skill context = %q", context) - } - for _, name := range []string{"system", "exit"} { - if _, err := catalog.Load(name); err == nil { - t.Fatalf("conflicting skill %q should be ignored", name) - } - } -} - -func TestAgentSelectionItemsUseLaunchSections(t *testing.T) { - items := agentSelectionItems([]agentchat.ModelOption{ - {Name: "glm-5.2:cloud", Description: "cloud", Recommended: true, Cloud: true}, - {Name: "llama3.2", Description: "local"}, - }) - - if len(items) != 2 { - t.Fatalf("items = %d, want 2", len(items)) - } - if !items[0].Recommended { - t.Fatalf("cloud recommendation should be pinned: %#v", items[0]) - } - if items[1].Recommended { - t.Fatalf("local selected model should stay in launch More section: %#v", items[1]) - } - if items[1].Description != "local" { - t.Fatalf("selected model description = %q, want plain description", items[1].Description) - } -} - -func TestContextWindowFromRecommendationsMatchesCloudModel(t *testing.T) { - got := contextWindowFromRecommendations("glm-5.2:cloud", []api.ModelRecommendation{ - {Model: "gemma4:cloud", ContextLength: 32768}, - {Model: "glm-5.2:cloud", ContextLength: 1048576}, - }) - if got != 1048576 { - t.Fatalf("context window = %d, want 1048576", got) - } -} - -func TestShowResponseContextWindowReadsArchitectureContextLength(t *testing.T) { - got := showResponseContextWindow(&api.ShowResponse{ - ModelInfo: map[string]any{ - "qwen3.context_length": uint32(262144), - "qwen3.rope.scaling.original_context_length": uint32(32768), - }, - }) - if got != 262144 { - t.Fatalf("context window = %d, want 262144", got) - } -} - -func TestSaveLastAgentModel(t *testing.T) { - setCmdTestHome(t, t.TempDir()) - - if err := saveLastAgentModel(" qwen3:8b "); err != nil { - t.Fatalf("saveLastAgentModel returned error: %v", err) - } - if got := config.LastModel(); got != "qwen3:8b" { - t.Fatalf("last model = %q, want qwen3:8b", got) - } - - if err := saveLastAgentModel(" "); err != nil { - t.Fatalf("saveLastAgentModel blank returned error: %v", err) - } - if got := config.LastModel(); got != "qwen3:8b" { - t.Fatalf("blank save changed last model to %q", got) - } -} diff --git a/cmd/cmd.go b/cmd/cmd.go index 5b3423593bc..212ee3227bb 100644 --- a/cmd/cmd.go +++ b/cmd/cmd.go @@ -2128,32 +2128,52 @@ Environment Variables: } func launchInteractiveModel(cmd *cobra.Command, modelName string) error { + opts := runOptions{ + Model: modelName, + WordWrap: os.Getenv("TERM") == "xterm-256color", + Options: map[string]any{}, + ShowConnect: true, + } + client, err := api.ClientFromEnvironment() if err != nil { return err } - opts := agentTUIOptions{ - Model: modelName, - Options: map[string]any{}, - } - info, err := prepareAgentModel(cmd, client, &opts, false) + info, resolvedModel, err := showOrPullModel(cmd, client, modelName, false, "run") if err != nil { if handleCloudAuthorizationError(err) { return nil } return err } - opts.System = info.System + opts.Model = resolvedModel + ensureCloudStub(cmd.Context(), client, opts.Model) - if err := saveLastAgentModel(opts.Model); err != nil { + opts.Think, err = inferThinkingOption(&info.Capabilities, &opts, false) + if err != nil { return err } - if err := GenerateAgentTUI(cmd, client, opts); err != nil { - if handleCloudAuthorizationError(err) { - return nil + + audioCapable := slices.Contains(info.Capabilities, model.CapabilityAudio) + opts.MultiModal = slices.Contains(info.Capabilities, model.CapabilityVision) || audioCapable + if len(info.ProjectorInfo) != 0 { + opts.MultiModal = true + } + for key := range info.ModelInfo { + if strings.Contains(key, ".vision.") { + opts.MultiModal = true + break } - return fmt.Errorf("error running agent: %w", err) + } + + applyShowResponseToRunOptions(&opts, info) + + if err := loadOrUnloadModel(cmd, &opts); err != nil { + return fmt.Errorf("error loading model: %w", err) + } + if err := generateInteractive(cmd, opts); err != nil { + return fmt.Errorf("error running model: %w", err) } return nil } diff --git a/cmd/cmd_test.go b/cmd/cmd_test.go index b2970d9768c..66e5b840ade 100644 --- a/cmd/cmd_test.go +++ b/cmd/cmd_test.go @@ -2390,3 +2390,44 @@ func TestIsLocalhost(t *testing.T) { }) } } + +func TestRunCommandHasNoAgentFlags(t *testing.T) { + root := NewCLI() + run, _, err := root.Find([]string{"run"}) + if err != nil { + t.Fatal(err) + } + + for _, name := range []string{"resume", "headless", "auto-approve-tools", "skill", "experimental", "experimental-yolo", "experimental-websearch"} { + if flag := run.Flags().Lookup(name); flag != nil { + t.Errorf("run command still exposes former agent flag --%s", name) + } + } +} + +func TestFormerAgentEntryPointsAreRejected(t *testing.T) { + tests := [][]string{ + {"run", "llama3", "--resume"}, + {"run", "llama3", "--headless"}, + {"run", "llama3", "--auto-approve-tools"}, + {"run", "llama3", "--skill", "release-notes"}, + {"run", "llama3", "--experimental"}, + {"run", "llama3", "--experimental-yolo"}, + {"run", "llama3", "--experimental-websearch"}, + {"agent"}, + } + + for _, args := range tests { + t.Run(strings.Join(args, " "), func(t *testing.T) { + root := NewCLI() + root.SetArgs(args) + err := root.Execute() + if err == nil { + t.Fatalf("former agent entry point %q succeeded", args) + } + if !strings.Contains(err.Error(), "unknown") { + t.Fatalf("former agent entry point %q returned %v, want unknown command or flag", args, err) + } + }) + } +} diff --git a/cmd/internal/filedata/filedata.go b/cmd/internal/filedata/filedata.go deleted file mode 100644 index 257fb609ada..00000000000 --- a/cmd/internal/filedata/filedata.go +++ /dev/null @@ -1,189 +0,0 @@ -package filedata - -import ( - "errors" - "fmt" - "io" - "net/http" - "net/url" - "os" - "path/filepath" - "regexp" - "slices" - "strings" - - "github.com/ollama/ollama/api" -) - -type File struct { - Path string - Data api.ImageData -} - -func NormalizePath(fp string) string { - fp = strings.Trim(fp, "\"") - fp = strings.NewReplacer( - "\\ ", " ", - "\\(", "(", - "\\)", ")", - "\\[", "[", - "\\]", "]", - "\\{", "{", - "\\}", "}", - "\\$", "$", - "\\&", "&", - "\\;", ";", - "\\'", "'", - "\\\\", "\\", - "\\*", "*", - "\\?", "?", - "\\~", "~", - ).Replace(fp) - - if u, err := url.Parse(fp); err == nil && strings.EqualFold(u.Scheme, "file") { - return normalizeFileURL(u) - } else if normalized, ok := normalizeMalformedFileURL(fp); ok { - return normalized - } - - return fp -} - -// fileExtractRe matches file:// URLs and filesystem paths ending in image/audio -// extensions. Hoisted to package scope so the per-keystroke slash-completion -// path (chat.slashInputIsMultimodalFile -> ExtractNames) doesn't recompile it -// on every call. -var fileExtractRe = regexp.MustCompile(`(?:file://\S+?\.(?i:jpg|jpeg|png|webp|wav)\b)|(?:(?:[a-zA-Z]:)?(?:\./|\.\\|/|\\)[\S\\ ]+?\.(?i:jpg|jpeg|png|webp|wav)\b)`) - -func ExtractNames(input string) []string { - return fileExtractRe.FindAllString(input, -1) -} - -func Extract(input string) (string, []api.ImageData, error) { - cleaned, files, err := ExtractWithFiles(input) - if err != nil { - return "", nil, err - } - data := make([]api.ImageData, 0, len(files)) - for _, file := range files { - data = append(data, file.Data) - } - return cleaned, data, nil -} - -func ExtractWithFiles(input string) (string, []File, error) { - filePaths := ExtractNames(input) - var files []File - - for _, fp := range filePaths { - nfp := NormalizePath(fp) - data, err := GetData(nfp) - if errors.Is(err, os.ErrNotExist) { - continue - } else if err != nil { - return "", nil, fmt.Errorf("couldn't process file %q: %w", nfp, err) - } - input = strings.ReplaceAll(input, "'"+nfp+"'", "") - input = strings.ReplaceAll(input, "'"+fp+"'", "") - input = strings.ReplaceAll(input, `"`+nfp+`"`, "") - input = strings.ReplaceAll(input, `"`+fp+`"`, "") - input = strings.ReplaceAll(input, fp, "") - files = append(files, File{Path: nfp, Data: data}) - } - return strings.TrimSpace(input), files, nil -} - -func GetData(filePath string) ([]byte, error) { - file, err := os.Open(filePath) - if err != nil { - return nil, err - } - defer file.Close() - - buf := make([]byte, 512) - _, err = file.Read(buf) - if err != nil { - return nil, err - } - - contentType := http.DetectContentType(buf) - allowedTypes := []string{"image/jpeg", "image/jpg", "image/png", "image/webp", "audio/wave"} - if !slices.Contains(allowedTypes, contentType) { - return nil, fmt.Errorf("invalid file type: %s", contentType) - } - - info, err := file.Stat() - if err != nil { - return nil, err - } - - var maxSize int64 = 100 * 1024 * 1024 - if info.Size() > maxSize { - return nil, errors.New("file size exceeds maximum limit (100MB)") - } - - buf = make([]byte, info.Size()) - _, err = file.Seek(0, 0) - if err != nil { - return nil, err - } - - _, err = io.ReadFull(file, buf) - if err != nil { - return nil, err - } - - return buf, nil -} - -func Kind(path string) string { - if strings.EqualFold(filepath.Ext(path), ".wav") { - return "audio" - } - return "image" -} - -func normalizeFileURL(u *url.URL) string { - path := u.Path - if unescaped, err := url.PathUnescape(path); err == nil { - path = unescaped - } - host := u.Host - if unescaped, err := url.PathUnescape(host); err == nil { - host = unescaped - } - if len(host) >= 2 && host[1] == ':' && isASCIIAlpha(host[0]) { - return filepath.Clean(filepath.FromSlash(host + path)) - } - if len(path) >= 4 && path[0] == '/' && path[2] == ':' && isASCIIAlpha(path[1]) { - path = path[1:] - } - if u.Host != "" && !strings.EqualFold(u.Host, "localhost") { - return `\\` + u.Host + filepath.FromSlash(path) - } - return filepath.FromSlash(path) -} - -func normalizeMalformedFileURL(raw string) (string, bool) { - const prefix = "file://" - if !strings.HasPrefix(strings.ToLower(raw), prefix) { - return "", false - } - - path := raw[len(prefix):] - if unescaped, err := url.PathUnescape(path); err == nil { - path = unescaped - } - path = strings.TrimPrefix(path, "localhost") - if len(path) >= 3 && path[0] == '/' && path[2] == ':' && isASCIIAlpha(path[1]) { - path = path[1:] - } - if len(path) >= 2 && path[1] == ':' && isASCIIAlpha(path[0]) { - return filepath.Clean(filepath.FromSlash(path)), true - } - return "", false -} - -func isASCIIAlpha(b byte) bool { - return (b >= 'a' && b <= 'z') || (b >= 'A' && b <= 'Z') -} diff --git a/cmd/internal/filedata/filedata_test.go b/cmd/internal/filedata/filedata_test.go deleted file mode 100644 index 0ce801a1b43..00000000000 --- a/cmd/internal/filedata/filedata_test.go +++ /dev/null @@ -1,223 +0,0 @@ -package filedata - -import ( - "net/url" - "os" - "path/filepath" - "strings" - "testing" -) - -func TestNormalizePathMalformedWindowsFileURL(t *testing.T) { - got := NormalizePath(`file://C:%5CUsers%5Cjdoe%5CPictures%5Cimg.png`) - want := filepath.Clean(`C:\Users\jdoe\Pictures\img.png`) - if got != want { - t.Fatalf("path = %q, want %q", got, want) - } -} - -func TestNormalizePathTwoSlashWindowsFileURL(t *testing.T) { - got := NormalizePath(`file://C:/Users/jdoe/Pictures/img.png`) - want := filepath.Clean(`C:/Users/jdoe/Pictures/img.png`) - if got != want { - t.Fatalf("path = %q, want %q", got, want) - } -} - -func TestNormalizePathLocalhostWindowsFileURL(t *testing.T) { - got := NormalizePath(`file://localhost/C:/Users/jdoe/Pictures/img.png`) - want := filepath.Clean(`C:/Users/jdoe/Pictures/img.png`) - if got != want { - t.Fatalf("path = %q, want %q", got, want) - } -} - -func TestExtractNames(t *testing.T) { - // Unix style paths - input := ` some preamble - ./relative\ path/one.png inbetween1 ./not a valid two.jpg inbetween2 ./1.svg -/unescaped space /three.jpeg inbetween3 /valid\ path/dir/four.png "./quoted with spaces/five.JPG -/unescaped space /six.webp inbetween6 /valid\ path/dir/seven.WEBP` - res := ExtractNames(input) - if len(res) != 7 { - t.Fatalf("len = %d, want 7", len(res)) - } - assertContains(t, res[0], "one.png") - assertContains(t, res[1], "two.jpg") - assertContains(t, res[2], "three.jpeg") - assertContains(t, res[3], "four.png") - assertContains(t, res[4], "five.JPG") - assertContains(t, res[5], "six.webp") - assertContains(t, res[6], "seven.WEBP") - assertNotContains(t, res[4], "\"") - for _, r := range res { - assertNotContains(t, r, "inbetween1") - } - assertNotContainsSlice(t, res, "./1.svg") -} - -func TestExtractNamesWindowsPaths(t *testing.T) { - input := ` some preamble - c:/users/jdoe/one.png inbetween1 c:/program files/someplace/two.jpg inbetween2 - /absolute/nospace/three.jpeg inbetween3 /absolute/with space/four.png inbetween4 -./relative\ path/five.JPG inbetween5 "./relative with/spaces/six.png inbetween6 -d:\path with\spaces\seven.JPEG inbetween7 c:\users\jdoe\eight.png inbetween8 - d:\program files\someplace\nine.png inbetween9 "E:\program files\someplace\ten.PNG -c:/users/jdoe/eleven.webp inbetween11 c:/program files/someplace/twelve.WebP inbetween12 -d:\path with\spaces\thirteen.WEBP some ending -` - res := ExtractNames(input) - if len(res) != 13 { - t.Fatalf("len = %d, want 13", len(res)) - } - assertNotContainsSlice(t, res, "inbetween2") - assertContains(t, res[0], "one.png") - assertContains(t, res[0], "c:") - assertContains(t, res[1], "two.jpg") - assertContains(t, res[1], "c:") - assertContains(t, res[2], "three.jpeg") - assertContains(t, res[3], "four.png") - assertContains(t, res[4], "five.JPG") - assertContains(t, res[5], "six.png") - assertContains(t, res[6], "seven.JPEG") - assertContains(t, res[6], "d:") - assertContains(t, res[7], "eight.png") - assertContains(t, res[7], "c:") - assertContains(t, res[8], "nine.png") - assertContains(t, res[8], "d:") - assertContains(t, res[9], "ten.PNG") - assertContains(t, res[9], "E:") - assertContains(t, res[10], "eleven.webp") - assertContains(t, res[10], "c:") - assertContains(t, res[11], "twelve.WebP") - assertContains(t, res[11], "c:") - assertContains(t, res[12], "thirteen.WEBP") - assertContains(t, res[12], "d:") -} - -func TestExtractNamesDragDropPaths(t *testing.T) { - input := `file:///Users/jdoe/Pictures/one.png file://localhost/C:/Users/jdoe/Pictures/two.webp file:///C:/Users/jdoe/Pictures/three.jpg .\relative\four.png` - res := ExtractNames(input) - if len(res) != 4 { - t.Fatalf("len = %d, want 4", len(res)) - } - assertContains(t, res[0], "file:///Users/jdoe/Pictures/one.png") - assertContains(t, res[1], "file://localhost/C:/Users/jdoe/Pictures/two.webp") - assertContains(t, res[2], "file:///C:/Users/jdoe/Pictures/three.jpg") - assertContains(t, res[3], `.\relative\four.png`) -} - -func TestNormalizePathFileURL(t *testing.T) { - got := NormalizePath("file:///C:/Users/jdoe/Pictures/img.png") - want := filepath.FromSlash("C:/Users/jdoe/Pictures/img.png") - if got != want { - t.Fatalf("path = %q, want %q", got, want) - } -} - -func TestExtractRemovesQuotedFilepath(t *testing.T) { - dir := t.TempDir() - fp := filepath.Join(dir, "img.jpg") - data := make([]byte, 600) - copy(data, []byte{ - 0xff, 0xd8, 0xff, 0xe0, 0x00, 0x10, 'J', 'F', 'I', 'F', - 0x00, 0x01, 0x01, 0x01, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, - 0xff, 0xd9, - }) - if err := os.WriteFile(fp, data, 0o600); err != nil { - t.Fatalf("failed to write test image: %v", err) - } - - input := "before '" + fp + "' after" - cleaned, imgs, err := Extract(input) - if err != nil { - t.Fatalf("err: %v", err) - } - if len(imgs) != 1 { - t.Fatalf("imgs = %d, want 1", len(imgs)) - } - if cleaned != "before after" { - t.Fatalf("cleaned = %q, want %q", cleaned, "before after") - } -} - -func TestExtractFileURL(t *testing.T) { - dir := t.TempDir() - fp := filepath.Join(dir, "img.png") - data := make([]byte, 600) - copy(data, []byte{0x89, 'P', 'N', 'G', '\r', '\n', 0x1a, '\n'}) - if err := os.WriteFile(fp, data, 0o600); err != nil { - t.Fatalf("failed to write test image: %v", err) - } - - fileURL := (&url.URL{Scheme: "file", Path: fp}).String() - cleaned, imgs, err := Extract("before " + fileURL + " after") - if err != nil { - t.Fatalf("err: %v", err) - } - if len(imgs) != 1 { - t.Fatalf("imgs = %d, want 1", len(imgs)) - } - if cleaned != "before after" { - t.Fatalf("cleaned = %q, want %q", cleaned, "before after") - } -} - -func TestExtractWAV(t *testing.T) { - dir := t.TempDir() - fp := filepath.Join(dir, "sample.wav") - data := make([]byte, 600) - copy(data[:44], []byte{ - 'R', 'I', 'F', 'F', - 0x58, 0x02, 0x00, 0x00, - 'W', 'A', 'V', 'E', - 'f', 'm', 't', ' ', - 0x10, 0x00, 0x00, 0x00, - 0x01, 0x00, - 0x01, 0x00, - 0x80, 0x3e, 0x00, 0x00, - 0x00, 0x7d, 0x00, 0x00, - 0x02, 0x00, - 0x10, 0x00, - 'd', 'a', 't', 'a', - 0x34, 0x02, 0x00, 0x00, - }) - if err := os.WriteFile(fp, data, 0o600); err != nil { - t.Fatalf("failed to write test audio: %v", err) - } - - input := "before " + fp + " after" - cleaned, imgs, err := Extract(input) - if err != nil { - t.Fatalf("err: %v", err) - } - if len(imgs) != 1 { - t.Fatalf("imgs = %d, want 1", len(imgs)) - } - if cleaned != "before after" { - t.Fatalf("cleaned = %q, want %q", cleaned, "before after") - } -} - -func assertContains(t *testing.T, s, want string) { - t.Helper() - if !strings.Contains(s, want) { - t.Fatalf("%q does not contain %q", s, want) - } -} - -func assertNotContains(t *testing.T, s, want string) { - t.Helper() - if strings.Contains(s, want) { - t.Fatalf("%q unexpectedly contains %q", s, want) - } -} - -func assertNotContainsSlice(t *testing.T, ss []string, want string) { - t.Helper() - for _, s := range ss { - if strings.Contains(s, want) { - t.Fatalf("slice unexpectedly contains %q in %q", want, s) - } - } -} diff --git a/cmd/tui/chat/approval.go b/cmd/tui/chat/approval.go deleted file mode 100644 index d9a76a02eaa..00000000000 --- a/cmd/tui/chat/approval.go +++ /dev/null @@ -1,475 +0,0 @@ -package chat - -import ( - "context" - "encoding/json" - "fmt" - "slices" - "strings" - "time" - - tea "github.com/charmbracelet/bubbletea" - - coreagent "github.com/ollama/ollama/agent" -) - -type chatApprovalChoice struct { - label string - key string - allow bool - allowTools bool - allowAll bool - reason string -} - -var chatApprovalChoices = []chatApprovalChoice{ - {label: "Approve once", key: "1", allow: true}, - {label: "Always allow tool", key: "2", allow: true, allowTools: true}, - {label: "Deny", key: "3", reason: "Tool execution denied."}, -} - -type chatApprovalPrompt struct { - request coreagent.ApprovalRequest - reply chan<- coreagent.Approval - cursor int -} - -func (m chatModel) approvalPrompterForRun(controller *chatApprovalController) coreagent.ApprovalPrompter { - if m.opts.ApprovalPrompter != nil { - return m.opts.ApprovalPrompter - } - return controller -} - -func (m *chatModel) ensureApprovalState() *coreagent.ApprovalState { - if m.approvalState == nil { - m.approvalState = &coreagent.ApprovalState{} - m.approvalState.Set(m.defaultAllowAll, nil) - } - return m.approvalState -} - -func (m *chatModel) resetApprovalState() { - m.approvalState = &coreagent.ApprovalState{} - m.approvalState.Set(m.defaultAllowAll, nil) -} - -func (m chatModel) allowAllToolsEnabled() bool { - if m.approvalState == nil { - return m.defaultAllowAll - } - return m.approvalState.AllGranted() -} - -func (m *chatModel) setAllowAllTools(allowAll bool) { - if allowAll { - m.ensureApprovalState().GrantAll() - } else { - m.ensureApprovalState().Set(false, nil) - } - m.opts.AllowAllTools = allowAll -} - -func (m *chatModel) openApprovalPrompt(msg chatApprovalPromptMsg) { - m.approvalPrompt = &chatApprovalPrompt{request: msg.request, reply: msg.reply} - m.status = "approval required" - m.thinking = false - m.thinkingTokens = 0 - m.upsertApprovalToolEntries(msg.request) -} - -func (m *chatModel) togglePermissionMode() (tea.Model, tea.Cmd) { - m.setAllowAllTools(!m.allowAllToolsEnabled()) - if m.allowAllToolsEnabled() { - m.permissionNotice = "full access enabled" - m.status = "full access enabled" - if m.approvalPrompt != nil { - updated, cmd := m.resolveApprovalPrompt(chatApprovalChoice{allow: true, allowAll: true}) - if model, ok := updated.(chatModel); ok { - model.permissionNotice = "full access enabled" - model.status = "full access enabled" - return model, cmd - } - return updated, cmd - } - return *m, nil - } - m.permissionNotice = "review mode enabled" - m.status = "review mode enabled" - return *m, nil -} - -func (m *chatModel) upsertApprovalToolEntries(request coreagent.ApprovalRequest) { - for _, call := range request.Calls { - idx := m.findToolEntry(call.ToolCallID) - if idx < 0 { - m.groupCompletedToolHistory() - m.entries = append(m.entries, newChatEntry(chatEntry{role: "tool"})) - idx = len(m.entries) - 1 - } - m.entries[idx].detail = call.ToolName - m.entries[idx].label = toolInvocationLabel(call.ToolName, call.Args) - m.entries[idx].status = "approval" - m.entries[idx].toolID = call.ToolCallID - m.entries[idx].args = call.Args - m.entries[idx].startedAt = time.Now() - m.applyToolOutputModeTo(idx) - m.markEntryDirty(idx) - } -} - -func (m chatModel) updateApprovalPrompt(msg tea.KeyMsg) (tea.Model, tea.Cmd) { - switch msg.Type { - case tea.KeyLeft, tea.KeyUp: - m.moveApprovalChoice(-1) - case tea.KeyRight, tea.KeyDown, tea.KeyTab: - m.moveApprovalChoice(1) - case tea.KeyRunes: - switch string(msg.Runes) { - case "1", "2", "3": - choice := chatApprovalChoices[int(msg.Runes[0]-'1')] - return m.resolveApprovalPrompt(choice) - } - case tea.KeyEnter: - choice := chatApprovalChoices[clamp(m.approvalPrompt.cursor, 0, len(chatApprovalChoices)-1)] - return m.resolveApprovalPrompt(choice) - case tea.KeyEsc, tea.KeyCtrlC: - return m.resolveApprovalPrompt(chatApprovalChoice{reason: "Tool execution denied."}) - } - return m, nil -} - -func (m *chatModel) moveApprovalChoice(delta int) { - if m.approvalPrompt == nil { - return - } - m.approvalPrompt.cursor = (m.approvalPrompt.cursor + delta) % len(chatApprovalChoices) - if m.approvalPrompt.cursor < 0 { - m.approvalPrompt.cursor += len(chatApprovalChoices) - } - m.markApprovalPromptEntryDirty() -} - -func (m *chatModel) markApprovalPromptEntryDirty() { - if m.approvalPrompt == nil { - return - } - for _, call := range m.approvalPrompt.request.Calls { - if idx := m.findToolEntry(call.ToolCallID); idx >= 0 { - m.markEntryDirty(idx) - } - } -} - -func (m chatModel) resolveApprovalPrompt(choice chatApprovalChoice) (tea.Model, tea.Cmd) { - if m.approvalPrompt == nil { - return m, nil - } - printedLines := m.flowPrintedLines - var printedTranscript []string - if printedLines > 0 { - printedTranscript = slices.Clone(m.transcriptLines(m.viewWidth())) - } - prompt := m.approvalPrompt - m.approvalPrompt = nil - m.status = "running" - if !choice.allow { - m.status = "denied" - } - if choice.allowAll { - m.setAllowAllTools(true) - } - allowScopes := approvalScopes(prompt.request) - if choice.allowTools { - m.ensureApprovalState().GrantScopes(allowScopes) - } - for _, call := range prompt.request.Calls { - if idx := m.findToolEntry(call.ToolCallID); idx >= 0 && m.entries[idx].status == "approval" { - if !choice.allow { - m.entries[idx].status = "error" - m.entries[idx].err = choice.reason - if m.entries[idx].err == "" { - m.entries[idx].err = "Tool execution denied." - } - } else { - m.entries[idx].status = "queued" - } - m.markEntryDirty(idx) - } - } - result := coreagent.Approval{Allow: choice.allow, AllowAll: choice.allowAll, Reason: choice.reason} - if choice.allowTools { - result.AllowScopes = allowScopes - } - prompt.reply <- result - return m.withFlowTranscriptRefreshAfter(printedTranscript, printedLines, waitForChatMsg(m.events)) -} - -func (m chatModel) renderApprovalPromptLines(width int) []string { - prompt := m.approvalPrompt - if prompt == nil { - return nil - } - if width <= 0 { - width = 80 - } - bodyWidth := max(20, width-2) - - var lines []string - if len(prompt.request.Calls) <= 1 { - detail := approvalRequestDetail(prompt.request, bodyWidth) - if detail == "" { - label := "Tool request" - if len(prompt.request.Calls) == 1 { - label = toolDisplayName(prompt.request.Calls[0].ToolName) - } - lines = append(lines, wrapChatText(fmt.Sprintf("%s wants to run", label), width)...) - } else { - lines = append(lines, indentLines(splitRenderedBody(detail), " ")...) - } - lines = append(lines, "") - } - - lines = append(lines, indentLines(renderApprovalChoices(prompt.request, prompt.cursor, bodyWidth), " ")...) - return lines -} - -func approvalRequestDetail(request coreagent.ApprovalRequest, width int) string { - if len(request.Calls) == 0 { - return "" - } - if len(request.Calls) == 1 { - return approvalToolCallDetail(request.Calls[0], width) - } - lines := make([]string, 0, len(request.Calls)) - for _, call := range request.Calls { - lines = append(lines, toolInvocationLabel(call.ToolName, call.Args)) - } - return chatMetaStyle.Render(strings.Join(lines, "\n")) -} - -func approvalToolCallDetail(call coreagent.ApprovalToolCall, width int) string { - if isShellToolName(call.ToolName) { - command, ok := rawStringArg(call.Args, "command") - if !ok { - return "" - } - return strings.Join(wrapChatText(shellPromptPrefix(call.ToolName)+command, width), "\n") - } - switch call.ToolName { - case "edit": - path, ok := rawStringArg(call.Args, "path") - if !ok { - return "" - } - var lines []string - lines = append(lines, "path: "+path) - if summary := editApprovalSummary(call.Args); summary != "" { - lines = append(lines, strings.Split(summary, "\n")...) - } - return chatMetaStyle.Render(strings.Join(lines, "\n")) - default: - if len(call.Args) == 0 { - return "" - } - return strings.Join(renderToolCallArgs(call.Args, width), "\n") - } -} - -// editApprovalSummary summarizes edit tool arguments for the approval dialog: -// the edit count and total sizes for the edits array form, or the legacy -// top-level old_text/new_text sizes. -func editApprovalSummary(args map[string]any) string { - edits := rawEditEntries(args) - if len(edits) > 0 { - oldChars, newChars := 0, 0 - for _, entry := range edits { - if oldText, ok := editEntryText(entry, "old_text"); ok { - oldChars += len([]rune(oldText)) - } - if newText, ok := editEntryText(entry, "new_text"); ok { - newChars += len([]rune(newText)) - } - } - return fmt.Sprintf("edits: %d (old: %d chars, new: %d chars)", len(edits), oldChars, newChars) - } - if raw, ok := args["edits"]; ok { - if s, ok := raw.(string); ok { - return fmt.Sprintf("edits: %d chars", len([]rune(s))) - } - } - - var lines []string - if oldText, ok := rawStringArg(args, "old_text"); ok { - lines = append(lines, fmt.Sprintf("old_text: %d chars", len([]rune(oldText)))) - } - if newText, ok := rawStringArg(args, "new_text"); ok { - lines = append(lines, fmt.Sprintf("new_text: %d chars", len([]rune(newText)))) - } - return strings.Join(lines, "\n") -} - -// rawEditEntries extracts the edits array entries from tool arguments, -// tolerating a JSON string encoding and camelCase keys. -func rawEditEntries(args map[string]any) []map[string]any { - raw, ok := args["edits"] - if !ok { - return nil - } - if s, ok := raw.(string); ok { - var decoded []map[string]any - if err := json.Unmarshal([]byte(s), &decoded); err != nil { - return nil - } - return decoded - } - items, ok := raw.([]any) - if !ok { - return nil - } - entries := make([]map[string]any, 0, len(items)) - for _, item := range items { - if entry, ok := item.(map[string]any); ok { - entries = append(entries, entry) - } - } - return entries -} - -func editEntryText(entry map[string]any, snake string) (string, bool) { - if value, ok := entry[snake].(string); ok { - return value, true - } - camel := strings.TrimSuffix(snake, "_text") + "Text" - if value, ok := entry[camel].(string); ok { - return value, true - } - return "", false -} - -func renderApprovalChoices(request coreagent.ApprovalRequest, cursor int, width int) []string { - var lines []string - for i, choice := range chatApprovalChoices { - label := choice.key + ". " + approvalChoiceLabel(choice, request) - wrapped := wrapChatText(label, max(20, width-2)) - if i == clamp(cursor, 0, len(chatApprovalChoices)-1) { - for j, line := range wrapped { - if j == 0 { - lines = append(lines, chatPickerSelectedStyle.Render("> "+line)) - } else { - lines = append(lines, chatPickerSelectedStyle.Render(" "+line)) - } - } - } else { - for _, line := range wrapped { - lines = append(lines, chatPickerTextStyle.Render(" "+line)) - } - } - } - return lines -} - -func approvalChoiceLabel(choice chatApprovalChoice, request coreagent.ApprovalRequest) string { - if !choice.allowTools { - return choice.label - } - scopes := approvalScopes(request) - if len(scopes) == 1 { - call := approvalCallForScope(request, scopes[0]) - if isShellToolName(call.ToolName) { - if command, ok := rawStringArg(call.Args, "command"); ok && strings.TrimSpace(command) != "" { - return "Always allow this command" - } - } - return "Always allow " + toolDisplayName(call.ToolName) - } - return "Always allow these requests" -} - -func approvalScopes(request coreagent.ApprovalRequest) []string { - seen := make(map[string]bool, len(request.Calls)) - var scopes []string - for _, call := range request.Calls { - scope := approvalScope(call) - if scope == "" || seen[scope] { - continue - } - seen[scope] = true - scopes = append(scopes, scope) - } - return scopes -} - -func approvalCallForScope(request coreagent.ApprovalRequest, scope string) coreagent.ApprovalToolCall { - for _, call := range request.Calls { - if approvalScope(call) == scope { - return call - } - } - return coreagent.ApprovalToolCall{} -} - -func approvalScope(call coreagent.ApprovalToolCall) string { - if scope := strings.TrimSpace(call.ApprovalScope); scope != "" { - return scope - } - return strings.TrimSpace(call.ToolName) -} - -type chatApprovalPrompter struct { - ch chan<- tea.Msg -} - -func (p chatApprovalPrompter) PromptApproval(ctx context.Context, request coreagent.ApprovalRequest) (coreagent.Approval, error) { - reply := make(chan coreagent.Approval, 1) - select { - case p.ch <- chatApprovalPromptMsg{request: request, reply: reply}: - case <-ctx.Done(): - return coreagent.Approval{Reason: "Tool approval canceled."}, nil - } - - select { - case result := <-reply: - return result, nil - case <-ctx.Done(): - return coreagent.Approval{Reason: "Tool approval canceled."}, nil - } -} - -type chatApprovalController struct { - ch chan<- tea.Msg - state *coreagent.ApprovalState -} - -func newChatApprovalController(ch chan<- tea.Msg, state *coreagent.ApprovalState) *chatApprovalController { - return &chatApprovalController{ - ch: ch, - state: state, - } -} - -func (c *chatApprovalController) PromptApproval(ctx context.Context, request coreagent.ApprovalRequest) (coreagent.Approval, error) { - if result, ok := c.preapproved(request); ok { - return result, nil - } - return chatApprovalPrompter{ch: c.ch}.PromptApproval(ctx, request) -} - -func (c *chatApprovalController) preapproved(request coreagent.ApprovalRequest) (coreagent.Approval, bool) { - if c == nil { - return coreagent.Approval{}, false - } - if c.state.AllGranted() { - return coreagent.Approval{Allow: true, AllowAll: true}, true - } - scopes := approvalScopes(request) - if len(scopes) == 0 { - return coreagent.Approval{}, false - } - for _, scope := range scopes { - if !c.state.Allows(scope) { - return coreagent.Approval{}, false - } - } - return coreagent.Approval{Allow: true, AllowScopes: scopes}, true -} diff --git a/cmd/tui/chat/approval_test.go b/cmd/tui/chat/approval_test.go deleted file mode 100644 index ca504b39a3d..00000000000 --- a/cmd/tui/chat/approval_test.go +++ /dev/null @@ -1,520 +0,0 @@ -package chat - -import ( - "context" - "strings" - "testing" - "time" - - tea "github.com/charmbracelet/bubbletea" - - coreagent "github.com/ollama/ollama/agent" -) - -func testApprovalRequest() coreagent.ApprovalRequest { - return coreagent.ApprovalRequest{ - WorkingDir: "/repo", - Calls: []coreagent.ApprovalToolCall{{ - ToolCallID: "call-1", - ToolName: "edit", - Args: map[string]any{"path": "note.txt"}, - ApprovalScope: "edit", - }}, - } -} - -func testApprovalState(allowAll bool, scopes map[string]bool) *coreagent.ApprovalState { - state := &coreagent.ApprovalState{} - state.Set(allowAll, scopes) - return state -} - -func TestChatApprovalApprovesOnce(t *testing.T) { - reply := make(chan coreagent.Approval, 1) - m := chatModel{ - approvalPrompt: &chatApprovalPrompt{ - request: testApprovalRequest(), - reply: reply, - }, - events: make(chan tea.Msg), - } - - updated, cmd := m.updateApprovalPrompt(tea.KeyMsg{Type: tea.KeyEnter}) - if cmd == nil { - t.Fatal("approval should resume waiting for agent events") - } - fm := updated.(chatModel) - if fm.approvalPrompt != nil { - t.Fatal("approval prompt should close") - } - result := <-reply - if !result.Allow || result.AllowAll { - t.Fatalf("approval = %#v, want allow once", result) - } -} - -func TestChatApprovalAllowsTool(t *testing.T) { - reply := make(chan coreagent.Approval, 1) - m := chatModel{ - approvalPrompt: &chatApprovalPrompt{ - request: testApprovalRequest(), - reply: reply, - cursor: 1, - }, - events: make(chan tea.Msg), - } - - updated, _ := m.updateApprovalPrompt(tea.KeyMsg{Type: tea.KeyEnter}) - fm := updated.(chatModel) - if fm.allowAllToolsEnabled() { - t.Fatal("allowing a tool should not enable full access") - } - if !fm.approvalState.Allows("edit") { - t.Fatal("edit scope was not saved") - } - result := <-reply - if !result.Allow || result.AllowAll || len(result.AllowScopes) != 1 || result.AllowScopes[0] != "edit" { - t.Fatalf("approval = %#v, want per-tool approval", result) - } -} - -func TestChatApprovalLabelsSecondChoiceAsPerTool(t *testing.T) { - lines := stripANSI(strings.Join(renderApprovalChoices(testApprovalRequest(), 1, 80), "\n")) - if !strings.Contains(lines, "2. Always allow Edit") { - t.Fatalf("approval choices = %q, want per-tool option", lines) - } - if strings.Contains(lines, "Approve all") { - t.Fatalf("approval choices = %q, should not offer approve all as option 2", lines) - } -} - -func TestChatApprovalLabelsShellChoiceAsCommandScoped(t *testing.T) { - request := coreagent.ApprovalRequest{ - WorkingDir: "/repo", - Calls: []coreagent.ApprovalToolCall{{ - ToolCallID: "call-1", - ToolName: "bash", - Args: map[string]any{"command": "pwd"}, - ApprovalScope: "bash\x00pwd", - }}, - } - lines := stripANSI(strings.Join(renderApprovalChoices(request, 1, 80), "\n")) - if !strings.Contains(lines, "2. Always allow this command") { - t.Fatalf("approval choices = %q, want command-scoped option", lines) - } - if strings.Contains(lines, "Always allow Bash") { - t.Fatalf("approval choices = %q, should not offer top-level Bash approval", lines) - } -} - -func TestChatApprovalUsesShellNameForPermissionPrompt(t *testing.T) { - request := coreagent.ApprovalRequest{ - WorkingDir: "/repo", - Calls: []coreagent.ApprovalToolCall{{ - ToolCallID: "call-1", - ToolName: "bash", - Args: map[string]any{"command": "pwd"}, - ApprovalScope: "bash\x00pwd", - }}, - } - - detail := stripANSI(approvalRequestDetail(request, 80)) - if !strings.Contains(detail, "$ pwd") { - t.Fatalf("approval detail should show command prompt, got %q", detail) - } - - m := chatModel{} - m.upsertApprovalToolEntries(request) - if len(m.entries) != 1 { - t.Fatalf("entries = %#v", m.entries) - } - line := stripANSI(toolStatusLine(m.entries[0])) - if !strings.Contains(line, `Bash("pwd")`) || !strings.Contains(line, "needs approval") { - t.Fatalf("approval status line = %q", line) - } -} - -func TestChatApprovalRendersSkillLoad(t *testing.T) { - request := coreagent.ApprovalRequest{ - WorkingDir: "/repo", - Calls: []coreagent.ApprovalToolCall{{ - ToolCallID: "call-skill-1", - ToolName: "skill", - Args: map[string]any{"name": "release-notes"}, - ApprovalScope: "skill", - }}, - } - - lines := stripANSI(strings.Join((&chatModel{approvalPrompt: &chatApprovalPrompt{request: request}}).renderApprovalPromptLines(80), "\n")) - for _, want := range []string{"name: release-notes", "2. Always allow skill"} { - if !strings.Contains(lines, want) { - t.Fatalf("skill approval prompt missing %q:\n%s", want, lines) - } - } - - m := chatModel{} - m.upsertApprovalToolEntries(request) - if len(m.entries) != 1 || !strings.Contains(stripANSI(toolStatusLine(m.entries[0])), `skill("release-notes") needs approval`) { - t.Fatalf("skill approval entry = %#v", m.entries) - } -} - -func TestChatApprovalPromptOmitsDuplicateBatchDetails(t *testing.T) { - request := coreagent.ApprovalRequest{ - WorkingDir: "/repo", - Calls: []coreagent.ApprovalToolCall{ - { - ToolCallID: "call-1", - ToolName: "bash", - Args: map[string]any{"command": "git rev-parse --abbrev-ref HEAD"}, - ApprovalScope: "bash\x00git rev-parse --abbrev-ref HEAD", - }, - { - ToolCallID: "call-2", - ToolName: "bash", - Args: map[string]any{"command": "git branch -a"}, - ApprovalScope: "bash\x00git branch -a", - }, - }, - } - m := chatModel{ - approvalPrompt: &chatApprovalPrompt{request: request}, - } - - lines := stripANSI(strings.Join(m.renderApprovalPromptLines(120), "\n")) - if strings.Contains(lines, `Bash("git rev-parse --abbrev-ref HEAD")`) || strings.Contains(lines, `Bash("git branch -a")`) { - t.Fatalf("batched approval prompt should not duplicate visible tool rows:\n%s", lines) - } - for _, want := range []string{"1. Approve once", "2. Always allow these requests", "3. Deny"} { - if !strings.Contains(lines, want) { - t.Fatalf("batched approval prompt missing %q:\n%s", want, lines) - } - } -} - -func TestChatApprovalKeepsQueuedBatchCallsVisible(t *testing.T) { - reply := make(chan coreagent.Approval, 1) - request := coreagent.ApprovalRequest{ - WorkingDir: "/repo", - Calls: []coreagent.ApprovalToolCall{ - { - ToolCallID: "call-1", - ToolName: "bash", - Args: map[string]any{"command": "git rev-parse --abbrev-ref HEAD"}, - ApprovalScope: "bash\x00git rev-parse --abbrev-ref HEAD", - }, - { - ToolCallID: "call-2", - ToolName: "bash", - Args: map[string]any{"command": "git branch -a"}, - ApprovalScope: "bash\x00git branch -a", - }, - }, - } - m := chatModel{ - running: true, - events: make(chan tea.Msg), - } - m.openApprovalPrompt(chatApprovalPromptMsg{request: request, reply: reply}) - - updated, _ := m.resolveApprovalPrompt(chatApprovalChoice{allow: true}) - m = updated.(chatModel) - m.applyAgentEvent(coreagent.Event{ - Type: coreagent.EventToolStarted, - ToolCallID: "call-1", - ToolName: "bash", - Args: request.Calls[0].Args, - }) - m.applyAgentEvent(coreagent.Event{ - Type: coreagent.EventToolFinished, - ToolCallID: "call-1", - ToolName: "bash", - Args: request.Calls[0].Args, - Content: "parth-agent-tui\n", - }) - m.applyAgentEvent(coreagent.Event{Type: coreagent.EventMessageDelta, Content: "I'm on branch parth-agent-tui."}) - - transcript := stripANSI(m.renderTranscript(180)) - for _, want := range []string{ - `Bash("git rev-parse --abbrev-ref HEAD")`, - `Bash("git branch -a")`, - "I'm on branch parth-agent-tui.", - } { - if !strings.Contains(transcript, want) { - t.Fatalf("transcript missing %q:\n%s", want, transcript) - } - } -} - -func TestChatApprovalPromptRepaintsFlowTranscript(t *testing.T) { - reply := make(chan coreagent.Approval, 1) - request := coreagent.ApprovalRequest{ - WorkingDir: "/repo", - Calls: []coreagent.ApprovalToolCall{ - { - ToolCallID: "call-1", - ToolName: "web_fetch", - Args: map[string]any{"url": "https://parthsareen.com/"}, - ApprovalScope: "web_fetch", - }, - { - ToolCallID: "call-2", - ToolName: "web_fetch", - Args: map[string]any{"url": "https://github.com/ParthSareen"}, - ApprovalScope: "web_fetch", - }, - }, - } - m := chatModel{ - running: true, - width: 160, - flowPrintedLines: 1, - entries: []chatEntry{ - {role: "user", content: "research parth"}, - }, - } - - updated, cmd := m.Update(chatApprovalPromptMsg{request: request, reply: reply}) - if cmd == nil { - t.Fatal("opening approval should repaint flow transcript") - } - fm := updated.(chatModel) - transcript := stripANSI(fm.renderTranscript(160)) - for _, want := range []string{ - `Web Fetch("https://parthsareen.com/") needs approval`, - `Web Fetch("https://github.com/ParthSareen") needs approval`, - } { - if !strings.Contains(transcript, want) { - t.Fatalf("transcript missing %q:\n%s", want, transcript) - } - } -} - -func TestChatApprovalResolutionRepaintsFlowTranscript(t *testing.T) { - reply := make(chan coreagent.Approval, 1) - request := coreagent.ApprovalRequest{ - WorkingDir: "/repo", - Calls: []coreagent.ApprovalToolCall{ - { - ToolCallID: "call-1", - ToolName: "web_fetch", - Args: map[string]any{"url": "https://parthsareen.com/"}, - ApprovalScope: "web_fetch", - }, - { - ToolCallID: "call-2", - ToolName: "web_fetch", - Args: map[string]any{"url": "https://github.com/ParthSareen"}, - ApprovalScope: "web_fetch", - }, - }, - } - m := chatModel{ - running: true, - width: 160, - events: make(chan tea.Msg), - entries: []chatEntry{ - {role: "user", content: "research parth"}, - }, - } - m.openApprovalPrompt(chatApprovalPromptMsg{request: request, reply: reply}) - printed := len(m.transcriptLines(160)) - m.flowPrintedLines = printed - - updated, cmd := m.resolveApprovalPrompt(chatApprovalChoice{allow: true}) - if cmd == nil { - t.Fatal("approval resolution should keep waiting for agent events") - } - fm := updated.(chatModel) - if fm.flowPrintedLines >= printed { - t.Fatalf("approval resolution should repaint and hold queued rows, flowPrintedLines = %d, was %d", fm.flowPrintedLines, printed) - } - if result := <-reply; !result.Allow { - t.Fatalf("approval = %#v, want allow", result) - } -} - -func TestChatApprovalBatchCollapsesAtNextToolBoundary(t *testing.T) { - reply := make(chan coreagent.Approval, 1) - request := coreagent.ApprovalRequest{ - WorkingDir: "/repo", - Calls: []coreagent.ApprovalToolCall{ - { - ToolCallID: "call-1", - ToolName: "bash", - Args: map[string]any{"command": "git rev-parse --abbrev-ref HEAD"}, - ApprovalScope: "bash\x00git rev-parse --abbrev-ref HEAD", - }, - { - ToolCallID: "call-2", - ToolName: "bash", - Args: map[string]any{"command": "git branch -a"}, - ApprovalScope: "bash\x00git branch -a", - }, - }, - } - m := chatModel{ - running: true, - events: make(chan tea.Msg), - } - m.openApprovalPrompt(chatApprovalPromptMsg{request: request, reply: reply}) - - updated, _ := m.resolveApprovalPrompt(chatApprovalChoice{allow: true}) - m = updated.(chatModel) - for _, call := range request.Calls { - m.applyAgentEvent(coreagent.Event{ - Type: coreagent.EventToolStarted, - ToolCallID: call.ToolCallID, - ToolName: call.ToolName, - Args: call.Args, - }) - m.applyAgentEvent(coreagent.Event{ - Type: coreagent.EventToolFinished, - ToolCallID: call.ToolCallID, - ToolName: call.ToolName, - Args: call.Args, - Content: "ok\n", - }) - } - m.applyAgentEvent(coreagent.Event{Type: coreagent.EventMessageDelta, Content: "I'm on branch parth-agent-tui."}) - - transcript := stripANSI(m.renderTranscript(180)) - if strings.Contains(transcript, "Ran 2 commands") { - t.Fatalf("completed batch should stay expanded until the next tool boundary:\n%s", transcript) - } - if !strings.Contains(transcript, `Bash("git branch -a")`) { - t.Fatalf("completed batch should keep concrete command rows before the next boundary:\n%s", transcript) - } - - m.applyAgentEvent(coreagent.Event{ - Type: coreagent.EventToolStarted, - ToolCallID: "call-3", - ToolName: "bash", - Args: map[string]any{"command": "git status --short"}, - }) - transcript = stripANSI(m.renderTranscript(180)) - if !strings.Contains(transcript, "Ran 2 commands") { - t.Fatalf("completed batch should collapse when a new tool starts:\n%s", transcript) - } - if !strings.Contains(transcript, `Bash("git status --short")`) { - t.Fatalf("new running command should remain concrete after previous batch collapses:\n%s", transcript) - } -} - -func TestChatApprovalPrompterCancels(t *testing.T) { - ctx, cancel := context.WithCancel(context.Background()) - cancel() - result, err := (chatApprovalPrompter{ch: make(chan tea.Msg)}).PromptApproval(ctx, testApprovalRequest()) - if err != nil { - t.Fatal(err) - } - if result.Allow || result.Reason == "" { - t.Fatalf("approval = %#v, want canceled denial", result) - } -} - -func TestChatApprovalControllerAutoApprovesAfterFullAccessToggle(t *testing.T) { - events := make(chan tea.Msg, 1) - state := testApprovalState(false, nil) - controller := newChatApprovalController(events, state) - state.GrantAll() - - result, err := controller.PromptApproval(context.Background(), testApprovalRequest()) - if err != nil { - t.Fatal(err) - } - if !result.Allow || !result.AllowAll { - t.Fatalf("approval = %#v, want full-access approval", result) - } - select { - case msg := <-events: - t.Fatalf("approval UI event should not be sent after full access toggle: %#v", msg) - default: - } -} - -func TestChatPermissionToggleSyncsRunningApprovalController(t *testing.T) { - events := make(chan tea.Msg, 1) - state := testApprovalState(false, nil) - m := chatModel{ - approvalState: state, - approvalController: newChatApprovalController(events, state), - } - - updated, _ := m.togglePermissionMode() - fm := updated.(chatModel) - result, err := fm.approvalController.PromptApproval(context.Background(), testApprovalRequest()) - if err != nil { - t.Fatal(err) - } - if !result.Allow || !result.AllowAll { - t.Fatalf("approval = %#v, want full-access approval", result) - } -} - -func TestChatPermissionToggleFromFullAccessRequiresReviewInRunningController(t *testing.T) { - events := make(chan tea.Msg, 1) - state := testApprovalState(true, nil) - m := chatModel{ - approvalState: state, - approvalController: newChatApprovalController(events, state), - } - - updated, _ := m.togglePermissionMode() - fm := updated.(chatModel) - if fm.allowAllToolsEnabled() { - t.Fatal("full access should be disabled") - } - - resultCh := make(chan coreagent.Approval, 1) - go func() { - result, err := fm.approvalController.PromptApproval(context.Background(), testApprovalRequest()) - if err != nil { - resultCh <- coreagent.Approval{Reason: err.Error()} - return - } - resultCh <- result - }() - - select { - case msg := <-events: - prompt, ok := msg.(chatApprovalPromptMsg) - if !ok { - t.Fatalf("event = %#v, want approval prompt", msg) - } - prompt.reply <- coreagent.Approval{Reason: "denied"} - case <-time.After(time.Second): - t.Fatal("expected approval prompt after toggling from full access to review") - } - - result := <-resultCh - if result.Allow { - t.Fatalf("approval = %#v, want review prompt result", result) - } -} - -func TestChatApprovalPromptSkippedWhenFullAccessEnabledInFlight(t *testing.T) { - reply := make(chan coreagent.Approval, 1) - // Full access is on by the time the buffered approval request reaches the - // UI (toggled after the agent sent the request but before Update ran). - // The stale prompt must not surface; the request is auto-approved. - m := chatModel{approvalState: testApprovalState(true, nil), running: true} - - updated, _ := m.Update(chatApprovalPromptMsg{request: testApprovalRequest(), reply: reply}) - fm := updated.(chatModel) - - if fm.approvalPrompt != nil { - t.Fatalf("approval prompt = %#v, want nil (full access on)", fm.approvalPrompt) - } - if got := fm.status; got == "approval required" { - t.Fatalf("status = %q, should not show approval required", got) - } - select { - case result := <-reply: - if !result.Allow || !result.AllowAll { - t.Fatalf("approval = %#v, want full-access approval", result) - } - default: - t.Fatal("expected auto-approval sent on the reply channel") - } -} diff --git a/cmd/tui/chat/chat.go b/cmd/tui/chat/chat.go deleted file mode 100644 index 8606f8770d6..00000000000 --- a/cmd/tui/chat/chat.go +++ /dev/null @@ -1,1246 +0,0 @@ -package chat - -import ( - "context" - "errors" - "fmt" - "hash/fnv" - "runtime" - "slices" - "strings" - - tea "github.com/charmbracelet/bubbletea" - - coreagent "github.com/ollama/ollama/agent" - "github.com/ollama/ollama/api" - "github.com/ollama/ollama/cmd/internal/filedata" -) - -var chatSpinnerFrames = []string{".", "..", "..."} - -var chatRuntimeGOOS = runtime.GOOS - -const ( - maxPickerItems = 8 - maxInlineModelPickerItems = 5 - maxSlashCompletions = 5 - maxPromptHistory = 50 - idleWorkingDelayTicks = 4 -) - -var chatEmptyPrompts = []string{ - `read this repo and tell me where to start`, - `what changed on this branch?`, - `run the tests and summarize failures`, - `find the riskiest code path in this folder`, - `search the web and compare the latest docs with this implementation`, - `summarize this file and suggest edits`, -} - -type ModelOption struct { - Name string - Description string - Recommended bool - RequiredPlan string - Cloud bool - AvailabilityBadge string - SignInURL string -} - -type Options struct { - Model string - OpenModelPicker bool - ChatID string - Messages []api.Message - Client coreagent.ChatClient - Tools *coreagent.Registry - Skills *coreagent.SkillCatalog - ImportSkills func(string) (coreagent.SkillImportResult, error) - ReloadSkills func() (*coreagent.SkillCatalog, error) - ToolRegistryForModel func(context.Context, string) *coreagent.Registry - ToolsDisabled bool - MultiModalForModel func(context.Context, string) bool - ModelOptions func(context.Context) ([]ModelOption, error) - OnModelSelected func(context.Context, string) error - SystemPromptForModel func(context.Context, string, *coreagent.Registry, bool) string - ApprovalPrompter coreagent.ApprovalPrompter - EventSinks []coreagent.EventSink - AllowAllTools bool - WorkingDir string - RootDir string - Format string - Options map[string]any - Think *api.ThinkValue - KeepAlive *api.Duration - Images []api.ImageData - MultiModal bool - Compactor coreagent.Compactor - ContextWindowTokens int - ContextWindowTokensForModel func(context.Context, string, int) int - PreloadModel func(context.Context, string, *api.ThinkValue) (int, error) - CheckCloudModel func(context.Context, string, string) error - OpenBrowser func(string) - PollCloudAuth func(context.Context) (string, bool, error) - CompactionThreshold float64 - SystemPrompt string -} - -type Result struct { - ChatID string - Messages []api.Message -} - -//nolint:containedctx // chatModel is a Bubble Tea session model; the context is the run-scoped cancellation root. -type chatModel struct { - ctx context.Context - opts Options - chatID string - messages []api.Message - liveMessages []api.Message - entries []chatEntry - workingDir string - - input []rune - inputCursor int - inputCursorSet bool - inputAttachments []chatInputAttachment - inputPastedTexts []chatInputPastedText - nextImageID int - nextAudioID int - nextPastedTextID int - promptHistory []string - promptCursor int - promptDraft []rune - promptActive bool - running bool - awaitingModel bool - compacting bool - cancel context.CancelFunc - events <-chan tea.Msg - compactEvents <-chan tea.Msg - detectedToolCalls []chatEntry - scroll int - toolOutputMode bool - toolOutputOpen bool - thinkingDetailsOpen bool - flowPrintedLines int - thinking bool - thinkingPhaseStart int - thinkingTokens int - compactingTokens int - contextTokens int - contextEstimate bool - modelPicker *chatModelPicker - modelPickerModels []ModelOption - thinkPicker *chatThinkPicker - promptDebug *chatPromptDebug - approvalPrompt *chatApprovalPrompt - approvalController *chatApprovalController - approvalState *coreagent.ApprovalState - cloudAuthPrompt *cloudAuthPrompt - pendingModel string - defaultAllowAll bool - permissionNotice string - selection chatSelection - - systemPromptDisabled bool - - width int - height int - status string - spinner int - tickActive bool - preloadingModel string - complete int - quitting bool - openModelOnInit bool - quitArmed bool - quitArmedKey string - escArmed bool - eventErrorRendered bool - err error -} - -type chatSelectionPoint struct { - line int - col int -} - -type chatSelection struct { - active bool - dragging bool - anchor chatSelectionPoint - cursor chatSelectionPoint -} - -func startChatSelection(selection *chatSelection, msg tea.MouseMsg, contains func(tea.MouseMsg) bool, point func(tea.MouseMsg) chatSelectionPoint) { - if !contains(msg) { - *selection = chatSelection{} - return - } - p := point(msg) - *selection = chatSelection{active: true, dragging: true, anchor: p, cursor: p} -} - -func dragChatSelection(selection *chatSelection, msg tea.MouseMsg, point func(tea.MouseMsg) chatSelectionPoint, scrollEdge func(tea.MouseMsg)) { - if !selection.active || !selection.dragging { - return - } - selection.cursor = point(msg) - scrollEdge(msg) -} - -func finishChatSelection(selection *chatSelection, msg tea.MouseMsg, point func(tea.MouseMsg) chatSelectionPoint) { - if !selection.active || !selection.dragging { - return - } - selection.cursor = point(msg) - selection.dragging = false - if selection.anchor == selection.cursor { - *selection = chatSelection{} - } -} - -type chatInputAttachment struct { - placeholder string - kind string - data api.ImageData -} - -func Run(ctx context.Context, opts Options) (*Result, error) { - if opts.RootDir == "" { - opts.RootDir = opts.WorkingDir - } - - approvalState := &coreagent.ApprovalState{} - approvalState.Set(opts.AllowAllTools, nil) - - m := chatModel{ - ctx: ctx, - opts: opts, - chatID: opts.ChatID, - messages: slices.Clone(opts.Messages), - workingDir: opts.WorkingDir, - approvalState: approvalState, - defaultAllowAll: opts.AllowAllTools, - promptHistory: initialPromptHistory(ctx, opts), - status: "ready", - openModelOnInit: opts.OpenModelPicker || (strings.TrimSpace(opts.Model) == "" && opts.ModelOptions != nil), - } - m.nextImageID, m.nextAudioID = nextInputAttachmentIDsFromMessages(m.messages) - m.nextPastedTextID = nextInputPastedTextIDFromMessages(m.messages) - m.entries = entriesFromMessages(m.messages) - // Context window is resolved post-load (chatModelPreloadDoneMsg) rather than - // here: for local models /api/ps only reports the running num_ctx after the - // model loads, and opts.ContextWindowTokens already holds Show's max as a - // pre-load fallback. Refreshing now would just re-derive that same value - // (and block construction on a network call). - m.contextTokens = m.estimatePromptTokens(m.messages, "") - m.contextEstimate = true - if m.openModelOnInit { - updated, cmd := m.openModelPicker("") - if cmd != nil { - return nil, errors.New("initial model picker returned an unexpected command") - } - m = updated.(chatModel) - } - if opts.PreloadModel != nil && strings.TrimSpace(opts.Model) != "" && !m.openModelOnInit { - m.preloadingModel = strings.TrimSpace(opts.Model) - } - - p := tea.NewProgram(m, tea.WithReportFocus()) - finalModel, err := p.Run() - if err != nil { - return nil, err - } - - fm := finalModel.(chatModel) - if fm.err != nil { - return nil, fm.err - } - return &Result{ChatID: fm.chatID, Messages: fm.messages}, nil -} - -func (m chatModel) Init() tea.Cmd { - var cmds []tea.Cmd - if m.preloadingModel != "" && m.opts.PreloadModel != nil { - cmds = append(cmds, preloadModelCmd(m.ctx, m.opts.PreloadModel, m.preloadingModel, m.opts.Think), chatTickCmd()) - } - if cmd := cloudModelPreflightCmd(m.ctx, m.opts, m.opts.Model, ""); cmd != nil && !m.openModelOnInit { - cmds = append(cmds, cmd) - } - return tea.Batch(cmds...) -} - -func (m chatModel) Update(msg tea.Msg) (tea.Model, tea.Cmd) { - if m.canEditInput() && isShiftEnterCSI(msg) { - m.insertInputNewline() - return m, nil - } - - switch msg := msg.(type) { - case tea.WindowSizeMsg: - wasSet := m.width > 0 || m.height > 0 - resized := wasSet && (m.width != msg.Width || m.height != msg.Height) - m.width = msg.Width - m.height = msg.Height - if resized { - return m.withFlowTranscriptRepaint(nil) - } - return m.withFlowTranscriptFlush(nil) - - case tea.FocusMsg: - return m, nil - - case chatTickMsg: - m.tickActive = false - if !m.running && !m.compacting && m.preloadingModel == "" { - return m, nil - } - m.spinner++ - cmd := m.scheduleTick() - return m, cmd - - case chatModelPreloadDoneMsg: - if msg.model != "" && msg.model != m.preloadingModel { - return m, nil - } - m.preloadingModel = "" - if msg.err != nil { - if isUnsupportedThinkingError(msg.err) && thinkRequestsThinking(m.opts.Think) { - m.opts.Think = &api.ThinkValue{Value: false} - m.status = fmt.Sprintf("Thinking disabled for %s", msg.model) - if msg.model != "" && m.opts.PreloadModel != nil { - m.preloadingModel = msg.model - return m, tea.Batch(preloadModelCmd(m.ctx, m.opts.PreloadModel, msg.model, m.opts.Think), m.scheduleTick()) - } - return m, nil - } - if !errors.Is(msg.err, context.Canceled) { - m.status = "error" - m.entries = append(m.entries, newChatEntry(chatEntry{role: "error", content: fmt.Sprintf("Could not load model: %v", msg.err), err: msg.err.Error()})) - } - return m, nil - } - if msg.contextWindowTokens > 0 { - m.updateContextWindowTokens(msg.contextWindowTokens) - } else { - m.refreshContextWindowTokens(msg.model) - } - return m, nil - - case chatAgentMsg: - printedLines := m.flowPrintedLines - var printedTranscript []string - if printedLines > 0 { - printedTranscript = slices.Clone(m.transcriptLines(m.viewWidth())) - } - m.applyAgentEvent(msg.event) - return m.withFlowTranscriptRefreshAfter(printedTranscript, printedLines, waitForChatMsg(m.events)) - - case chatClipboardErrorMsg: - if msg.err != nil { - m.status = "copy failed: " + msg.err.Error() - } - return m, nil - - case chatApprovalPromptMsg: - // Full access may have been enabled while this request was in flight - // (the user toggled it on after the agent sent the approval request - // but before this buffered message was handled). In that window - // togglePermissionMode's auto-resolve sees no open prompt yet, so - // short-circuit here rather than surfacing a stale approval prompt. - if m.allowAllToolsEnabled() { - m.approvalPrompt = &chatApprovalPrompt{request: msg.request, reply: msg.reply} - return m.resolveApprovalPrompt(chatApprovalChoice{allow: true, allowAll: true}) - } - printedLines := m.flowPrintedLines - var printedTranscript []string - if printedLines > 0 { - printedTranscript = slices.Clone(m.transcriptLines(m.viewWidth())) - } - m.modelPicker = nil - m.modelPickerModels = nil - m.openApprovalPrompt(msg) - return m.withFlowTranscriptRefreshAfter(printedTranscript, printedLines, nil) - - case chatRunDoneMsg: - wasCanceling := m.status == "canceling" - m.finishThinkingEntry() - m.running = false - m.awaitingModel = false - m.compacting = false - m.compactingTokens = 0 - m.cancel = nil - m.events = nil - m.approvalController = nil - m.thinking = false - m.thinkingTokens = 0 - m.approvalPrompt = nil - if msg.result != nil { - m.messages = msg.result.Messages - m.liveMessages = nil - if msg.result.WorkingDir != "" { - m.workingDir = msg.result.WorkingDir - } - // Context window is settled by preload (local num_ctx) or is - // static (cloud); no refresh needed post-run. - m.contextTokens = m.estimatePromptTokens(m.messages, "") - m.contextEstimate = true - if !messagesEndWithCompactionResult(m.messages) { - m.applyResponseMetrics(&msg.result.Latest) - } - } - if msg.result == nil { - m.finishLiveMessagesForStoppedRun(msg.newMessagesPersisted, msg.persistedMessages) - } - if wasCanceling || isChatContextCanceledError(msg.err) { - m.status = "Tell the model what to do instead." - return m.withFlowTranscriptFlush(nil) - } - if msg.err != nil { - if !m.eventErrorRendered { - m.entries = append(m.entries, newChatEntry(chatEntry{role: "error", content: msg.err.Error(), err: msg.err.Error()})) - } - m.status = "error" - return m.withFlowTranscriptFlush(nil) - } - m.status = "ready" - return m.withFlowTranscriptFlush(nil) - - case chatCompactDoneMsg: - return m.finishManualCompaction(msg) - - case chatCompactProgressMsg: - if msg.tokens > m.compactingTokens { - m.compactingTokens = msg.tokens - } - return m, waitForChatMsg(m.compactEvents) - - case chatEventsClosedMsg: - if m.compacting { - return m.Update(chatCompactDoneMsg{err: context.Canceled}) - } - if m.running { - return m.Update(chatRunDoneMsg{err: context.Canceled, newMessagesPersisted: true}) - } - return m, nil - - case cloudAuthTickMsg: - if m.cloudAuthPrompt != nil { - return m.updateCloudAuthPrompt(msg) - } - return m, nil - - case cloudAuthCheckMsg: - if m.cloudAuthPrompt != nil { - return m.updateCloudAuthPrompt(msg) - } - return m, nil - - case cloudModelPreflightMsg: - return m.updateCloudModelPreflight(msg) - - case cloudAuthPollMsg: - if m.cloudAuthPrompt != nil { - return m.updateCloudAuthPrompt(msg) - } - return m, nil - - case tea.MouseMsg: - return m.updateMouse(msg) - - case tea.KeyMsg: - return m.updateKey(msg) - } - return m, nil -} - -func (m chatModel) updateMouse(msg tea.MouseMsg) (tea.Model, tea.Cmd) { - if m.promptDebug != nil { - switch msg.Type { - case tea.MouseWheelUp: - m.promptDebug.scroll = clamp(m.promptDebug.scroll-3, 0, m.promptDebugMaxScroll()) - case tea.MouseWheelDown: - m.promptDebug.scroll = clamp(m.promptDebug.scroll+3, 0, m.promptDebugMaxScroll()) - } - return m, nil - } - if m.modelPicker != nil { - switch msg.Type { - case tea.MouseWheelUp: - m.modelPicker.Move(-3) - case tea.MouseWheelDown: - m.modelPicker.Move(3) - } - return m, nil - } - if !m.mouseInTranscript(msg) && !m.selection.active { - switch msg.Type { - case tea.MouseWheelUp: - m.scrollBy(3) - case tea.MouseWheelDown: - m.scrollBy(-3) - } - return m, nil - } - switch msg.Type { - case tea.MouseWheelUp: - m.scrollBy(3) - case tea.MouseWheelDown: - m.scrollBy(-3) - case tea.MouseLeft: - switch msg.Action { - case tea.MouseActionPress: - m.startTranscriptSelection(msg) - case tea.MouseActionMotion: - m.dragTranscriptSelection(msg) - default: - if msg.Action == 0 { - m.startTranscriptSelection(msg) - } - } - case tea.MouseMotion: - m.dragTranscriptSelection(msg) - case tea.MouseRelease: - return m, m.finishTranscriptSelection(msg) - } - return m, nil -} - -func (m chatModel) mouseInTranscript(msg tea.MouseMsg) bool { - top, height := m.transcriptLayout() - if msg.X < 0 || msg.X >= m.viewWidth() { - return false - } - return msg.Y >= top && msg.Y < top+height -} - -func (m chatModel) mouseTranscriptPoint(msg tea.MouseMsg) chatSelectionPoint { - top, height := m.transcriptLayout() - visibleY := clamp(msg.Y-top, 0, max(0, height-1)) - line := m.visibleTranscriptStartLine(m.viewWidth(), height) + visibleY - col := max(0, msg.X) - return chatSelectionPoint{line: line, col: col} -} - -func (m *chatModel) startTranscriptSelection(msg tea.MouseMsg) { - startChatSelection(&m.selection, msg, m.mouseInTranscript, m.mouseTranscriptPoint) -} - -func (m *chatModel) dragTranscriptSelection(msg tea.MouseMsg) { - dragChatSelection(&m.selection, msg, m.mouseTranscriptPoint, func(msg tea.MouseMsg) { - top, height := m.transcriptLayout() - if msg.Y <= top { - m.scrollBy(1) - } else if msg.Y >= top+height-1 { - m.scrollBy(-1) - } - }) -} - -func (m *chatModel) finishTranscriptSelection(msg tea.MouseMsg) tea.Cmd { - finishChatSelection(&m.selection, msg, m.mouseTranscriptPoint) - if chatRuntimeGOOS != "windows" || !m.selection.active { - return nil - } - selected := m.selectedTranscriptText(m.viewWidth()) - if strings.TrimSpace(selected) == "" { - m.selection = chatSelection{} - return nil - } - m.selection = chatSelection{} - m.status = "copied" - return copyTextCmd(m.ctx, selected) -} - -func (m chatModel) updateKey(msg tea.KeyMsg) (tea.Model, tea.Cmd) { - if m.promptDebug != nil { - return m.updatePromptDebug(msg) - } - if msg.Type == tea.KeyCtrlO { - m.toggleInlineTranscriptDetails() - m.disarmQuit() - m.disarmEsc() - return m.withFlowTranscriptRepaint(nil) - } - if msg.Type == tea.KeyShiftTab { - return m.togglePermissionMode() - } - if m.approvalPrompt != nil { - return m.updateApprovalPrompt(msg) - } - if m.cloudAuthPrompt != nil { - return m.updateCloudAuthPrompt(msg) - } - if m.modelPicker != nil { - return m.updateModelPicker(msg) - } - if m.thinkPicker != nil { - return m.updateThinkPicker(msg) - } - if msg.Type != tea.KeyCtrlC && msg.Type != tea.KeyCtrlD { - m.disarmQuit() - } - if msg.Type != tea.KeyEsc { - m.disarmEsc() - } - - switch msg.Type { - case tea.KeyCtrlC: - return m.updateCtrlC() - case tea.KeyCtrlD: - return m.updateCtrlD() - case tea.KeyEsc: - return m.updateEsc() - case tea.KeyEnter: - if msg.Alt { - m.insertInputNewline() - return m, nil - } - if m.applySlashCompletion() { - return m, nil - } - if m.applyMentionCompletion() { - return m, nil - } - return m.handleSubmit() - case tea.KeyCtrlJ: - m.insertInputNewline() - return m, nil - case tea.KeyCtrlG: - return m, nil - case tea.KeyCtrlP: - return m.updateUpKey() - case tea.KeyCtrlN: - return m.updateDownKey() - case tea.KeyUp: - return m.updateUpKey() - case tea.KeyDown: - return m.updateDownKey() - case tea.KeyLeft: - if msg.Alt { - m.moveInputCursorWord(-1) - } else { - m.moveInputCursorHorizontal(-1) - } - case tea.KeyRight: - if msg.Alt { - m.moveInputCursorWord(1) - } else { - m.moveInputCursorHorizontal(1) - } - case tea.KeyCtrlA: - m.moveInputCursorToLineStart() - case tea.KeyCtrlE: - m.moveInputCursorToLineEnd() - case tea.KeyCtrlB: - m.moveInputCursorHorizontal(-1) - case tea.KeyCtrlF: - m.moveInputCursorHorizontal(1) - case tea.KeyCtrlLeft: - m.moveInputCursorWord(-1) - case tea.KeyCtrlRight: - m.moveInputCursorWord(1) - case tea.KeyPgUp: - m.scrollBy(max(1, m.transcriptHeight()-1)) - case tea.KeyPgDown: - m.scrollBy(-max(1, m.transcriptHeight()-1)) - case tea.KeyHome: - m.moveInputCursorToLineStart() - case tea.KeyEnd: - m.moveInputCursorToLineEnd() - case tea.KeyCtrlHome: - m.scroll = m.maxScroll() - case tea.KeyCtrlEnd: - m.scroll = 0 - case tea.KeyBackspace, tea.KeyCtrlH: - m.resetPromptHistoryCursor() - if msg.Alt { - m.deleteInputWordBackward() - } else { - m.deleteInputBackward() - } - case tea.KeyCtrlW: - m.resetPromptHistoryCursor() - m.deleteInputWordBackward() - case tea.KeyCtrlU: - m.resetPromptHistoryCursor() - m.clearInput() - case tea.KeyCtrlK: - m.resetPromptHistoryCursor() - m.deleteInputForward() - case tea.KeyDelete: - m.resetPromptHistoryCursor() - m.deleteInputForward() - case tea.KeyTab: - m.applyCompletion() - case tea.KeySpace: - m.insertInputRunes([]rune{' '}) - case tea.KeyRunes: - if !msg.Alt || !m.handleInputAltRunes(msg.Runes) { - m.insertInputRunesFromKey(msg.Runes, msg.Paste) - } - } - return m, nil -} - -func (m *chatModel) toggleInlineTranscriptDetails() { - m.toolOutputMode = true - m.toolOutputOpen = !m.toolOutputOpen - m.applyToolOutputMode() - m.thinkingDetailsOpen = !m.thinkingDetailsOpen - m.applyThinkingDetails() - m.selection = chatSelection{} - m.scroll = 0 -} - -func (m chatModel) updateCtrlC() (tea.Model, tea.Cmd) { - if len(m.input) > 0 { - m.clearInput() - m.resetPromptHistoryCursor() - m.disarmQuit() - m.status = "ready" - return m, nil - } - if (m.running || m.compacting) && m.cancel != nil { - m.cancel() - m.disarmQuit() - m.status = "canceling" - return m, nil - } - if !m.quitArmed || m.quitArmedKey != "ctrl+c" { - m.armQuit("ctrl+c", "press ctrl+c again to quit") - return m, nil - } - m.quitting = true - return m, m.quitCmd() -} - -func (m chatModel) updateCtrlD() (tea.Model, tea.Cmd) { - if len(m.input) > 0 { - return m, nil - } - if !m.quitArmed || m.quitArmedKey != "ctrl+d" { - m.armQuit("ctrl+d", "press ctrl+d again to quit") - return m, nil - } - m.quitting = true - if (m.running || m.compacting) && m.cancel != nil { - m.cancel() - } - return m, m.quitCmd() -} - -func (m chatModel) quitCmd() tea.Cmd { - return tea.Quit -} - -func (m *chatModel) armQuit(key, status string) { - m.quitArmed = true - m.quitArmedKey = key - m.status = status -} - -func (m chatModel) updateEsc() (tea.Model, tea.Cmd) { - if (m.running || m.compacting) && m.cancel != nil { - m.cancel() - m.disarmQuit() - m.escArmed = false - m.status = "canceling" - return m, nil - } - if !m.escArmed { - m.escArmed = true - switch { - case len(m.input) > 0: - m.status = "press esc again to clear input" - default: - m.status = "ready" - } - return m, nil - } - - m.escArmed = false - cleared := false - if len(m.input) > 0 { - m.clearInput() - m.resetPromptHistoryCursor() - cleared = true - } - if cleared { - m.status = "ready" - } else { - m.status = "ready" - } - return m, nil -} - -func (m *chatModel) clearInput() { - m.input = nil - m.inputCursor = 0 - m.inputCursorSet = false - m.inputAttachments = nil - m.inputPastedTexts = nil - m.complete = 0 -} - -func (m chatModel) updateUpKey() (tea.Model, tea.Cmd) { - if m.promptActive && m.movePromptHistory(-1) { - return m, nil - } - if m.moveCompletion(-1) { - return m, nil - } - if m.moveInputCursorVertical(-1) { - return m, nil - } - if slices.Contains(m.input, '\n') { - return m, nil - } - m.movePromptHistory(-1) - return m, nil -} - -func (m chatModel) updateDownKey() (tea.Model, tea.Cmd) { - if m.promptActive && m.movePromptHistory(1) { - return m, nil - } - if m.moveCompletion(1) { - return m, nil - } - if m.moveInputCursorVertical(1) { - return m, nil - } - if slices.Contains(m.input, '\n') { - return m, nil - } - m.movePromptHistory(1) - return m, nil -} - -func (m chatModel) View() string { - if m.quitting { - return "" - } - - width, height := m.viewSize() - - if m.promptDebug != nil { - return m.renderPromptDebug(width, height) - } - if m.modelPicker != nil && m.openModelOnInit { - return m.renderModelPicker(width) - } - if m.cloudAuthPrompt != nil { - return m.renderCloudAuthPrompt(width) - } - if m.thinkPicker != nil { - return m.renderThinkPicker(width) - } - return m.flowView(width) -} - -func (m chatModel) flowView(width int) string { - allTranscriptLines := m.transcriptLines(width) - bottomLines := m.bottomLines(width, 0) - bottomGap := transcriptInputGap(0, len(bottomLines), len(allTranscriptLines)) - - printed := clamp(m.flowPrintedLines, 0, len(allTranscriptLines)) - lines := slices.Clone(allTranscriptLines[printed:]) - for range bottomGap { - lines = append(lines, "") - } - lines = append(lines, bottomLines...) - return strings.Join(lines, "\n") -} - -func (m chatModel) withFlowTranscriptFlush(cmd tea.Cmd) (tea.Model, tea.Cmd) { - next, printCmd := m.flowTranscriptFlushCmd() - return next, tea.Batch(printCmd, cmd) -} - -func (m chatModel) withFlowTranscriptRepaint(cmd tea.Cmd) (tea.Model, tea.Cmd) { - if m.flowPrintedLines == 0 { - return m.withFlowTranscriptFlush(cmd) - } - m.flowPrintedLines = 0 - next, printCmd := m.flowTranscriptFlushCmd() - return next, tea.Sequence(tea.ClearScreen, printCmd, cmd) -} - -func (m chatModel) withFlowTranscriptRefreshAfter(before []string, printed int, cmd tea.Cmd) (tea.Model, tea.Cmd) { - next, printCmd := m.flowTranscriptRefreshAfterCmd(before, printed) - return next, tea.Batch(printCmd, cmd) -} - -func (m chatModel) flowTranscriptRefreshAfterCmd(before []string, printed int) (chatModel, tea.Cmd) { - width := m.viewWidth() - after := m.transcriptLines(width) - start := flowTranscriptChangedPrefixStart(before, after, printed) - if start < 0 { - return m.flowTranscriptFlushCmd() - } - - printed = clamp(printed, 0, len(before)) - start = clamp(start, 0, len(after)) - rewind := printed - start - if rewind <= 0 { - return m.flowTranscriptFlushCmd() - } - - flushCount := m.flowTranscriptFlushCount(after, width) - flushCount = clamp(flushCount, start, len(after)) - rewriteLines := slices.Clone(after[start:flushCount]) - m.flowPrintedLines = flushCount - return m, tea.Printf("%s", flowTranscriptRewriteSequence(rewind, rewriteLines)) -} - -func flowTranscriptChangedPrefixStart(before, after []string, printed int) int { - printed = clamp(printed, 0, len(before)) - if printed == 0 { - return -1 - } - limit := min(printed, len(after)) - for i := range limit { - if before[i] != after[i] { - return i - } - } - if len(after) < printed { - return len(after) - } - return -1 -} - -func flowTranscriptRewriteSequence(rewind int, lines []string) string { - var b strings.Builder - if rewind > 0 { - fmt.Fprintf(&b, "\x1b[%dA", rewind) - } - b.WriteString("\r\x1b[J") - b.WriteString(strings.Join(lines, "\n")) - return b.String() -} - -func (m chatModel) flowTranscriptFlushCmd() (chatModel, tea.Cmd) { - if m.promptDebug != nil || m.modelPicker != nil || m.thinkPicker != nil || m.cloudAuthPrompt != nil { - return m, nil - } - width := m.viewWidth() - lines := m.transcriptLines(width) - if len(lines) == 0 { - m.flowPrintedLines = 0 - return m, nil - } - m.flowPrintedLines = clamp(m.flowPrintedLines, 0, len(lines)) - flushCount := m.flowTranscriptFlushCount(lines, width) - if flushCount <= m.flowPrintedLines { - return m, nil - } - pending := strings.Join(lines[m.flowPrintedLines:flushCount], "\n") - m.flowPrintedLines = flushCount - return m, tea.Println(pending) -} - -func (m chatModel) flowTranscriptFlushCount(lines []string, width int) int { - holdFrom := m.flowTranscriptHoldEntryIndex() - if holdFrom < 0 { - return len(lines) - } - clone := m - clone.entries = slices.Clone(m.entries[:holdFrom]) - clone.selection = chatSelection{} - return len(clone.transcriptLines(width)) -} - -func (m chatModel) flowTranscriptHoldEntryIndex() int { - if !m.running && !m.compacting { - return -1 - } - if len(m.entries) == 0 { - return -1 - } - index := len(m.entries) - 1 - entry := m.entries[index] - switch entry.role { - case "assistant": - if entry.content != "" { - return index - } - case "tool": - if isToolActiveStatus(entry.status) { - return index - } - case "thinking": - if entry.status == "running" { - return index - } - } - return -1 -} - -func (m chatModel) emptyChatHint() string { - if len(chatEmptyPrompts) == 0 { - return "Try asking the agent to inspect files, run tools, or explain a repo." - } - h := fnv.New32a() - _, _ = h.Write([]byte(m.chatID)) - prompt := chatEmptyPrompts[int(h.Sum32())%len(chatEmptyPrompts)] - return prompt -} - -func (m chatModel) headerLines() []string { - return nil -} - -func (m *chatModel) resetChat(status string) (tea.Model, tea.Cmd) { - m.messages = nil - m.liveMessages = nil - m.entries = nil - m.inputAttachments = nil - m.inputPastedTexts = nil - m.modelPicker = nil - m.modelPickerModels = nil - m.approvalController = nil - m.resetApprovalState() - m.nextImageID = 0 - m.nextAudioID = 0 - m.nextPastedTextID = 1 - m.resetPromptHistoryCursor() - m.resetWorkingDir() - m.opts.AllowAllTools = m.defaultAllowAll - m.permissionNotice = "" - m.thinking = false - m.thinkingTokens = 0 - m.contextTokens = 0 - m.contextEstimate = true - m.scroll = 0 - m.flowPrintedLines = 0 - m.status = status - return *m, tea.ClearScreen -} - -func (m *chatModel) resetWorkingDir() { - if m.opts.RootDir != "" { - m.workingDir = m.opts.RootDir - return - } - m.workingDir = m.opts.WorkingDir -} - -func (m *chatModel) startRun(input string) (tea.Model, tea.Cmd) { - displayInput, message, err := m.userMessageFromInput(input, input) - if err != nil { - m.entries = append(m.entries, newChatEntry(chatEntry{role: "error", content: err.Error(), err: err.Error()})) - m.status = "error" - return *m, nil - } - return m.startRunWithMessages(displayInput, message.Content, []api.Message{message}, "", "") -} - -func (m *chatModel) userMessageFromInput(displayInput, userInput string) (string, api.Message, error) { - content := m.expandPastedTextPlaceholders(userInput) - images := slices.Clone(m.opts.Images) - preloaded := len(images) - placeholderAttachments := m.activeInputAttachmentsFor(userInput) - for _, attachment := range placeholderAttachments { - images = append(images, attachment.data) - } - extracted := 0 - - if m.opts.MultiModal { - cleaned, files, err := filedata.ExtractWithFiles(content) - if err != nil { - return "", api.Message{}, err - } - content = cleaned - extracted = len(files) - for _, file := range files { - images = append(images, file.Data) - } - } - m.opts.Images = nil - m.inputAttachments = nil - m.inputPastedTexts = nil - - if len(images) > 0 { - base := displayInput - if extracted > 0 { - base = content - } - if len(placeholderAttachments) > 0 && extracted == 0 && preloaded == 0 { - displayInput = strings.TrimSpace(content) - } else { - displayInput = chatDisplayInputWithAttachments(base, len(images)) - } - } - - return displayInput, api.Message{Role: "user", Content: content, Images: images}, nil -} - -func chatDisplayInputWithAttachments(input string, count int) string { - note := fmt.Sprintf("[attached %d file%s]", count, pluralSuffix(count)) - input = strings.TrimSpace(input) - if input == "" { - return note - } - return input + "\n\n" + note -} - -func pluralSuffix(count int) string { - if count == 1 { - return "" - } - return "s" -} - -func (m *chatModel) startSkillRun(name, prompt string) (tea.Model, tea.Cmd) { - name = strings.TrimSpace(name) - prompt = strings.TrimSpace(prompt) - // The skill instructions are delivered via the synthetic tool result that - // follows this user turn, so this message orients the model rather than - // re-requesting the skill. When a prompt is supplied it becomes the task. - content := prompt - if content == "" { - content = "The " + name + " skill is loaded; follow its instructions for this request." - } - message := api.Message{Role: "user", Content: content} - displayInput := "/" + name - if prompt != "" { - displayInput = "/" + name + " " + prompt - } - return m.startRunWithMessages(displayInput, "", []api.Message{message}, "", name) -} - -func (m *chatModel) startRunWithMessages(displayInput, historyInput string, newMessages []api.Message, extraSystemPrompt, skillName string) (tea.Model, tea.Cmd) { - m.addPromptHistory(historyInput) - m.entries = append(m.entries, newChatEntry(chatEntry{role: "user", content: displayInput})) - if len(newMessages) > 1 { - m.entries = append(m.entries, entriesFromMessages(newMessages[1:])...) - } - m.running = true - m.awaitingModel = true - m.status = "running" - m.spinner = 0 - m.scroll = 0 - m.detectedToolCalls = nil - m.thinking = false - m.thinkingTokens = 0 - m.eventErrorRendered = false - systemPrompt := m.systemPrompt(extraSystemPrompt) - m.liveMessages = append(slices.Clone(m.messages), newMessages...) - m.contextTokens = m.estimatePromptTokens(m.liveMessages, systemPrompt) - m.contextEstimate = true - - runCtx, cancel := context.WithCancel(m.ctx) - m.cancel = cancel - events := make(chan tea.Msg, 128) - m.events = events - m.approvalController = newChatApprovalController(events, m.ensureApprovalState()) - - var newMessagesPersisted bool - eventSinks := []coreagent.EventSink{chatEventSink{ctx: runCtx, ch: events, newMessagesPersisted: &newMessagesPersisted}} - eventSinks = append(eventSinks, m.opts.EventSinks...) - - session := &coreagent.Session{ - Client: m.opts.Client, - EventSinks: eventSinks, - Tools: m.opts.Tools, - Skills: m.opts.Skills, - DisableTools: m.opts.ToolsDisabled, - ApprovalPrompter: m.approvalPrompterForRun(m.approvalController), - ApprovalState: m.ensureApprovalState(), - WorkingDir: m.currentWorkingDir(), - Compactor: m.opts.Compactor, - } - opts := coreagent.RunOptions{ - ChatID: m.chatID, - Model: m.opts.Model, - SystemPrompt: systemPrompt, - Messages: slices.Clone(m.messages), - NewMessages: slices.Clone(newMessages), - Format: m.opts.Format, - Options: m.opts.Options, - Think: m.opts.Think, - KeepAlive: m.opts.KeepAlive, - SkillName: skillName, - } - - persistedMessages := make([]api.Message, 0, len(m.messages)+len(newMessages)) - persistedMessages = append(persistedMessages, slices.Clone(m.messages)...) - persistedMessages = append(persistedMessages, slices.Clone(newMessages)...) - go func() { - defer close(events) - result, err := session.Run(runCtx, opts) - select { - case events <- chatRunDoneMsg{result: result, err: err, newMessagesPersisted: newMessagesPersisted, persistedMessages: persistedMessages}: - case <-runCtx.Done(): - } - }() - - tickCmd := m.scheduleTick() - flushModel, flushCmd := m.flowTranscriptFlushCmd() - *m = flushModel - return *m, tea.Batch(flushCmd, waitForChatMsg(events), tickCmd) -} - -func (m *chatModel) finishLiveMessagesForStoppedRun(promote bool, persistedMessages []api.Message) { - if len(m.liveMessages) == 0 { - return - } - if promote { - if len(persistedMessages) > 0 { - m.messages = slices.Clone(persistedMessages) - } else if !messagesHavePendingToolCalls(m.liveMessages) { - m.messages = slices.Clone(m.liveMessages) - } - } - m.liveMessages = nil - m.contextTokens = m.estimatePromptTokens(m.messages, "") - m.contextEstimate = true -} - -func messagesHavePendingToolCalls(messages []api.Message) bool { - pending := map[string]struct{}{} - for _, msg := range messages { - if msg.Role == "assistant" { - for _, call := range msg.ToolCalls { - pending[call.ID] = struct{}{} - } - } - if msg.Role == "tool" && msg.ToolCallID != "" { - delete(pending, msg.ToolCallID) - } - } - return len(pending) > 0 -} - -func (m *chatModel) disarmQuit() { - if !m.quitArmed { - return - } - m.quitArmed = false - m.quitArmedKey = "" - if strings.HasPrefix(m.status, "press ctrl+") && strings.HasSuffix(m.status, "again to quit") { - m.status = "ready" - } -} - -func (m *chatModel) disarmEsc() { - if !m.escArmed { - return - } - m.escArmed = false - if strings.HasPrefix(m.status, "press esc again") { - m.status = "ready" - } -} - -func (m chatModel) canEditInput() bool { - return m.promptDebug == nil && m.approvalPrompt == nil && m.cloudAuthPrompt == nil && m.modelPicker == nil && m.thinkPicker == nil -} - -func isChatContextCanceledError(err error) bool { - return err != nil && (errors.Is(err, context.Canceled) || strings.Contains(err.Error(), "context canceled")) -} diff --git a/cmd/tui/chat/clipboard.go b/cmd/tui/chat/clipboard.go deleted file mode 100644 index 337cd37f6b3..00000000000 --- a/cmd/tui/chat/clipboard.go +++ /dev/null @@ -1,66 +0,0 @@ -package chat - -import ( - "context" - "errors" - "fmt" - "os/exec" - "runtime" - "strings" - - tea "github.com/charmbracelet/bubbletea" -) - -type chatClipboardErrorMsg struct { - err error -} - -var writeClipboard = writeSystemClipboard - -func copyTextCmd(ctx context.Context, text string) tea.Cmd { - return func() tea.Msg { - if err := writeClipboard(ctx, text); err != nil { - return chatClipboardErrorMsg{err: err} - } - return nil - } -} - -func writeSystemClipboard(ctx context.Context, text string) error { - if ctx == nil { - ctx = context.Background() - } - switch runtime.GOOS { - case "darwin": - return runClipboardCommand(ctx, text, "pbcopy") - case "windows": - return runClipboardCommand(ctx, text, "clip") - default: - for _, candidate := range []struct { - name string - args []string - }{ - {name: "wl-copy"}, - {name: "xclip", args: []string{"-selection", "clipboard"}}, - {name: "xsel", args: []string{"--clipboard", "--input"}}, - } { - if _, err := exec.LookPath(candidate.name); err != nil { - continue - } - return runClipboardCommand(ctx, text, candidate.name, candidate.args...) - } - return errors.New("no clipboard command found") - } -} - -func runClipboardCommand(ctx context.Context, text, name string, args ...string) error { - cmd := exec.CommandContext(ctx, name, args...) - cmd.Stdin = strings.NewReader(text) - if output, err := cmd.CombinedOutput(); err != nil { - if len(output) > 0 { - return fmt.Errorf("%s: %w: %s", name, err, strings.TrimSpace(string(output))) - } - return fmt.Errorf("%s: %w", name, err) - } - return nil -} diff --git a/cmd/tui/chat/cloudauth.go b/cmd/tui/chat/cloudauth.go deleted file mode 100644 index adf74cab858..00000000000 --- a/cmd/tui/chat/cloudauth.go +++ /dev/null @@ -1,434 +0,0 @@ -package chat - -import ( - "context" - "errors" - "fmt" - "net/http" - "strings" - "time" - - tea "github.com/charmbracelet/bubbletea" - "github.com/ollama/ollama/api" - "github.com/ollama/ollama/cmd/launch" - "github.com/ollama/ollama/internal/modelref" -) - -type cloudAuthKind string - -const ( - cloudAuthSignIn cloudAuthKind = "signin" - cloudAuthUpgrade cloudAuthKind = "upgrade" - cloudAuthChecking cloudAuthKind = "checking" -) - -const cloudPlanVerificationUnavailable = "Could not verify Ollama plan. Try again in a moment or use a local model." - -// Sign-in/upgrade verification polling bounds. While the check is healthy but -// the user hasn't signed in yet, polling stays prompt so completion is detected -// quickly. When the check itself fails, polling backs off so a down server -// isn't hammered, and gives up after maxPollFailures consecutive errors (or -// pollHardCap elapsed) so the user isn't stuck on a spinner with no recourse -// beyond Esc. -const ( - maxPollFailures = 6 - pollBackoffBase = 3 * time.Second - pollBackoffCap = 30 * time.Second - pollHardCap = 2 * time.Minute -) - -// cloudAuthPrompt is an inline modal that handles sign-in and plan-upgrade -// flows when a user selects a cloud model from the picker. -type cloudAuthPrompt struct { - modelName string - requiredPlan string - signInURL string - upgradeURL string - kind cloudAuthKind - spinner int - openNow bool - polling bool - // pollStarted tracks when sign-in/upgrade verification polling began, for - // the hard-cap timeout. Lazily set on the first poll response. - pollStarted time.Time - // pollFailures counts consecutive verification-check errors; once it - // reaches maxPollFailures the modal gives up and surfaces an error. - pollFailures int - // pollErr holds the last verification error, rendered while retrying. - pollErr string -} - -type cloudAuthCheckMsg struct { - err error - signInURL string -} - -type cloudModelPreflightMsg struct { - model string - err error - signInURL string -} - -type cloudAuthTickMsg struct{} - -type cloudAuthPollMsg struct { - done bool - err error -} - -func checkCloudModelCmd(ctx context.Context, check func(context.Context, string, string) error, model, requiredPlan string) tea.Cmd { - if check == nil { - return nil - } - return func() tea.Msg { - if ctx == nil { - ctx = context.Background() - } - err := check(ctx, model, requiredPlan) - var signInURL string - if err != nil { - var authErr api.AuthorizationError - if errors.As(err, &authErr) && authErr.SigninURL != "" { - signInURL = authErr.SigninURL - } - } - return cloudAuthCheckMsg{err: err, signInURL: signInURL} - } -} - -func cloudModelPreflightCmd(ctx context.Context, opts Options, modelName, requiredPlan string) tea.Cmd { - modelName = strings.TrimSpace(modelName) - if opts.CheckCloudModel == nil || modelName == "" || !modelref.HasExplicitCloudSource(modelName) { - return nil - } - return func() tea.Msg { - if ctx == nil { - ctx = context.Background() - } - plan := strings.TrimSpace(requiredPlan) - if plan == "" && opts.ModelOptions != nil { - models, err := opts.ModelOptions(ctx) - if err == nil { - for _, model := range models { - if strings.EqualFold(strings.TrimSpace(model.Name), modelName) { - plan = strings.TrimSpace(model.RequiredPlan) - break - } - } - } - } - err := opts.CheckCloudModel(ctx, modelName, plan) - return cloudModelPreflightMsg{ - model: modelName, - err: err, - signInURL: cloudAuthSignInURL(err), - } - } -} - -func cloudAuthSignInURL(err error) string { - if err == nil { - return "" - } - var authErr api.AuthorizationError - if errors.As(err, &authErr) && (authErr.StatusCode == http.StatusUnauthorized || authErr.SigninURL != "") { - return authErr.SigninURL - } - return "" -} - -func cloudAuthTickCmd() tea.Cmd { - return tea.Tick(200*time.Millisecond, func(t time.Time) tea.Msg { - return cloudAuthTickMsg{} - }) -} - -func (m chatModel) updateCloudModelPreflight(msg cloudModelPreflightMsg) (tea.Model, tea.Cmd) { - if msg.model == "" || !strings.EqualFold(strings.TrimSpace(m.opts.Model), strings.TrimSpace(msg.model)) { - return m, nil - } - if msg.err == nil { - if m.status == cloudPlanVerificationUnavailable { - m.status = "ready" - } - return m, nil - } - if msg.signInURL != "" { - return m.startCloudAuthSignIn(msg.model, "", msg.signInURL) - } - m.status = cloudPlanVerificationUnavailable - return m, nil -} - -func pollCloudAuthCmd(ctx context.Context, poll func(context.Context) (string, bool, error), delay time.Duration) tea.Cmd { - if poll == nil { - return nil - } - return func() tea.Msg { - if ctx == nil { - ctx = context.Background() - } - // Back off before the next check when the previous one failed. Honor - // context cancellation so an abandoned modal doesn't block on the - // full delay. - if delay > 0 { - timer := time.NewTimer(delay) - defer timer.Stop() - select { - case <-ctx.Done(): - case <-timer.C: - } - } - pollCtx, cancel := context.WithTimeout(ctx, 3*time.Second) - defer cancel() - _, done, err := poll(pollCtx) - return cloudAuthPollMsg{done: done, err: err} - } -} - -func (m *chatModel) startCloudAuthSignIn(modelName, requiredPlan, signInURL string) (tea.Model, tea.Cmd) { - // When no sign-in URL is available yet, show the "checking" state while - // we verify the plan, rather than rendering a blank "Navigate to:" URL. - kind := cloudAuthSignIn - if signInURL == "" { - kind = cloudAuthChecking - } - m.cloudAuthPrompt = &cloudAuthPrompt{ - modelName: modelName, - requiredPlan: requiredPlan, - kind: kind, - signInURL: signInURL, - polling: true, - } - m.status = "cloud-auth" - m.modelPicker = nil - m.modelPickerModels = nil - if m.opts.OpenBrowser != nil && signInURL != "" { - m.opts.OpenBrowser(signInURL) - } - if signInURL == "" { - return m, checkCloudModelCmd(m.ctx, m.opts.CheckCloudModel, modelName, requiredPlan) - } - return m, tea.Batch(cloudAuthTickCmd(), pollCloudAuthCmd(m.ctx, m.opts.PollCloudAuth, 0)) -} - -func (m *chatModel) startCloudAuthUpgrade(modelName, requiredPlan string) (tea.Model, tea.Cmd) { - m.cloudAuthPrompt = &cloudAuthPrompt{ - modelName: modelName, - requiredPlan: requiredPlan, - kind: cloudAuthUpgrade, - upgradeURL: launch.DefaultUpgradeURL, - openNow: true, - } - m.status = "cloud-auth" - m.modelPicker = nil - m.modelPickerModels = nil - return m, nil -} - -func (m chatModel) updateCloudAuthPrompt(msg tea.Msg) (tea.Model, tea.Cmd) { - switch msg := msg.(type) { - case cloudAuthCheckMsg: - if msg.err == nil { - // Auth passed — apply the pending model. - return m.completeCloudAuth() - } - // Determine if sign-in or upgrade is needed. - if msg.signInURL != "" { - m.cloudAuthPrompt.kind = cloudAuthSignIn - m.cloudAuthPrompt.signInURL = msg.signInURL - m.cloudAuthPrompt.polling = true - if m.opts.OpenBrowser != nil { - m.opts.OpenBrowser(msg.signInURL) - } - return m, tea.Batch(cloudAuthTickCmd(), pollCloudAuthCmd(m.ctx, m.opts.PollCloudAuth, 0)) - } - // Could be a plan upgrade error or unknown error. - m.cloudAuthPrompt = nil - m.openModelOnInit = false - m.status = "ready" - m.entries = append(m.entries, newChatEntry(chatEntry{role: "error", content: fmt.Sprintf("Could not switch model: %v", msg.err), err: msg.err.Error()})) - return m, nil - - case cloudAuthTickMsg: - if m.cloudAuthPrompt == nil { - return m, nil - } - m.cloudAuthPrompt.spinner++ - return m, cloudAuthTickCmd() - - case cloudAuthPollMsg: - if m.cloudAuthPrompt == nil { - return m, nil - } - if msg.done { - // Signed in — re-check auth to see if plan is satisfied. - m.cloudAuthPrompt.polling = false - m.cloudAuthPrompt.pollFailures = 0 - m.cloudAuthPrompt.pollErr = "" - return m, checkCloudModelCmd(m.ctx, m.opts.CheckCloudModel, m.cloudAuthPrompt.modelName, m.cloudAuthPrompt.requiredPlan) - } - // Lazily mark the start of the polling window on the first response. - if m.cloudAuthPrompt.pollStarted.IsZero() { - m.cloudAuthPrompt.pollStarted = time.Now() - } - // Hard cap: give up if verification drags on too long for any reason. - if time.Since(m.cloudAuthPrompt.pollStarted) > pollHardCap { - return m.failCloudAuthPoll(errors.New("sign-in is taking longer than expected; check your connection and try again")) - } - if msg.err != nil { - // The verification check itself failed (network down, server 5xx). - // Back off and retry, but give up after a handful of consecutive - // failures so the user isn't stuck on a spinner with no signal. - m.cloudAuthPrompt.pollFailures++ - m.cloudAuthPrompt.pollErr = msg.err.Error() - if m.cloudAuthPrompt.pollFailures >= maxPollFailures { - return m.failCloudAuthPoll(fmt.Errorf("couldn't verify sign-in: %w", msg.err)) - } - delay := pollBackoffCap - if d := pollBackoffBase << (m.cloudAuthPrompt.pollFailures - 1); d < pollBackoffCap { - delay = d - } - return m, pollCloudAuthCmd(m.ctx, m.opts.PollCloudAuth, delay) - } - // Healthy but not signed in yet — keep polling promptly so sign-in - // completion is detected without added latency. - m.cloudAuthPrompt.pollFailures = 0 - m.cloudAuthPrompt.pollErr = "" - return m, pollCloudAuthCmd(m.ctx, m.opts.PollCloudAuth, 0) - - case tea.KeyMsg: - if msg.Type == tea.KeyEsc || msg.Type == tea.KeyCtrlC { - m.cloudAuthPrompt = nil - m.pendingModel = "" - m.openModelOnInit = false - m.status = "ready" - return m, nil - } - if m.cloudAuthPrompt.kind == cloudAuthUpgrade && !m.cloudAuthPrompt.polling { - switch msg.Type { - case tea.KeyLeft, tea.KeyRight, tea.KeyTab: - m.cloudAuthPrompt.openNow = !m.cloudAuthPrompt.openNow - case tea.KeyEnter: - if m.cloudAuthPrompt.openNow { - m.cloudAuthPrompt.polling = true - if m.opts.OpenBrowser != nil && m.cloudAuthPrompt.upgradeURL != "" { - m.opts.OpenBrowser(m.cloudAuthPrompt.upgradeURL) - } - return m, tea.Batch(cloudAuthTickCmd(), pollCloudAuthCmd(m.ctx, m.opts.PollCloudAuth, 0)) - } - m.cloudAuthPrompt = nil - m.pendingModel = "" - m.openModelOnInit = false - m.status = "ready" - return m, nil - } - } - } - - return m, nil -} - -// failCloudAuthPoll abandons the sign-in/upgrade verification modal, surfaces -// an error entry to the user, and returns to the ready state so they can -// re-pick a model and retry. -func (m chatModel) failCloudAuthPoll(err error) (tea.Model, tea.Cmd) { - m.cloudAuthPrompt = nil - m.pendingModel = "" - m.openModelOnInit = false - m.status = "ready" - m.entries = append(m.entries, newChatEntry(chatEntry{role: "error", content: fmt.Sprintf("Could not switch model: %v", err), err: err.Error()})) - return m, nil -} - -func (m chatModel) completeCloudAuth() (tea.Model, tea.Cmd) { - pending := m.cloudAuthPrompt.modelName - m.cloudAuthPrompt = nil - m.pendingModel = "" - m.modelPicker = nil - m.modelPickerModels = nil - m.openModelOnInit = false - m.status = "ready" - if err := m.applyModelSelection(pending, true); err != nil { - m.entries = append(m.entries, newChatEntry(chatEntry{role: "error", content: fmt.Sprintf("Could not switch model: %v", err), err: err.Error()})) - m.status = "error" - return m, nil - } - return m, m.startModelPreload(pending) -} - -func (m chatModel) renderCloudAuthPrompt(width int) string { - if m.cloudAuthPrompt == nil { - return "" - } - if width <= 0 { - width = 80 - } - - p := m.cloudAuthPrompt - spinnerFrames := []string{"⠋", "⠙", "⠹", "⠸", "⠼", "⠴", "⠦", "⠧", "⠇", "⠏"} - frame := spinnerFrames[p.spinner%len(spinnerFrames)] - - var b strings.Builder - - switch p.kind { - case cloudAuthChecking: - fmt.Fprintf(&b, "%s Checking %s...\n\n", frame, chatPickerSelectedStyle.Render(p.modelName)) - b.WriteString(chatPickerMetaStyle.Render("esc cancel")) - case cloudAuthSignIn: - fmt.Fprintf(&b, "To use %s, please sign in.\n\n", chatPickerSelectedStyle.Render(p.modelName)) - b.WriteString("Navigate to:\n") - urlWrap := chatPickerTextStyle - if width > 4 { - urlWrap = chatPickerTextStyle.Width(width - 4) - } - b.WriteString(urlWrap.Render(p.signInURL)) - b.WriteString("\n\n") - if p.pollErr != "" { - b.WriteString(chatPickerMetaStyle.Render(frame + " Couldn't verify sign-in: " + p.pollErr + " — retrying...")) - } else { - b.WriteString(chatPickerMetaStyle.Render(frame + " Waiting for sign in to complete...")) - } - b.WriteString("\n\n") - b.WriteString(chatPickerMetaStyle.Render("esc cancel")) - case cloudAuthUpgrade: - fmt.Fprintf(&b, "To use %s, upgrade your Ollama plan.\n\n", chatPickerSelectedStyle.Render(p.modelName)) - if !p.polling { - var yesBtn, noBtn string - if p.openNow { - yesBtn = chatPickerSelectedStyle.Render("› Yes ") - noBtn = chatPickerMetaStyle.Render(" No ") - } else { - yesBtn = chatPickerMetaStyle.Render(" Yes ") - noBtn = chatPickerSelectedStyle.Render("› No ") - } - b.WriteString("Open upgrade page now?\n") - b.WriteString(yesBtn + " " + noBtn) - b.WriteString("\n\n") - if !p.openNow { - b.WriteString("Or navigate to:\n") - urlWrap := chatPickerTextStyle - if width > 4 { - urlWrap = chatPickerTextStyle.Width(width - 4) - } - if u := p.upgradeURL; u != "" { - b.WriteString(urlWrap.Render(u)) - } else { - b.WriteString(urlWrap.Render(launch.DefaultUpgradeURL)) - } - b.WriteString("\n\n") - } - b.WriteString(chatPickerMetaStyle.Render("←/→ navigate • enter confirm • esc cancel")) - } else { - if p.pollErr != "" { - b.WriteString(chatPickerMetaStyle.Render(frame + " Couldn't verify upgrade: " + p.pollErr + " — retrying...")) - } else { - b.WriteString(chatPickerMetaStyle.Render(frame + " Waiting for upgrade to complete...")) - } - b.WriteString("\n\n") - b.WriteString(chatPickerMetaStyle.Render("esc cancel")) - } - } - - return b.String() -} diff --git a/cmd/tui/chat/cloudauth_test.go b/cmd/tui/chat/cloudauth_test.go deleted file mode 100644 index 74f85c86762..00000000000 --- a/cmd/tui/chat/cloudauth_test.go +++ /dev/null @@ -1,247 +0,0 @@ -package chat - -import ( - "context" - "errors" - "strings" - "testing" -) - -func TestCloudAuthTickDoesNotPoll(t *testing.T) { - polls := 0 - m := chatModel{ - cloudAuthPrompt: &cloudAuthPrompt{polling: true}, - opts: Options{ - PollCloudAuth: func(context.Context) (string, bool, error) { - polls++ - return "", false, nil - }, - }, - } - - updated, cmd := m.updateCloudAuthPrompt(cloudAuthTickMsg{}) - m = updated.(chatModel) - - if m.cloudAuthPrompt.spinner != 1 { - t.Fatalf("spinner = %d, want 1", m.cloudAuthPrompt.spinner) - } - if polls != 0 { - t.Fatalf("polls = %d, want 0 before running returned tick command", polls) - } - if cmd == nil { - t.Fatal("tick should schedule the next tick") - } - if _, ok := cmd().(cloudAuthTickMsg); !ok { - t.Fatal("tick should schedule another tick, not a poll") - } - if polls != 0 { - t.Fatalf("polls = %d, want 0 after running returned tick command", polls) - } -} - -func TestCloudAuthPollSchedulesNextPoll(t *testing.T) { - polls := 0 - m := chatModel{ - cloudAuthPrompt: &cloudAuthPrompt{polling: true}, - opts: Options{ - PollCloudAuth: func(context.Context) (string, bool, error) { - polls++ - return "", false, nil - }, - }, - } - - _, cmd := m.updateCloudAuthPrompt(cloudAuthPollMsg{}) - if cmd == nil { - t.Fatal("poll should schedule the next poll") - } - msg, ok := cmd().(cloudAuthPollMsg) - if !ok { - t.Fatal("poll should schedule another poll, not a tick") - } - if msg.done { - t.Fatal("poll should report not done") - } - if polls != 1 { - t.Fatalf("polls = %d, want 1", polls) - } -} - -func TestCloudModelPreflightFailureShowsPlanVerificationNotice(t *testing.T) { - m := chatModel{ - opts: Options{ - Model: "glm-5.2:cloud", - }, - } - - updated, cmd := m.updateCloudModelPreflight(cloudModelPreflightMsg{ - model: "glm-5.2:cloud", - err: errors.New("temporary network failure"), - }) - if cmd != nil { - t.Fatal("transient preflight failure should not start an auth modal") - } - m = updated.(chatModel) - - if got := m.status; got != cloudPlanVerificationUnavailable { - t.Fatalf("status = %q", got) - } - if m.cloudAuthPrompt != nil { - t.Fatalf("cloud auth prompt = %#v, want nil", m.cloudAuthPrompt) - } -} - -func TestCloudModelPreflightIgnoresStaleModel(t *testing.T) { - m := chatModel{ - opts: Options{ - Model: "glm-5.2:cloud", - }, - status: "ready", - } - - updated, _ := m.updateCloudModelPreflight(cloudModelPreflightMsg{ - model: "kimi-k2.7-code:cloud", - err: errors.New("temporary network failure"), - }) - m = updated.(chatModel) - - if got := m.status; got != "ready" { - t.Fatalf("status = %q, want unchanged", got) - } -} - -func TestCloudModelPreflightCommandChecksCloudModel(t *testing.T) { - var checkedModel, checkedPlan string - cmd := cloudModelPreflightCmd(context.Background(), Options{ - CheckCloudModel: func(_ context.Context, model, requiredPlan string) error { - checkedModel = model - checkedPlan = requiredPlan - return errors.New("temporary network failure") - }, - ModelOptions: func(context.Context) ([]ModelOption, error) { - return []ModelOption{{Name: "glm-5.2:cloud", RequiredPlan: "pro", Cloud: true}}, nil - }, - }, "glm-5.2:cloud", "") - if cmd == nil { - t.Fatal("cloud preflight command should be scheduled") - } - raw := cmd() - msg, ok := raw.(cloudModelPreflightMsg) - if !ok { - t.Fatalf("message = %T, want cloudModelPreflightMsg", raw) - } - if checkedModel != "glm-5.2:cloud" || checkedPlan != "pro" { - t.Fatalf("checked model/plan = %q/%q", checkedModel, checkedPlan) - } - if msg.model != "glm-5.2:cloud" || msg.err == nil || !strings.Contains(msg.err.Error(), "temporary") { - t.Fatalf("message = %#v", msg) - } -} - -func TestCloudAuthPollGivesUpAfterConsecutiveFailures(t *testing.T) { - pollErr := errors.New("whoami: connection refused") - m := chatModel{ - cloudAuthPrompt: &cloudAuthPrompt{polling: true, kind: cloudAuthSignIn}, - opts: Options{ - PollCloudAuth: func(context.Context) (string, bool, error) { - return "", false, pollErr - }, - }, - } - - // The first maxPollFailures-1 failures should keep retrying. - for i := 1; i < maxPollFailures; i++ { - updated, _ := m.updateCloudAuthPrompt(cloudAuthPollMsg{done: false, err: pollErr}) - m = updated.(chatModel) - if m.cloudAuthPrompt == nil { - t.Fatalf("failure %d: prompt cleared early", i) - } - if got := m.cloudAuthPrompt.pollFailures; got != i { - t.Fatalf("failure %d: pollFailures = %d, want %d", i, got, i) - } - if m.cloudAuthPrompt.pollErr != pollErr.Error() { - t.Fatalf("failure %d: pollErr = %q, want %q", i, m.cloudAuthPrompt.pollErr, pollErr.Error()) - } - } - - // The threshold failure gives up: prompt cleared, back to ready, error entry. - updated, cmd := m.updateCloudAuthPrompt(cloudAuthPollMsg{done: false, err: pollErr}) - m = updated.(chatModel) - if cmd != nil { - t.Fatalf("threshold failure should not reschedule, got cmd %T", cmd) - } - if m.cloudAuthPrompt != nil { - t.Fatalf("prompt = %#v, want nil after give-up", m.cloudAuthPrompt) - } - if m.status != "ready" { - t.Fatalf("status = %q, want ready", m.status) - } - if len(m.entries) == 0 { - t.Fatal("expected an error entry after give-up") - } - last := m.entries[len(m.entries)-1] - if last.role != "error" || !strings.Contains(last.content, "couldn't verify sign-in") { - t.Fatalf("last entry = %+v, want error containing sign-in failure", last) - } -} - -func TestCloudAuthPollResetsFailuresOnHealthyResponse(t *testing.T) { - pollErr := errors.New("whoami: timeout") - m := chatModel{ - cloudAuthPrompt: &cloudAuthPrompt{polling: true, kind: cloudAuthSignIn}, - opts: Options{ - PollCloudAuth: func(context.Context) (string, bool, error) { - return "", false, pollErr - }, - }, - } - - // Accumulate some failures without hitting the threshold. - for range maxPollFailures - 2 { - updated, _ := m.updateCloudAuthPrompt(cloudAuthPollMsg{done: false, err: pollErr}) - m = updated.(chatModel) - } - if got := m.cloudAuthPrompt.pollFailures; got != maxPollFailures-2 { - t.Fatalf("pollFailures = %d, want %d", got, maxPollFailures-2) - } - - // A healthy (no-error, not-done) response resets the streak so a later - // transient blip isn't counted against a recovered connection. - updated, _ := m.updateCloudAuthPrompt(cloudAuthPollMsg{done: false, err: nil}) - m = updated.(chatModel) - if m.cloudAuthPrompt == nil { - t.Fatal("healthy response should keep the prompt open") - } - if got := m.cloudAuthPrompt.pollFailures; got != 0 { - t.Fatalf("pollFailures = %d, want 0 after healthy response", got) - } - if m.cloudAuthPrompt.pollErr != "" { - t.Fatalf("pollErr = %q, want empty after healthy response", m.cloudAuthPrompt.pollErr) - } -} - -func TestCloudAuthPollCompletesAfterFailures(t *testing.T) { - pollErr := errors.New("whoami: timeout") - m := chatModel{ - cloudAuthPrompt: &cloudAuthPrompt{ - modelName: "glm-5.2:cloud", - polling: true, - kind: cloudAuthSignIn, - pollFailures: maxPollFailures - 1, - }, - opts: Options{ - CheckCloudModel: func(context.Context, string, string) error { return nil }, - PollCloudAuth: func(context.Context) (string, bool, error) { return "", false, pollErr }, - }, - } - - // A successful sign-in mid-retry should clear the failure state and re-check. - updated, _ := m.updateCloudAuthPrompt(cloudAuthPollMsg{done: true}) - m = updated.(chatModel) - if m.cloudAuthPrompt.polling { - t.Fatal("done should stop polling") - } - if m.cloudAuthPrompt.pollFailures != 0 || m.cloudAuthPrompt.pollErr != "" { - t.Fatalf("failure state not reset: failures=%d err=%q", m.cloudAuthPrompt.pollFailures, m.cloudAuthPrompt.pollErr) - } -} diff --git a/cmd/tui/chat/compaction.go b/cmd/tui/chat/compaction.go deleted file mode 100644 index 51e035b5fc1..00000000000 --- a/cmd/tui/chat/compaction.go +++ /dev/null @@ -1,101 +0,0 @@ -package chat - -import ( - "context" - "slices" - - tea "github.com/charmbracelet/bubbletea" - - coreagent "github.com/ollama/ollama/agent" - "github.com/ollama/ollama/api" -) - -func (m *chatModel) startManualCompaction() (tea.Model, tea.Cmd) { - if m.running || m.compacting { - m.entries = append(m.entries, newChatEntry(chatEntry{role: "system", content: "Wait for the current response to finish before compacting."})) - return *m, nil - } - m.refreshContextWindowTokens(m.opts.Model) - if m.opts.Compactor == nil { - m.entries = append(m.entries, newChatEntry(chatEntry{role: "system", content: coreagent.CompactionSkippedMessage("compaction is unavailable")})) - m.status = "compact skipped" - return *m, nil - } - - ctx := m.ctx - if ctx == nil { - ctx = context.Background() - } - runCtx, cancel := context.WithCancel(ctx) - compactor := m.opts.Compactor - events := make(chan tea.Msg, 128) - m.compacting = true - m.compactingTokens = 0 - m.cancel = cancel - m.compactEvents = events - m.status = "compacting" - messages := slices.Clone(m.messages) - var tools api.Tools - if m.opts.Tools != nil { - tools = m.opts.Tools.Tools() - } - req := coreagent.CompactionRequest{ - ChatID: m.chatID, - Model: m.opts.Model, - SystemPrompt: m.systemPrompt(""), - Messages: messages, - Tools: tools, - Format: m.opts.Format, - Options: m.opts.Options, - KeepAlive: m.opts.KeepAlive, - Force: true, - Progress: func(progress coreagent.CompactionProgress) { - select { - case events <- chatCompactProgressMsg{tokens: progress.Tokens}: - case <-runCtx.Done(): - } - }, - } - go func() { - defer close(events) - result, err := compactor.MaybeCompact(runCtx, req) - select { - case events <- chatCompactDoneMsg{result: result, err: err}: - case <-runCtx.Done(): - } - }() - tickCmd := m.scheduleTick() - return *m, tea.Batch(waitForChatMsg(events), tickCmd) -} - -func (m chatModel) finishManualCompaction(msg chatCompactDoneMsg) (tea.Model, tea.Cmd) { - wasCanceling := m.status == "canceling" - m.compacting = false - m.compactEvents = nil - m.cancel = nil - m.compactingTokens = 0 - if wasCanceling || isChatContextCanceledError(msg.err) { - m.status = "compact canceled" - return m.withFlowTranscriptFlush(nil) - } - if msg.err != nil { - m.entries = append(m.entries, newChatEntry(chatEntry{role: "system", content: coreagent.CompactionSkippedMessage(msg.err.Error())})) - m.status = "compact skipped" - return m.withFlowTranscriptFlush(nil) - } - if !msg.result.Compacted { - m.entries = append(m.entries, newChatEntry(chatEntry{role: "system", content: coreagent.CompactionSkippedMessage(msg.result.Reason)})) - m.status = "compact skipped" - return m.withFlowTranscriptFlush(nil) - } - - m.messages = msg.result.Messages - m.liveMessages = nil - m.entries = entriesFromMessages(m.messages) - m.contextTokens = m.estimatePromptTokens(m.messages, "") - m.contextEstimate = true - m.scroll = 0 - m.flowPrintedLines = 0 - m.status = "compacted" - return m.withFlowTranscriptFlush(nil) -} diff --git a/cmd/tui/chat/debug.go b/cmd/tui/chat/debug.go deleted file mode 100644 index c61fa1067bf..00000000000 --- a/cmd/tui/chat/debug.go +++ /dev/null @@ -1,565 +0,0 @@ -package chat - -import ( - "encoding/json" - "fmt" - "os" - "path/filepath" - "slices" - "strings" - - tea "github.com/charmbracelet/bubbletea" - "github.com/charmbracelet/lipgloss" - - coreagent "github.com/ollama/ollama/agent" - "github.com/ollama/ollama/api" -) - -type chatPromptDebug struct { - request api.ChatRequest - tokens int - scroll int - lines []string - linesWidth int -} - -const maxPromptDebugToolResultRunes = 400 - -func (m *chatModel) handleSaveCommand(args string) (tea.Model, tea.Cmd) { - filename, err := saveRequestFilename(args) - if err != nil { - return m.addDebugError(err) - } - raw, err := m.rawRequestJSON() - if err != nil { - return m.addDebugError(err) - } - - dir, err := m.debugWorkingDir() - if err != nil { - return m.addDebugError(err) - } - path := filepath.Join(dir, filename) - if err := os.WriteFile(path, []byte(raw+"\n"), 0o644); err != nil { - return m.addDebugError(err) - } - m.entries = append(m.entries, newSlashEntry(fmt.Sprintf("saved as %s", filename))) - m.status = "saved" - return *m, nil -} - -func (m *chatModel) handlePromptCommand(args string) (tea.Model, tea.Cmd) { - if strings.TrimSpace(args) != "" { - return m.addDebugError(fmt.Errorf("usage: /prompt")) - } - req, tokens := m.requestPreview() - m.promptDebug = &chatPromptDebug{ - request: req, - tokens: tokens, - } - m.flowPrintedLines = 0 - m.selection = chatSelection{} - m.status = "prompt" - return *m, tea.Batch(tea.ClearScreen, tea.EnableMouseCellMotion) -} - -func (m chatModel) updatePromptDebug(msg tea.KeyMsg) (tea.Model, tea.Cmd) { - if m.promptDebug == nil { - return m, nil - } - switch msg.Type { - case tea.KeyEsc, tea.KeyCtrlC, tea.KeyEnter: - return m.closePromptDebug() - case tea.KeyUp, tea.KeyCtrlP: - m.promptDebug.scroll-- - case tea.KeyDown, tea.KeyCtrlN: - m.promptDebug.scroll++ - case tea.KeyPgUp: - m.promptDebug.scroll -= max(1, m.promptDebugPageSize()) - case tea.KeyPgDown: - m.promptDebug.scroll += max(1, m.promptDebugPageSize()) - case tea.KeyHome, tea.KeyCtrlHome: - m.promptDebug.scroll = 0 - case tea.KeyEnd, tea.KeyCtrlEnd: - m.promptDebug.scroll = m.promptDebugMaxScroll() - } - if m.promptDebug != nil { - m.promptDebug.scroll = clamp(m.promptDebug.scroll, 0, m.promptDebugMaxScroll()) - } - return m, nil -} - -func (m chatModel) closePromptDebug() (tea.Model, tea.Cmd) { - m.promptDebug = nil - m.status = "ready" - m.flowPrintedLines = 0 - next, printCmd := m.flowTranscriptFlushCmd() - return next, tea.Sequence(tea.DisableMouse, tea.ClearScreen, printCmd) -} - -func (m chatModel) renderPromptDebug(width, height int) string { - if width <= 0 { - width = 80 - } - if height <= 0 { - height = 24 - } - if m.promptDebug == nil { - return renderFullFrame("", width, height) - } - - header := []string{ - chatPickerTitleStyle.Render("Prompt"), - chatPickerMetaStyle.Render("full request preview • /save saved as .json"), - "", - } - footer := chatPickerMetaStyle.Render("↑/↓ scroll • pgup/pgdn page • enter/esc close") - bodyHeight := max(0, height-len(header)-1) - body := m.promptDebugLines(width) - maxScroll := max(0, len(body)-bodyHeight) - scroll := clamp(m.promptDebug.scroll, 0, maxScroll) - if bodyHeight < len(body) { - body = body[scroll:min(len(body), scroll+bodyHeight)] - } - - lines := slices.Clone(header) - lines = append(lines, body...) - for len(lines) < height-1 { - lines = append(lines, "") - } - lines = append(lines, footer) - return renderFrameLines(lines, width, height) -} - -func (m chatModel) promptDebugPageSize() int { - height := m.height - if height <= 0 { - height = 24 - } - return max(1, height-5) -} - -func (m chatModel) promptDebugMaxScroll() int { - if m.promptDebug == nil { - return 0 - } - width := m.viewWidth() - height := m.height - if height <= 0 { - height = 24 - } - bodyHeight := max(0, height-4) - return max(0, len(m.promptDebugLines(width))-bodyHeight) -} - -func (m *chatModel) addDebugError(err error) (tea.Model, tea.Cmd) { - m.entries = append(m.entries, newChatEntry(chatEntry{role: "error", content: err.Error(), err: err.Error()})) - m.status = "error" - return *m, nil -} - -func (m chatModel) rawRequestJSON() (string, error) { - req, _ := m.requestPreview() - data, err := json.MarshalIndent(req, "", " ") - if err != nil { - return "", err - } - return string(data), nil -} - -func (m chatModel) requestPreview() (api.ChatRequest, int) { - opts := m.previewRunOptions() - messages := m.previewMessages() - req := m.previewChatRequest(opts, messages) - return req, m.estimatePromptTokens(messages, opts.SystemPrompt) -} - -func (m chatModel) previewRunOptions() coreagent.RunOptions { - return coreagent.RunOptions{ - ChatID: m.chatID, - Model: m.opts.Model, - SystemPrompt: m.systemPrompt(""), - Format: m.opts.Format, - Options: m.opts.Options, - Think: m.opts.Think, - KeepAlive: m.opts.KeepAlive, - } -} - -func (m chatModel) previewMessages() []api.Message { - if len(m.liveMessages) > 0 { - return slices.Clone(m.liveMessages) - } - return slices.Clone(m.messages) -} - -func (m chatModel) previewChatRequest(opts coreagent.RunOptions, messages []api.Message) api.ChatRequest { - requestMessages := slices.Clone(messages) - if strings.TrimSpace(opts.SystemPrompt) != "" { - withSystem := make([]api.Message, 0, len(requestMessages)+1) - withSystem = append(withSystem, api.Message{Role: "system", Content: opts.SystemPrompt}) - requestMessages = append(withSystem, requestMessages...) - } - - format := opts.Format - if format == "json" { - format = `"` + format + `"` - } - - req := api.ChatRequest{ - Model: opts.Model, - Messages: requestMessages, - Format: json.RawMessage(format), - Options: opts.Options, - Think: opts.Think, - } - if opts.KeepAlive != nil { - req.KeepAlive = opts.KeepAlive - } - if m.opts.Tools != nil && !m.opts.ToolsDisabled { - req.Tools = m.opts.Tools.Tools() - } - return req -} - -func (m *chatModel) promptDebugLines(width int) []string { - if m.promptDebug == nil { - return nil - } - if m.promptDebug.lines != nil && m.promptDebug.linesWidth == width { - return m.promptDebug.lines - } - req := m.promptDebug.request - innerWidth := max(20, width-2) - lines := []string{ - chatHeaderStyle.Render("Request"), - promptDebugFieldLine("model", req.Model, innerWidth), - promptDebugFieldLine("estimated prompt", m.promptTokenText(m.promptDebug.tokens), innerWidth), - promptDebugFieldLine("messages", fmt.Sprint(len(req.Messages)), innerWidth), - promptDebugFieldLine("tools", fmt.Sprint(len(req.Tools)), innerWidth), - } - if len(req.Format) > 0 { - lines = append(lines, promptDebugFieldLine("format", strings.TrimSpace(string(req.Format)), innerWidth)) - } - if req.Options != nil { - lines = append(lines, promptDebugMapLines("options", req.Options, innerWidth)...) - } - if req.Think != nil { - lines = append(lines, promptDebugBlockLines("think", req.Think.String(), innerWidth, chatHistoryTextStyle)...) - } - if req.KeepAlive != nil { - lines = append(lines, promptDebugFieldLine("keep_alive", req.KeepAlive.String(), innerWidth)) - } - lines = append(lines, "", chatHeaderStyle.Render("Messages")) - if len(req.Messages) == 0 { - lines = append(lines, chatMetaStyle.Render("none")) - } else { - for i, msg := range req.Messages { - if i > 0 { - lines = append(lines, "") - } - lines = append(lines, promptDebugMessageLines(i+1, msg, innerWidth)...) - } - } - lines = append(lines, "", chatHeaderStyle.Render("Tools")) - if len(req.Tools) == 0 { - lines = append(lines, chatMetaStyle.Render("none")) - } else { - for i, tool := range req.Tools { - if i > 0 { - lines = append(lines, "") - } - lines = append(lines, promptDebugToolLines(i+1, tool, innerWidth)...) - } - } - m.promptDebug.lines = lines - m.promptDebug.linesWidth = width - return m.promptDebug.lines -} - -func promptDebugFieldLine(label, value string, width int) string { - labelText := label + ":" - value = strings.TrimSpace(value) - if value == "" { - value = "_empty_" - } - line := chatHistoryLabelStyle.Render(labelText) + " " + chatHistoryTextStyle.Render(value) - return truncateRenderedLine(line, width) -} - -func promptDebugMessageLines(index int, msg api.Message, width int) []string { - role := promptMessageLabel(msg) - header := fmt.Sprintf("%d. %s", index, role) - lines := []string{historyRoleStyle(msg.Role).Render(header)} - - if strings.TrimSpace(msg.Thinking) != "" { - lines = append(lines, promptDebugBlockLines("thinking", msg.Thinking, width, chatHistoryTextStyle)...) - } - if msg.Role != "tool" && (strings.TrimSpace(msg.Content) != "" || (msg.Role != "assistant" && len(msg.ToolCalls) == 0 && len(msg.Images) == 0 && msg.Thinking == "")) { - lines = append(lines, promptDebugBlockLines("content", msg.Content, width, chatHistoryTextStyle)...) - } - if len(msg.ToolCalls) > 0 { - for i, call := range msg.ToolCalls { - lines = append(lines, promptDebugToolCallLines(i+1, call, width)...) - } - } - if msg.Role == "tool" { - if msg.ToolName != "" { - lines = append(lines, " "+chatHistoryLabelStyle.Render("tool_name:")+" "+chatHistoryTextStyle.Render(msg.ToolName)) - } - if msg.ToolCallID != "" { - lines = append(lines, " "+chatHistoryLabelStyle.Render("tool_call_id:")+" "+chatHistoryTextStyle.Render(msg.ToolCallID)) - } - lines = append(lines, promptDebugBlockLines("tool result", promptDebugToolResult(msg.Content), width, chatHistoryTextStyle)...) - } - if len(msg.Images) > 0 { - lines = append(lines, " "+chatHistoryLabelStyle.Render(fmt.Sprintf("%d image%s", len(msg.Images), pluralSuffix(len(msg.Images))))) - } - return lines -} - -func promptDebugToolResult(content string) string { - runes := []rune(content) - if len(runes) <= maxPromptDebugToolResultRunes { - return content - } - return string(runes[:maxPromptDebugToolResultRunes-3]) + "..." -} - -func promptDebugMapLines(label string, values map[string]any, width int) []string { - lines := []string{" " + chatHistoryLabelStyle.Render(label+":")} - if len(values) == 0 { - return append(lines, " "+chatMetaStyle.Render("_empty_")) - } - keys := make([]string, 0, len(values)) - for key := range values { - keys = append(keys, key) - } - slices.Sort(keys) - for _, key := range keys { - lines = append(lines, promptDebugValueLine(4, key, values[key], width)...) - } - return lines -} - -func promptDebugToolLines(index int, tool api.Tool, width int) []string { - name := strings.TrimSpace(tool.Function.Name) - if name == "" { - name = "_unnamed_" - } - lines := []string{historyRoleStyle("tool").Render(fmt.Sprintf("%d. %s", index, name))} - if strings.TrimSpace(tool.Function.Description) != "" { - lines = append(lines, promptDebugBlockLines("description", tool.Function.Description, width, chatHistoryTextStyle)...) - } - - params := tool.Function.Parameters - if params.Type != "" || params.Properties != nil { - kind := params.Type - if kind == "" { - kind = "object" - } - lines = append(lines, " "+chatHistoryLabelStyle.Render("parameters:")+" "+chatHistoryTextStyle.Render(kind)) - } - if params.Properties == nil || params.Properties.Len() == 0 { - return lines - } - - lines = append(lines, " "+chatHistoryLabelStyle.Render("properties:")) - required := map[string]bool{} - for _, name := range params.Required { - required[name] = true - } - for name, property := range params.Properties.All() { - label := name - propertyType := property.ToTypeScriptType() - switch { - case propertyType != "" && required[name]: - label += " (" + propertyType + ", required)" - case propertyType != "": - label += " (" + propertyType + ")" - case required[name]: - label += " (required)" - } - value := strings.TrimSpace(property.Description) - if value == "" { - value = promptDebugPropertyDetails(property) - } - lines = append(lines, promptDebugTextLine(4, label, value, width)...) - } - return lines -} - -func promptDebugToolCallLines(index int, call api.ToolCall, width int) []string { - name := strings.TrimSpace(call.Function.Name) - if name == "" { - name = "_unnamed_" - } - lines := []string{" " + chatHistoryLabelStyle.Render(fmt.Sprintf("tool call %d:", index)) + " " + chatHistoryTextStyle.Render(name)} - if strings.TrimSpace(call.ID) != "" { - lines = append(lines, promptDebugTextLine(4, "id", call.ID, width)...) - } - if call.Function.Arguments.Len() == 0 { - lines = append(lines, " "+chatHistoryLabelStyle.Render("arguments:")+" "+chatMetaStyle.Render("none")) - return lines - } - lines = append(lines, " "+chatHistoryLabelStyle.Render("arguments:")) - for key, value := range call.Function.Arguments.All() { - lines = append(lines, promptDebugValueLine(6, key, value, width)...) - } - return lines -} - -func promptDebugPropertyDetails(property api.ToolProperty) string { - var parts []string - if len(property.Enum) > 0 { - values := make([]string, 0, len(property.Enum)) - for _, value := range property.Enum { - values = append(values, promptDebugValueText(value)) - } - parts = append(parts, "one of "+strings.Join(values, ", ")) - } - if property.Properties != nil && property.Properties.Len() > 0 { - count := property.Properties.Len() - noun := "property" - if count != 1 { - noun = "properties" - } - parts = append(parts, fmt.Sprintf("%d nested %s", count, noun)) - } - if property.Items != nil { - parts = append(parts, "array items: "+promptDebugValueText(property.Items)) - } - if len(parts) == 0 { - return "_empty_" - } - return strings.Join(parts, "; ") -} - -func promptDebugValueLine(indent int, label string, value any, width int) []string { - return promptDebugTextLine(indent, label, promptDebugValueText(value), width) -} - -func promptDebugTextLine(indent int, label, value string, width int) []string { - prefix := strings.Repeat(" ", indent) + chatHistoryLabelStyle.Render(label+":") - value = strings.TrimSpace(value) - if value == "" { - value = "_empty_" - } - wrapWidth := max(20, width-indent-lipgloss.Width(label)-2) - wrapped := wrapChatText(value, wrapWidth) - if len(wrapped) == 0 { - return []string{prefix + " " + chatMetaStyle.Render("_empty_")} - } - lines := []string{prefix + " " + chatHistoryTextStyle.Render(wrapped[0])} - for _, line := range wrapped[1:] { - lines = append(lines, strings.Repeat(" ", indent+2)+chatHistoryTextStyle.Render(line)) - } - return lines -} - -func promptDebugValueText(value any) string { - switch v := value.(type) { - case nil: - return "null" - case string: - return v - case fmt.Stringer: - return v.String() - case []any: - parts := make([]string, 0, len(v)) - for _, item := range v { - parts = append(parts, promptDebugValueText(item)) - } - return strings.Join(parts, ", ") - case map[string]any: - keys := make([]string, 0, len(v)) - for key := range v { - keys = append(keys, key) - } - slices.Sort(keys) - parts := make([]string, 0, len(keys)) - for _, key := range keys { - parts = append(parts, key+": "+promptDebugValueText(v[key])) - } - return strings.Join(parts, ", ") - default: - return fmt.Sprint(value) - } -} - -func promptDebugBlockLines(label, value string, width int, style lipgloss.Style) []string { - lines := []string{" " + chatHistoryLabelStyle.Render(label+":")} - if value == "" { - return append(lines, " "+chatMetaStyle.Render("_empty_")) - } - for _, raw := range strings.Split(strings.TrimRight(value, "\n"), "\n") { - if raw == "" { - lines = append(lines, "") - continue - } - for _, wrapped := range wrapChatText(raw, max(20, width-4)) { - lines = append(lines, " "+style.Render(wrapped)) - } - } - return lines -} - -func (m chatModel) promptTokenText(tokens int) string { - window := m.displayContextWindowTokens() - if window > 0 { - return fmt.Sprintf("%s / %s tokens", formatPromptTokenCount(max(tokens, 0)), formatPromptTokenCount(window)) - } - return formatTokenCount(tokens) -} - -func formatPromptTokenCount(count int) string { - sign := "" - if count < 0 { - sign = "-" - count = -count - } - if count < 100_000 { - return sign + fmt.Sprint(count) - } - if count >= 950_000 { - return fmt.Sprintf("%s%dM", sign, int(float64(count)/1_000_000+0.5)) - } - return fmt.Sprintf("%s%dk", sign, int(float64(count)/1024+0.5)) -} - -func promptMessageLabel(msg api.Message) string { - if msg.Role == "tool" && msg.ToolName != "" { - return msg.Role + ":" + msg.ToolName - } - return msg.Role -} - -func saveRequestFilename(args string) (string, error) { - args = strings.TrimSpace(args) - if args == "" { - return "", fmt.Errorf("usage: /save ") - } - if strings.HasPrefix(args, ">") { - args = strings.TrimSpace(strings.TrimPrefix(args, ">")) - } - fields := strings.Fields(args) - if len(fields) != 1 { - return "", fmt.Errorf("usage: /save ") - } - filename := strings.TrimSpace(fields[0]) - if filename == "" || filename == "." || filename == ".." || strings.ContainsAny(filename, `/\`) || filepath.IsAbs(filename) { - return "", fmt.Errorf("save filename must be a file name, not a path") - } - if !strings.HasSuffix(strings.ToLower(filename), ".json") { - filename += ".json" - } - return filename, nil -} - -func (m chatModel) debugWorkingDir() (string, error) { - dir := strings.TrimSpace(m.currentWorkingDir()) - if dir != "" { - return dir, nil - } - return os.Getwd() -} diff --git a/cmd/tui/chat/events.go b/cmd/tui/chat/events.go deleted file mode 100644 index f498868e22c..00000000000 --- a/cmd/tui/chat/events.go +++ /dev/null @@ -1,374 +0,0 @@ -package chat - -import ( - "context" - "slices" - "strings" - "time" - - tea "github.com/charmbracelet/bubbletea" - - coreagent "github.com/ollama/ollama/agent" - "github.com/ollama/ollama/api" -) - -type chatAgentMsg struct { - event coreagent.Event -} - -type chatApprovalPromptMsg struct { - request coreagent.ApprovalRequest - reply chan<- coreagent.Approval -} - -type chatRunDoneMsg struct { - result *coreagent.RunResult - err error - newMessagesPersisted bool - persistedMessages []api.Message -} - -type chatCompactDoneMsg struct { - result coreagent.CompactionResult - err error -} - -type chatCompactProgressMsg struct { - tokens int -} - -// resetStreamingState clears the transient streaming flags that every -// non-streaming event resets before applying its own state. -func (m *chatModel) resetStreamingState() { - m.finishThinkingEntry() - m.awaitingModel = false - m.thinking = false - m.thinkingTokens = 0 -} - -// resetRunState clears all run-progress flags (streaming plus compaction -// progress) for terminal events that fully reset the run view. -func (m *chatModel) resetRunState() { - m.finishThinkingEntry() - m.awaitingModel = false - m.compacting = false - m.compactingTokens = 0 - m.detectedToolCalls = nil - m.thinking = false - m.thinkingTokens = 0 -} - -type chatModelPreloadDoneMsg struct { - model string - contextWindowTokens int - err error -} - -type chatEventsClosedMsg struct{} - -type chatTickMsg struct{} - -func (m *chatModel) applyAgentEvent(event coreagent.Event) { - contextChanged := false - - switch event.Type { - case coreagent.EventThinkingDelta: - m.awaitingModel = false - if event.Thinking != "" { - if event.Tokens > 0 { - m.thinkingTokens = max(m.thinkingTokens, event.Tokens) - } else { - m.thinkingTokens += approximateTokenCount(event.Thinking) - } - idx := m.ensureLiveAssistantMessage() - if !m.thinking { - m.thinkingPhaseStart = len(m.liveMessages[idx].Thinking) - } - m.thinking = true - m.liveMessages[idx].Thinking += event.Thinking - m.syncThinkingEntry(m.liveMessages[idx].Thinking[m.thinkingPhaseStart:]) - contextChanged = true - } - case coreagent.EventMessageDelta: - m.resetStreamingState() - m.spinner = 0 - m.detectedToolCalls = nil - idx := m.ensureAssistantEntry() - m.entries[idx].content += event.Content - m.markEntryDirty(idx) - msgIdx := m.ensureLiveAssistantMessage() - m.liveMessages[msgIdx].Content += event.Content - contextChanged = true - case coreagent.EventToolCallDetected: - m.finishThinkingEntry() - m.awaitingModel = m.running - m.thinking = false - m.thinkingTokens = 0 - m.groupCompletedToolHistory() - m.detectedToolCalls = nil - m.addDetectedToolCalls(event.ToolCalls) - idx := m.ensureLiveAssistantMessage() - m.liveMessages[idx].ToolCalls = append(m.liveMessages[idx].ToolCalls, event.ToolCalls...) - contextChanged = true - case coreagent.EventToolStarted: - m.resetStreamingState() - startedAt := time.Now() - idx := m.findActiveToolEntry(event.ToolCallID) - if idx < 0 { - m.groupCompletedToolHistory() - m.entries = append(m.entries, newChatEntry(chatEntry{role: "tool"})) - idx = len(m.entries) - 1 - } - m.entries[idx].detail = event.ToolName - m.entries[idx].label = toolInvocationLabel(event.ToolName, event.Args) - m.entries[idx].status = "running" - m.entries[idx].toolID = event.ToolCallID - m.entries[idx].args = event.Args - m.entries[idx].startedAt = startedAt - m.applyToolOutputModeTo(idx) - m.markEntryDirty(idx) - case coreagent.EventToolFinished: - m.resetStreamingState() - if event.WorkingDir != "" { - m.workingDir = event.WorkingDir - } - startedAt := m.toolStartedAt(event.ToolCallID) - status := toolFinishedStatus(event) - idx := m.findToolEntry(event.ToolCallID) - if idx < 0 { - m.entries = append(m.entries, newChatEntry(chatEntry{role: "tool"})) - idx = len(m.entries) - 1 - } - m.entries[idx].content = event.Content - m.entries[idx].label = toolInvocationLabel(event.ToolName, event.Args) - m.entries[idx].detail = event.ToolName - m.entries[idx].status = status - if status != "denied" { - m.entries[idx].err = event.Error - } - m.entries[idx].toolID = event.ToolCallID - m.entries[idx].args = event.Args - m.entries[idx].startedAt = startedAt - m.entries[idx].finishedAt = time.Now() - m.applyToolOutputModeTo(idx) - m.markEntryDirty(idx) - m.liveMessages = append(m.liveMessages, api.Message{ - Role: "tool", - Content: event.Content, - ToolName: event.ToolName, - ToolCallID: event.ToolCallID, - }) - if m.running && status != "denied" && !m.hasPendingDetectedToolCalls() { - m.awaitingModel = true - } - contextChanged = true - case coreagent.EventCompacted: - m.resetRunState() - if len(event.Messages) > 0 { - m.liveMessages = slices.Clone(event.Messages) - m.messages = slices.Clone(event.Messages) - contextChanged = true - } - m.status = "compacted" - case coreagent.EventCompactionStarted: - m.awaitingModel = false - m.compacting = true - m.compactingTokens = 0 - m.thinking = false - m.thinkingTokens = 0 - m.status = "compacting" - case coreagent.EventCompactionProgress: - m.awaitingModel = false - m.compacting = true - m.thinking = false - m.thinkingTokens = 0 - if event.Tokens > m.compactingTokens { - m.compactingTokens = event.Tokens - } - case coreagent.EventCompactionSkipped: - m.resetRunState() - message := event.Content - if strings.TrimSpace(message) == "" { - message = coreagent.CompactionSkippedMessage(event.Error) - } - m.entries = append(m.entries, newChatEntry(chatEntry{role: "system", content: message})) - m.status = "compact skipped" - case coreagent.EventError: - m.resetRunState() - m.eventErrorRendered = true - m.entries = append(m.entries, newChatEntry(chatEntry{role: "error", content: event.Error, err: event.Error})) - } - - if contextChanged { - m.refreshLiveContextEstimate() - } -} - -func (m *chatModel) addDetectedToolCalls(calls []api.ToolCall) { - if len(calls) == 0 { - return - } - seen := make(map[string]struct{}, len(m.detectedToolCalls)+len(calls)) - for _, entry := range m.detectedToolCalls { - if entry.toolID != "" { - seen[entry.toolID] = struct{}{} - } - } - for _, call := range calls { - if call.ID != "" { - if _, ok := seen[call.ID]; ok { - continue - } - seen[call.ID] = struct{}{} - } - args := call.Function.Arguments.ToMap() - m.detectedToolCalls = append(m.detectedToolCalls, newChatEntry(chatEntry{ - role: "tool", - label: toolInvocationLabel(call.Function.Name, args), - detail: call.Function.Name, - status: "queued", - toolID: call.ID, - args: args, - })) - } -} - -func toolFinishedStatus(event coreagent.Event) string { - switch event.ToolStatus { - case coreagent.ToolStatusDenied: - return "denied" - case coreagent.ToolStatusDisabled: - return "disabled" - case coreagent.ToolStatusDone: - return "done" - } - // failed/skipped/unknown: derive from content and error fields. - if isDeniedToolResult(event.Content) || isDeniedToolResult(event.Error) { - return "denied" - } - if event.Error != "" { - return "error" - } - return "done" -} - -func messagesEndWithCompactionResult(messages []api.Message) bool { - if len(messages) == 0 { - return false - } - return coreagent.IsCompactionToolResult(messages[len(messages)-1]) -} - -func (m chatModel) awaitingToolStart() bool { - for i := len(m.liveMessages) - 1; i >= 0; i-- { - msg := m.liveMessages[i] - if msg.Role != "assistant" { - continue - } - if len(msg.ToolCalls) == 0 { - return false - } - for _, call := range msg.ToolCalls { - if call.ID == "" || m.findToolEntry(call.ID) < 0 { - return true - } - } - return false - } - return false -} - -func (m *chatModel) ensureLiveAssistantMessage() int { - if len(m.liveMessages) > 0 && m.liveMessages[len(m.liveMessages)-1].Role == "assistant" { - return len(m.liveMessages) - 1 - } - m.liveMessages = append(m.liveMessages, api.Message{Role: "assistant"}) - return len(m.liveMessages) - 1 -} - -func (m *chatModel) refreshLiveContextEstimate() { - messages := m.liveMessages - if len(messages) == 0 { - messages = m.messages - } - m.contextTokens = m.estimatePromptTokens(messages, "") - m.contextEstimate = true -} - -//nolint:containedctx // event sinks need the session context to unblock sends on cancellation. -type chatEventSink struct { - ctx context.Context - ch chan<- tea.Msg - newMessagesPersisted *bool -} - -func (s chatEventSink) Emit(event coreagent.Event) error { - if s.newMessagesPersisted != nil { - *s.newMessagesPersisted = true - } - select { - case s.ch <- chatAgentMsg{event: event}: - return nil - case <-s.ctx.Done(): - return s.ctx.Err() - } -} - -func waitForChatMsg(ch <-chan tea.Msg) tea.Cmd { - if ch == nil { - return nil - } - return func() tea.Msg { - msg, ok := <-ch - if !ok { - return chatEventsClosedMsg{} - } - return msg - } -} - -func (m *chatModel) scheduleTick() tea.Cmd { - if m.tickActive { - return nil - } - m.tickActive = true - return chatTickCmd() -} - -func chatTickCmd() tea.Cmd { - return tea.Tick(350*time.Millisecond, func(time.Time) tea.Msg { - return chatTickMsg{} - }) -} - -func preloadModelCmd(ctx context.Context, preload func(context.Context, string, *api.ThinkValue) (int, error), model string, think *api.ThinkValue) tea.Cmd { - if preload == nil || strings.TrimSpace(model) == "" { - return nil - } - if think != nil { - copied := *think - think = &copied - } - return func() tea.Msg { - if ctx == nil { - ctx = context.Background() - } - tokens, err := preload(ctx, model, think) - return chatModelPreloadDoneMsg{model: model, contextWindowTokens: tokens, err: err} - } -} - -func isUnsupportedThinkingError(err error) bool { - if err == nil { - return false - } - text := strings.ToLower(err.Error()) - return strings.Contains(text, "does not support thinking") -} - -func thinkRequestsThinking(think *api.ThinkValue) bool { - if think == nil { - return false - } - return think.Bool() -} diff --git a/cmd/tui/chat/events_test.go b/cmd/tui/chat/events_test.go deleted file mode 100644 index 4c57de910cc..00000000000 --- a/cmd/tui/chat/events_test.go +++ /dev/null @@ -1,454 +0,0 @@ -package chat - -import ( - "strings" - "testing" - - coreagent "github.com/ollama/ollama/agent" - "github.com/ollama/ollama/api" -) - -func TestApplyAgentEventStreamsAssistantContent(t *testing.T) { - m := chatModel{running: true} - - m.applyAgentEvent(coreagent.Event{Type: coreagent.EventMessageDelta, Content: "hello"}) - m.applyAgentEvent(coreagent.Event{Type: coreagent.EventMessageDelta, Content: " world"}) - - if len(m.entries) != 1 || m.entries[0].role != "assistant" || m.entries[0].content != "hello world" { - t.Fatalf("entries = %#v", m.entries) - } - if len(m.liveMessages) != 1 || m.liveMessages[0].Content != "hello world" { - t.Fatalf("live messages = %#v", m.liveMessages) - } -} - -func TestApplyAgentEventStreamsThinkingThenCollapsesOnAssistantOrTool(t *testing.T) { - m := chatModel{running: true} - - m.applyAgentEvent(coreagent.Event{Type: coreagent.EventThinkingDelta, Thinking: "first "}) - m.applyAgentEvent(coreagent.Event{Type: coreagent.EventThinkingDelta, Thinking: "second", Tokens: 7}) - if len(m.entries) != 1 || m.entries[0].role != "thinking" || !m.entries[0].expanded || m.entries[0].content != "first second" { - t.Fatalf("live thinking entry = %#v", m.entries) - } - if got := m.liveMessages[0].Thinking; got != "first second" { - t.Fatalf("live message thinking = %q, want full streamed value", got) - } - if view := stripANSI(m.renderTranscript(100)); !strings.Contains(view, "Thinking ↓ 7 tokens") || !strings.Contains(view, "first second") { - t.Fatalf("live thinking trace missing from transcript:\n%s", view) - } - - m.applyAgentEvent(coreagent.Event{Type: coreagent.EventMessageDelta, Content: "answer"}) - if m.entries[0].status != "done" || m.entries[0].expanded { - t.Fatalf("assistant content should collapse thinking: %#v", m.entries[0]) - } - collapsed := stripANSI(m.renderTranscript(100)) - if !strings.Contains(collapsed, "Thought") || strings.Contains(collapsed, "7 tokens") || strings.Contains(collapsed, "first second") { - t.Fatalf("collapsed thinking should remain as a thought row without trace content:\n%s", collapsed) - } - if got := m.liveMessages[0].Thinking; got != "first second" { - t.Fatalf("collapsing display must not change request history: %q", got) - } - - m.applyAgentEvent(coreagent.Event{Type: coreagent.EventThinkingDelta, Thinking: "tool plan"}) - if entry := m.entries[len(m.entries)-1]; entry.role != "thinking" || entry.content != "tool plan" { - t.Fatalf("second thinking phase should contain only its own deltas: %#v", entry) - } - if got := m.liveMessages[0].Thinking; got != "first secondtool plan" { - t.Fatalf("message history should retain both thinking phases exactly: %q", got) - } - m.applyAgentEvent(coreagent.Event{Type: coreagent.EventToolStarted, ToolCallID: "call-1", ToolName: "bash"}) - if entry := m.entries[len(m.entries)-2]; entry.role != "thinking" || entry.status != "done" || entry.expanded { - t.Fatalf("tool transition should collapse thinking: %#v", entry) - } -} - -func TestApplyAgentEventDoesNotCreateThinkingEntryWithoutThinking(t *testing.T) { - m := chatModel{running: true} - - m.applyAgentEvent(coreagent.Event{Type: coreagent.EventThinkingDelta, Tokens: 12}) - m.applyAgentEvent(coreagent.Event{Type: coreagent.EventMessageDelta, Content: "answer"}) - if len(m.entries) != 1 || m.entries[0].role != "assistant" { - t.Fatalf("empty thinking event should not create a trace: %#v", m.entries) - } - if len(m.liveMessages) != 1 || m.liveMessages[0].Thinking != "" { - t.Fatalf("empty thinking event should not alter message history: %#v", m.liveMessages) - } -} - -func TestApplyAgentEventPreservesCollapsedThoughtsAcrossToolGrouping(t *testing.T) { - m := chatModel{running: true} - - m.applyAgentEvent(coreagent.Event{Type: coreagent.EventThinkingDelta, Thinking: "first plan", Tokens: 1}) - m.applyAgentEvent(coreagent.Event{Type: coreagent.EventToolStarted, ToolCallID: "call-1", ToolName: "bash"}) - m.applyAgentEvent(coreagent.Event{Type: coreagent.EventToolFinished, ToolCallID: "call-1", ToolName: "bash", Content: "one"}) - m.applyAgentEvent(coreagent.Event{Type: coreagent.EventThinkingDelta, Thinking: "second plan", Tokens: 1}) - m.applyAgentEvent(coreagent.Event{Type: coreagent.EventToolStarted, ToolCallID: "call-2", ToolName: "bash"}) - m.applyAgentEvent(coreagent.Event{Type: coreagent.EventToolFinished, ToolCallID: "call-2", ToolName: "bash", Content: "two"}) - m.applyAgentEvent(coreagent.Event{Type: coreagent.EventToolStarted, ToolCallID: "call-3", ToolName: "bash"}) - - if len(m.entries) != 5 { - t.Fatalf("entries = %#v, want two thought rows and three tool rows", m.entries) - } - for _, index := range []int{0, 2} { - entry := m.entries[index] - if entry.role != "thinking" || entry.status != "done" || entry.expanded { - t.Fatalf("collapsed thought %d = %#v", index, entry) - } - } - if transcript := stripANSI(m.renderTranscript(100)); strings.Count(transcript, "Thought") != 2 || strings.Contains(transcript, "1 token") { - t.Fatalf("transcript should retain both thought rows:\n%s", transcript) - } -} - -func TestApplyAgentEventTracksToolLifecycle(t *testing.T) { - m := chatModel{running: true} - args := map[string]any{"command": "pwd"} - - m.applyAgentEvent(coreagent.Event{ - Type: coreagent.EventToolStarted, - ToolCallID: "call-1", - ToolName: "bash", - Args: args, - }) - m.applyAgentEvent(coreagent.Event{ - Type: coreagent.EventToolFinished, - ToolCallID: "call-1", - ToolName: "bash", - Args: args, - Content: "ok", - }) - - if len(m.entries) != 1 { - t.Fatalf("entries = %#v", m.entries) - } - entry := m.entries[0] - if entry.status != "done" || entry.content != "ok" || !strings.Contains(entry.label, "Bash") { - t.Fatalf("tool entry = %#v", entry) - } - if line := stripANSI(toolStatusLine(entry)); line != `Bash("pwd")` { - t.Fatalf("tool status line = %q, want command label", line) - } - if len(m.liveMessages) != 1 || m.liveMessages[0].Role != "tool" || m.liveMessages[0].Content != "ok" { - t.Fatalf("live messages = %#v", m.liveMessages) - } -} - -func TestApplyAgentEventRendersDeniedCommandAsDenied(t *testing.T) { - m := chatModel{running: true} - args := map[string]any{"command": "pwd"} - - m.applyAgentEvent(coreagent.Event{ - Type: coreagent.EventToolFinished, - ToolStatus: coreagent.ToolStatusDenied, - ToolCallID: "call-1", - ToolName: "bash", - Args: args, - Content: "Tool execution denied.", - Error: "Tool execution denied.", - }) - - if len(m.entries) != 1 { - t.Fatalf("entries = %#v", m.entries) - } - entry := m.entries[0] - if entry.status != "denied" { - t.Fatalf("tool status = %q, want denied: %#v", entry.status, entry) - } - if line := stripANSI(toolStatusLine(entry)); line != `Bash("pwd") denied` { - t.Fatalf("tool status line = %q, want denied command label", line) - } -} - -func TestApplyAgentEventShowsWorkingWhileAwaitingCloudToolStart(t *testing.T) { - args := api.NewToolCallFunctionArguments() - args.Set("command", "pwd") - m := chatModel{ - running: true, - } - - m.applyAgentEvent(coreagent.Event{ - Type: coreagent.EventToolCallDetected, - ToolCalls: []api.ToolCall{{ - ID: "call-1", - Function: api.ToolCallFunction{ - Name: "bash", - Arguments: args, - }, - }}, - }) - - if line := stripANSI(m.activityLine()); !strings.Contains(line, "Working") { - t.Fatalf("activityLine = %q, want Working while tool call is pending", line) - } -} - -func TestActivityLineShowsWorkingWhileAwaitingModelBeforeFirstEvent(t *testing.T) { - m := chatModel{ - running: true, - awaitingModel: true, - spinner: 0, - } - - if line := stripANSI(m.activityLine()); !strings.Contains(line, "Working") { - t.Fatalf("activityLine = %q, want Working while stream is open before first event", line) - } -} - -func TestActivityLineShowsWorkingAfterAssistantContentGoesIdle(t *testing.T) { - m := chatModel{ - running: true, - spinner: idleWorkingDelayTicks, - } - - m.applyAgentEvent(coreagent.Event{Type: coreagent.EventMessageDelta, Content: "I will inspect that next."}) - if line := strings.TrimSpace(stripANSI(m.activityLine())); line != "" { - t.Fatalf("activityLine immediately after content = %q, want quiet until the idle delay", line) - } - - m.spinner = idleWorkingDelayTicks - if line := stripANSI(m.activityLine()); !strings.Contains(line, "Working") { - t.Fatalf("activityLine after idle content stream = %q, want Working while stream remains open", line) - } -} - -func TestApplyAgentEventKeepsDetectedBatchStableUntilComplete(t *testing.T) { - firstArgs := api.NewToolCallFunctionArguments() - firstArgs.Set("command", "pwd") - secondArgs := api.NewToolCallFunctionArguments() - secondArgs.Set("command", "ls") - m := chatModel{ - running: true, - spinner: idleWorkingDelayTicks, - } - - m.applyAgentEvent(coreagent.Event{ - Type: coreagent.EventToolCallDetected, - ToolCalls: []api.ToolCall{ - {ID: "call-1", Function: api.ToolCallFunction{Name: "bash", Arguments: firstArgs}}, - {ID: "call-2", Function: api.ToolCallFunction{Name: "bash", Arguments: secondArgs}}, - }, - }) - m.applyAgentEvent(coreagent.Event{Type: coreagent.EventToolStarted, ToolCallID: "call-1", ToolName: "bash", Args: firstArgs.ToMap()}) - m.applyAgentEvent(coreagent.Event{Type: coreagent.EventToolFinished, ToolCallID: "call-1", ToolName: "bash", Args: firstArgs.ToMap(), Content: "one"}) - - if len(m.entries) != 1 { - t.Fatalf("entries = %d, want first completed command row: %#v", len(m.entries), m.entries) - } - if m.entries[0].role != "tool" || m.entries[0].status != "done" { - t.Fatalf("first command should remain stable while second is pending: %#v", m.entries[0]) - } - if line := stripANSI(toolStatusLine(m.entries[0])); line != `Bash("pwd")` { - t.Fatalf("completed command line = %q", line) - } - if line := stripANSI(m.activityLine()); !strings.Contains(line, "Working") { - t.Fatalf("activityLine = %q, want Working while second command is pending", line) - } - - m.applyAgentEvent(coreagent.Event{Type: coreagent.EventToolStarted, ToolCallID: "call-2", ToolName: "bash", Args: secondArgs.ToMap()}) - if len(m.entries) != 2 { - t.Fatalf("entries after second start = %d, want finished command plus running command: %#v", len(m.entries), m.entries) - } - if line := stripANSI(toolStatusLine(m.entries[0])); line != `Bash("pwd")` { - t.Fatalf("finished command line after second start = %q", line) - } - if line := stripANSI(toolStatusLine(m.entries[1])); line != `Bash("ls")` { - t.Fatalf("running command line = %q", line) - } - - m.applyAgentEvent(coreagent.Event{Type: coreagent.EventToolFinished, ToolCallID: "call-2", ToolName: "bash", Args: secondArgs.ToMap(), Content: "two"}) - if line := stripANSI(m.activityLine()); !strings.Contains(line, "Working") { - t.Fatalf("activityLine after completed batch = %q, want Working while waiting for next model response", line) - } - if len(m.entries) != 2 { - t.Fatalf("entries after batch completion = %d, want stable command rows until the next tool boundary: %#v", len(m.entries), m.entries) - } - for i, want := range []string{`Bash("pwd")`, `Bash("ls")`} { - if line := stripANSI(toolStatusLine(m.entries[i])); line != want { - t.Fatalf("completed command row %d = %q, want %q", i, line, want) - } - } - - thirdArgs := api.NewToolCallFunctionArguments() - thirdArgs.Set("command", "date") - m.applyAgentEvent(coreagent.Event{ - Type: coreagent.EventToolCallDetected, - ToolCalls: []api.ToolCall{ - {ID: "call-3", Function: api.ToolCallFunction{Name: "bash", Arguments: thirdArgs}}, - }, - }) - m.applyAgentEvent(coreagent.Event{Type: coreagent.EventToolStarted, ToolCallID: "call-3", ToolName: "bash", Args: thirdArgs.ToMap()}) - - if len(m.entries) != 2 { - t.Fatalf("entries after next tool boundary = %d, want grouped history plus running command: %#v", len(m.entries), m.entries) - } - if m.entries[0].role != "tool_group" || len(m.entries[0].tools) != 2 { - t.Fatalf("completed detected batch should collapse at the next tool boundary: %#v", m.entries[0]) - } - if line := stripANSI(toolGroupStatusLine(m.entries[0])); line != "Ran 2 commands" { - t.Fatalf("grouped command line = %q", line) - } - if line := stripANSI(toolStatusLine(m.entries[1])); line != `Bash("date")` { - t.Fatalf("running command line = %q", line) - } -} - -func TestApplyAgentEventDoesNotCollapsePartialDetectedBatch(t *testing.T) { - firstArgs := api.NewToolCallFunctionArguments() - firstArgs.Set("command", "pwd") - secondArgs := api.NewToolCallFunctionArguments() - secondArgs.Set("command", "ls") - thirdArgs := api.NewToolCallFunctionArguments() - thirdArgs.Set("command", "date") - m := chatModel{running: true} - - m.applyAgentEvent(coreagent.Event{ - Type: coreagent.EventToolCallDetected, - ToolCalls: []api.ToolCall{ - {ID: "call-1", Function: api.ToolCallFunction{Name: "bash", Arguments: firstArgs}}, - {ID: "call-2", Function: api.ToolCallFunction{Name: "bash", Arguments: secondArgs}}, - {ID: "call-3", Function: api.ToolCallFunction{Name: "bash", Arguments: thirdArgs}}, - }, - }) - m.applyAgentEvent(coreagent.Event{Type: coreagent.EventToolStarted, ToolCallID: "call-1", ToolName: "bash", Args: firstArgs.ToMap()}) - m.applyAgentEvent(coreagent.Event{Type: coreagent.EventToolFinished, ToolCallID: "call-1", ToolName: "bash", Args: firstArgs.ToMap(), Content: "one"}) - m.applyAgentEvent(coreagent.Event{Type: coreagent.EventToolStarted, ToolCallID: "call-2", ToolName: "bash", Args: secondArgs.ToMap()}) - m.applyAgentEvent(coreagent.Event{Type: coreagent.EventToolFinished, ToolCallID: "call-2", ToolName: "bash", Args: secondArgs.ToMap(), Content: "two"}) - - if line := stripANSI(m.activityLine()); !strings.Contains(line, "Working") { - t.Fatalf("activityLine before final detected call = %q, want Working while final tool is pending", line) - } - if len(m.entries) != 2 { - t.Fatalf("entries before final detected call = %d, want two stable rows: %#v", len(m.entries), m.entries) - } - for i, want := range []string{`Bash("pwd")`, `Bash("ls")`} { - if line := stripANSI(toolStatusLine(m.entries[i])); line != want { - t.Fatalf("tool row %d = %q, want %q", i, line, want) - } - } - - m.applyAgentEvent(coreagent.Event{Type: coreagent.EventToolStarted, ToolCallID: "call-3", ToolName: "bash", Args: thirdArgs.ToMap()}) - m.applyAgentEvent(coreagent.Event{Type: coreagent.EventToolFinished, ToolCallID: "call-3", ToolName: "bash", Args: thirdArgs.ToMap(), Content: "three"}) - - if len(m.entries) != 3 { - t.Fatalf("entries after full detected batch = %#v, want stable tool rows until the next tool boundary", m.entries) - } - for i, want := range []string{`Bash("pwd")`, `Bash("ls")`, `Bash("date")`} { - if line := stripANSI(toolStatusLine(m.entries[i])); line != want { - t.Fatalf("tool row %d = %q, want %q", i, line, want) - } - } - - fourthArgs := api.NewToolCallFunctionArguments() - fourthArgs.Set("command", "whoami") - m.applyAgentEvent(coreagent.Event{ - Type: coreagent.EventToolCallDetected, - ToolCalls: []api.ToolCall{ - {ID: "call-4", Function: api.ToolCallFunction{Name: "bash", Arguments: fourthArgs}}, - }, - }) - m.applyAgentEvent(coreagent.Event{Type: coreagent.EventToolStarted, ToolCallID: "call-4", ToolName: "bash", Args: fourthArgs.ToMap()}) - - if len(m.entries) != 2 || m.entries[0].role != "tool_group" || len(m.entries[0].tools) != 3 { - t.Fatalf("entries after next detected batch starts = %#v, want one grouped history entry plus active tool", m.entries) - } - if line := stripANSI(toolGroupStatusLine(m.entries[0])); line != "Ran 3 commands" { - t.Fatalf("grouped command line = %q", line) - } -} - -func TestApplyAgentEventGroupsCompletedCommandsAtNextToolBoundary(t *testing.T) { - m := chatModel{running: true} - firstArgs := map[string]any{"command": "pwd"} - secondArgs := map[string]any{"command": "ls"} - thirdArgs := map[string]any{"command": "date"} - - m.applyAgentEvent(coreagent.Event{Type: coreagent.EventToolStarted, ToolCallID: "call-1", ToolName: "bash", Args: firstArgs}) - m.applyAgentEvent(coreagent.Event{Type: coreagent.EventToolFinished, ToolCallID: "call-1", ToolName: "bash", Args: firstArgs, Content: "one"}) - m.applyAgentEvent(coreagent.Event{Type: coreagent.EventToolStarted, ToolCallID: "call-2", ToolName: "bash", Args: secondArgs}) - m.applyAgentEvent(coreagent.Event{Type: coreagent.EventToolFinished, ToolCallID: "call-2", ToolName: "bash", Args: secondArgs, Content: "two"}) - - if len(m.entries) != 2 { - t.Fatalf("entries after second finish = %d, want two stable command rows: %#v", len(m.entries), m.entries) - } - m.applyAgentEvent(coreagent.Event{Type: coreagent.EventToolStarted, ToolCallID: "call-3", ToolName: "bash", Args: thirdArgs}) - - if len(m.entries) != 2 { - t.Fatalf("entries = %d, want grouped command history plus active command: %#v", len(m.entries), m.entries) - } - if m.entries[0].role != "tool_group" || len(m.entries[0].tools) != 2 { - t.Fatalf("completed commands should be grouped when the next command starts: %#v", m.entries[0]) - } - if line := stripANSI(toolGroupStatusLine(m.entries[0])); line != "Ran 2 commands" { - t.Fatalf("grouped command line = %q", line) - } - - m.applyAgentEvent(coreagent.Event{Type: coreagent.EventToolFinished, ToolCallID: "call-3", ToolName: "bash", Args: thirdArgs, Content: "three"}) - m.applyAgentEvent(coreagent.Event{Type: coreagent.EventMessageDelta, Content: "done"}) - - if len(m.entries) != 3 { - t.Fatalf("entries after assistant content = %d, want grouped history, last command, assistant: %#v", len(m.entries), m.entries) - } - transcript := stripANSI(m.renderTranscript(100)) - if !strings.Contains(transcript, "• Ran 2 commands\n\n• Bash(\"date\")\n\n done") { - t.Fatalf("tool history should stay visually separated from assistant content:\n%s", transcript) - } -} - -func TestApplyAgentEventDoesNotGroupCompletedCommandsOnMessageDelta(t *testing.T) { - m := chatModel{running: true} - firstArgs := map[string]any{"command": "pwd"} - secondArgs := map[string]any{"command": "ls"} - - m.applyAgentEvent(coreagent.Event{Type: coreagent.EventToolStarted, ToolCallID: "call-1", ToolName: "bash", Args: firstArgs}) - m.applyAgentEvent(coreagent.Event{Type: coreagent.EventToolFinished, ToolCallID: "call-1", ToolName: "bash", Args: firstArgs, Content: "one"}) - m.applyAgentEvent(coreagent.Event{Type: coreagent.EventToolStarted, ToolCallID: "call-2", ToolName: "bash", Args: secondArgs}) - m.applyAgentEvent(coreagent.Event{Type: coreagent.EventToolFinished, ToolCallID: "call-2", ToolName: "bash", Args: secondArgs, Content: "two"}) - m.applyAgentEvent(coreagent.Event{Type: coreagent.EventMessageDelta, Content: "done"}) - - if len(m.entries) != 3 { - t.Fatalf("entries = %d, want two command rows plus assistant content: %#v", len(m.entries), m.entries) - } - if m.entries[0].role != "tool" || m.entries[1].role != "tool" || m.entries[2].role != "assistant" { - t.Fatalf("completed commands should not collapse on assistant content: %#v", m.entries) - } -} - -func TestApplyAgentEventGroupsPreviouslyDeniedCommandsAtNextToolBoundary(t *testing.T) { - m := chatModel{running: true} - firstArgs := map[string]any{"command": "pwd"} - secondArgs := map[string]any{"command": "ls"} - thirdArgs := map[string]any{"command": "date"} - - m.applyAgentEvent(coreagent.Event{Type: coreagent.EventToolFinished, ToolStatus: coreagent.ToolStatusDenied, ToolCallID: "call-1", ToolName: "bash", Args: firstArgs, Content: "Tool execution denied.", Error: "Tool execution denied."}) - m.applyAgentEvent(coreagent.Event{Type: coreagent.EventToolFinished, ToolStatus: coreagent.ToolStatusDenied, ToolCallID: "call-2", ToolName: "bash", Args: secondArgs, Content: "Tool execution denied.", Error: "Tool execution denied."}) - - if len(m.entries) != 2 { - t.Fatalf("entries = %d, want two stable denied command rows: %#v", len(m.entries), m.entries) - } - m.applyAgentEvent(coreagent.Event{Type: coreagent.EventToolStarted, ToolCallID: "call-3", ToolName: "bash", Args: thirdArgs}) - - if len(m.entries) != 2 { - t.Fatalf("entries = %d, want grouped denied command entry plus active command: %#v", len(m.entries), m.entries) - } - if m.entries[0].role != "tool_group" || len(m.entries[0].tools) != 2 { - t.Fatalf("denied commands should be grouped at the next tool boundary: %#v", m.entries[0]) - } - if line := stripANSI(toolGroupStatusLine(m.entries[0])); line != "Denied 2 commands" { - t.Fatalf("grouped command line = %q", line) - } - if line := stripANSI(toolStatusLine(m.entries[1])); line != `Bash("date")` { - t.Fatalf("running command line = %q", line) - } -} - -func TestMessagesEndWithCompactionResult(t *testing.T) { - messages := []api.Message{{ - Role: "tool", - ToolName: coreagent.CompactionToolName, - ToolCallID: coreagent.CompactionToolCallID, - Content: coreagent.CompactionSummaryMessagePrefix + "summary", - }} - if !messagesEndWithCompactionResult(messages) { - t.Fatal("expected compaction result") - } -} diff --git a/cmd/tui/chat/input.go b/cmd/tui/chat/input.go deleted file mode 100644 index 56414500257..00000000000 --- a/cmd/tui/chat/input.go +++ /dev/null @@ -1,1630 +0,0 @@ -package chat - -import ( - "context" - "fmt" - "os" - "path/filepath" - "regexp" - "slices" - "sort" - "strconv" - "strings" - "unicode" - "unicode/utf8" - - tea "github.com/charmbracelet/bubbletea" - "github.com/charmbracelet/lipgloss" - "github.com/mattn/go-runewidth" - - coreagent "github.com/ollama/ollama/agent" - "github.com/ollama/ollama/api" - "github.com/ollama/ollama/cmd/internal/filedata" -) - -type chatSlashCommand struct { - name string - usage string - description string - aliases []string - hidden bool -} - -type chatCompletion struct { - value string - label string - description string - directory bool -} - -const ( - chatPromptPrefix = "" - inputBoxHorizontalPadding = 1 - inputCursorGlyph = "█" - inputCursorMarker = "\x00" - maxInputBoxBodyLines = 12 -) - -const ( - pastedTextPlaceholderMinRunes = 1000 - pastedTextPlaceholderMinLines = 8 -) - -var chatSlashCommands = []chatSlashCommand{ - {name: "/model", description: "switch models"}, - {name: "/new", description: "start a new chat"}, - {name: "/think", description: "set thinking mode"}, - {name: "/tools", description: "toggle tools on or off"}, - {name: "/system", usage: "/system [on|off]", description: "show or set the built-in system prompt"}, - {name: "/skills", usage: "/skills [import codex|claude|pi]", description: "list or import skills"}, - {name: "/compact", description: "summarize older context"}, - {name: "/help", description: "show commands", aliases: []string{"/?"}}, - {name: "/bye", description: "exit", aliases: []string{"/exit"}}, - {name: "/prompt", description: "show full prompt, tools, and messages"}, - {name: "/save", usage: "/save ", description: "save request JSON; saved as .json"}, -} - -var skillsImportCompletions = []chatCompletion{ - {value: "/skills import codex", label: "/skills import codex", description: "import from ~/.codex/skills"}, - {value: "/skills import claude", label: "/skills import claude", description: "import from ~/.claude/skills"}, - {value: "/skills import pi", label: "/skills import pi", description: "import from ~/.pi/agent/skills"}, -} - -// BuiltinSlashCommandNames returns the names reserved by built-in slash -// commands, including aliases. -func BuiltinSlashCommandNames() []string { - names := make(map[string]struct{}) - for _, command := range chatSlashCommands { - names[strings.TrimPrefix(command.name, "/")] = struct{}{} - for _, alias := range command.aliases { - names[strings.TrimPrefix(alias, "/")] = struct{}{} - } - } - reserved := make([]string, 0, len(names)) - for name := range names { - reserved = append(reserved, name) - } - sort.Strings(reserved) - return reserved -} - -func (m *chatModel) handleSubmit() (tea.Model, tea.Cmd) { - m.syncInputPlaceholders() - input := strings.TrimSpace(string(m.input)) - if input == "" { - return *m, nil - } - _, _, hasSlashCommand := slashCommandInvocation(input) - if (m.running || m.compacting) && !hasSlashCommand { - m.status = "wait for current response" - return *m, nil - } - - attachments := cloneInputAttachments(m.inputAttachments) - pastedTexts := cloneInputPastedTexts(m.inputPastedTexts) - m.input = nil - m.inputCursor = 0 - m.inputCursorSet = false - m.inputAttachments = attachments - m.inputPastedTexts = pastedTexts - m.complete = 0 - m.resetPromptHistoryCursor() - return m.submitInput(input) -} - -func (m *chatModel) applySlashCompletion() bool { - rawInput := string(m.input) - input := strings.TrimSpace(rawInput) - if !strings.HasPrefix(input, "/") { - return false - } - if _, _, known := slashCommandInvocation(input); known && !hasSystemCommandArgument(rawInput) { - return false - } - completions := m.slashCompletions() - if len(completions) == 0 || !completionIsSelectable(completions) { - return false - } - selected := completions[clamp(m.complete, 0, len(completions)-1)] - if strings.EqualFold(selected.value, input) { - return false - } - // Reset prompt-history state: Up/Down is shared between history recall and - // slash completion, and a recalled prompt may start with "/" and trigger - // completion. Keep the two in sync when we accept a completion. - m.resetPromptHistoryCursor() - m.input = []rune(selected.value) - m.inputCursor = len(m.input) - m.inputCursorSet = true - m.complete = 0 - return true -} - -func (m *chatModel) submitInput(input string) (tea.Model, tea.Cmd) { - command, args, _ := slashCommandInvocation(input) - if command != "" { - input = strings.TrimSpace(command + " " + args) - } - skillName, skillPrompt, skillOK := m.skillSlashInvocation(input) - - switch { - case command == "/bye": - m.quitting = true - return *m, m.quitCmd() - case command == "/help": - m.entries = append(m.entries, newSlashEntry(m.helpSummary())) - return *m, nil - case command == "/model": - return m.openModelPicker(args) - case command == "/think" && args == "": - return m.openThinkPicker() - case command == "/think": - return m.handleThinkCommand(args) - case command == "/tools": - return m.handleToolsCommand(args) - case command == "/system": - return m.handleSystemCommand(args) - case command == "/skills": - return m.handleSkillsCommand(args) - case command == "/prompt": - return m.handlePromptCommand(args) - case command == "/save": - return m.handleSaveCommand(args) - case command == "/new" && args == "": - return m.resetChat("new chat") - case command == "/compact" && args == "": - return m.startManualCompaction() - case skillOK: - return m.startSkillRun(skillName, skillPrompt) - case strings.HasPrefix(input, "/") && m.slashInputIsMultimodalFile(input): - return m.startRun(input) - case strings.HasPrefix(input, "/"): - m.entries = append(m.entries, newChatEntry(chatEntry{role: "error", content: fmt.Sprintf("Unknown command %q", strings.Fields(input)[0])})) - return *m, nil - } - - return m.startRun(input) -} - -func (m *chatModel) handleSkillsCommand(args string) (tea.Model, tea.Cmd) { - if fields := strings.Fields(args); len(fields) == 2 && fields[0] == "import" { - return m.handleSkillsImport(fields[1]) - } else if len(fields) != 0 { - m.entries = append(m.entries, newChatEntry(chatEntry{role: "error", content: "usage: /skills [import codex|claude|pi]"})) - return *m, nil - } - skills := m.opts.Skills.List() - if len(skills) == 0 { - m.entries = append(m.entries, newSlashEntry("No skills found. Add directories containing SKILL.md under "+skillsDirForDisplay(m.opts.Skills)+".")) - return *m, nil - } - lines := []string{"Available skills:"} - for _, skill := range skills { - description := skill.Description - if description == "" { - description = "No description provided." - } - lines = append(lines, fmt.Sprintf("- `%s`: %s", skill.Name, description)) - } - lines = append(lines, "\nType `/` to load a skill into the conversation.") - m.entries = append(m.entries, newSlashEntry(strings.Join(lines, "\n"))) - return *m, nil -} - -func (m *chatModel) handleSkillsImport(source string) (tea.Model, tea.Cmd) { - importSkills := m.opts.ImportSkills - if importSkills == nil { - importSkills = coreagent.ImportSkills - } - result, err := importSkills(source) - if err != nil { - m.status = "error" - m.entries = append(m.entries, newChatEntry(chatEntry{role: "error", content: fmt.Sprintf("Could not import %s skills: %v", source, err)})) - return *m, nil - } - - if len(result.Imported) != 0 || len(result.Existing) != 0 { - reload := m.opts.ReloadSkills - if reload == nil { - reload = func() (*coreagent.SkillCatalog, error) { - return coreagent.LoadDefaultSkills(m.currentWorkingDir()) - } - } - catalog, err := reload() - if err != nil { - m.status = "error" - m.entries = append(m.entries, newChatEntry(chatEntry{role: "error", content: fmt.Sprintf("%s\n\nCould not reload skills: %v", skillsImportSummary(result), err)})) - return *m, nil - } - m.opts.Skills = catalog - if m.opts.ToolRegistryForModel != nil && m.opts.Model != "" { - m.opts.Tools = m.opts.ToolRegistryForModel(m.ctx, m.opts.Model) - } - if m.opts.SystemPromptForModel != nil { - m.opts.SystemPrompt = m.opts.SystemPromptForModel(m.ctx, m.opts.Model, m.opts.Tools, m.opts.ToolsDisabled) - } - m.status = "skills reloaded" - } - m.entries = append(m.entries, newSlashEntry(skillsImportSummary(result))) - return *m, nil -} - -func skillsImportSummary(result coreagent.SkillImportResult) string { - if len(result.Imported) == 0 && len(result.Existing) == 0 && len(result.Failures) == 0 { - return fmt.Sprintf("No %s skills found at %s.", result.Source, result.SourceDir) - } - var lines []string - if len(result.Imported) != 0 { - lines = append(lines, fmt.Sprintf("Imported %d skill%s from %s.", len(result.Imported), pluralSuffix(len(result.Imported)), result.SourceDir)) - } - if len(result.Existing) != 0 { - lines = append(lines, "Already present (left unchanged): "+strings.Join(result.Existing, ", ")+".") - } - for _, failure := range result.Failures { - lines = append(lines, fmt.Sprintf("Skipped %s: %v.", failure.Name, failure.Err)) - } - return strings.Join(lines, "\n") -} - -func skillsDirForDisplay(catalog *coreagent.SkillCatalog) string { - if catalog != nil && catalog.Dir() != "" { - return catalog.Dir() - } - dir, err := coreagent.SkillsDir() - if err != nil { - return "the Ollama skills directory" - } - return dir -} - -// skillSlashInvocation parses "/" or "/ ". -// It returns the skill name, any trailing prompt, and ok when the first token is -// a catalog skill. Built-in slash commands take precedence over same-named -// skills, so they are never claimed here. -func (m *chatModel) skillSlashInvocation(input string) (name, prompt string, ok bool) { - input = strings.TrimSpace(input) - if !strings.HasPrefix(input, "/") { - return "", "", false - } - token, args, _ := strings.Cut(input, " ") - name = strings.TrimPrefix(token, "/") - if name == "" { - return "", "", false - } - if _, _, known := slashCommandInvocation(input); known { - return "", "", false - } - if _, err := m.opts.Skills.Load(name); err != nil { - return "", "", false - } - return name, strings.TrimSpace(args), true -} - -func (m *chatModel) handleToolsCommand(args string) (tea.Model, tea.Cmd) { - if strings.TrimSpace(args) != "" { - m.status = "error" - m.entries = append(m.entries, newChatEntry(chatEntry{role: "error", content: "usage: /tools"})) - return *m, nil - } - if m.opts.ToolsDisabled { - m.opts.ToolsDisabled = false - if m.opts.ToolRegistryForModel != nil { - m.opts.Tools = m.opts.ToolRegistryForModel(m.ctx, m.opts.Model) - } - m.status = "tools on" - } else { - m.opts.ToolsDisabled = true - m.status = "tools off" - } - if m.opts.SystemPromptForModel != nil { - m.opts.SystemPrompt = m.opts.SystemPromptForModel(m.ctx, m.opts.Model, m.opts.Tools, m.opts.ToolsDisabled) - } - return *m, nil -} - -func (m *chatModel) handleSystemCommand(args string) (tea.Model, tea.Cmd) { - switch strings.ToLower(strings.TrimSpace(args)) { - case "": - m.entries = append(m.entries, newSlashEntry(m.systemCommandOutput())) - case "on": - m.systemPromptDisabled = false - m.status = "system prompt on" - m.entries = append(m.entries, newSlashEntry(m.systemCommandOutput())) - case "off": - m.systemPromptDisabled = true - m.status = "system prompt off" - m.entries = append(m.entries, newSlashEntry(m.systemCommandOutput())) - default: - m.status = "error" - m.entries = append(m.entries, newChatEntry(chatEntry{role: "error", content: "usage: /system [on|off]"})) - } - return *m, nil -} - -func (m chatModel) systemPromptState() string { - if m.systemPromptDisabled { - return "off" - } - return "on" -} - -func (m chatModel) systemCommandOutput() string { - prompt := strings.TrimSpace(m.opts.SystemPrompt) - if prompt == "" { - prompt = "(empty)" - } - return "Built-in system prompt is " + m.systemPromptState() + ".\n\n" + prompt + "\n\nWarning: Changing the system prompt during a session breaks the prompt cache." -} - -func (m chatModel) slashInputIsMultimodalFile(input string) bool { - if !m.opts.MultiModal { - return false - } - fields := strings.Fields(input) - if len(fields) == 0 { - return false - } - for _, file := range filedata.ExtractNames(input) { - if strings.HasPrefix(file, fields[0]) { - return true - } - } - return false -} - -func initialPromptHistory(ctx context.Context, opts Options) []string { - var prompts []string - for _, msg := range opts.Messages { - if msg.Role == "user" { - prompts = append(prompts, msg.Content) - } - } - return normalizePromptHistory(prompts) -} - -func normalizePromptHistory(prompts []string) []string { - history := make([]string, 0, min(len(prompts), maxPromptHistory)) - for _, prompt := range prompts { - prompt = strings.TrimSpace(prompt) - if prompt == "" || coreagent.IsCompactionSummary(api.Message{Role: "user", Content: prompt}) { - continue - } - history = append(history, prompt) - } - if len(history) > maxPromptHistory { - history = history[len(history)-maxPromptHistory:] - } - return history -} - -func (m *chatModel) addPromptHistory(prompt string) { - prompt = strings.TrimSpace(prompt) - if prompt == "" { - return - } - m.promptHistory = append(m.promptHistory, prompt) - if len(m.promptHistory) > maxPromptHistory { - m.promptHistory = m.promptHistory[len(m.promptHistory)-maxPromptHistory:] - } - m.resetPromptHistoryCursor() -} - -func (m *chatModel) movePromptHistory(delta int) bool { - if len(m.promptHistory) == 0 || delta == 0 { - return false - } - if !m.promptActive { - if delta > 0 { - return false - } - m.promptDraft = slices.Clone(m.input) - m.promptCursor = len(m.promptHistory) - 1 - m.promptActive = true - } else { - m.promptCursor += delta - if m.promptCursor >= len(m.promptHistory) { - m.input = slices.Clone(m.promptDraft) - m.inputCursor = len(m.input) - m.inputCursorSet = true - m.inputAttachments = nil - m.inputPastedTexts = nil - m.resetPromptHistoryCursor() - m.complete = 0 - return true - } - if m.promptCursor < 0 { - m.promptCursor = 0 - } - } - - m.inputPastedTexts = nil - input := m.promptHistory[m.promptCursor] - if placeholder, ok := m.pastedTextPlaceholder(input); ok { - input = placeholder - } - m.input = []rune(input) - m.inputCursor = len(m.input) - m.inputCursorSet = true - m.inputAttachments = nil - m.complete = 0 - return true -} - -func (m *chatModel) resetPromptHistoryCursor() { - m.promptActive = false - m.promptCursor = 0 - m.promptDraft = nil -} - -func (m *chatModel) insertInputNewline() { - m.insertInputRunes([]rune{'\n'}) -} - -func (m *chatModel) insertInputRunesFromKey(runes []rune, pasted bool) { - if len(runes) == 0 { - return - } - if m.opts.MultiModal && (pasted || len(runes) > 1) && m.insertInputFilePlaceholders(string(runes)) { - return - } - if pasted && m.insertPastedTextPlaceholder(string(runes)) { - return - } - m.insertInputRunes(runes) -} - -func (m *chatModel) insertInputFilePlaceholders(input string) bool { - cleaned, files, err := filedata.ExtractWithFiles(input) - if err != nil { - m.entries = append(m.entries, newChatEntry(chatEntry{role: "error", content: err.Error(), err: err.Error()})) - m.status = "attachment failed" - return true - } - if len(files) == 0 { - return false - } - - var parts []string - if strings.TrimSpace(cleaned) != "" { - if placeholder, ok := m.pastedTextPlaceholder(cleaned); ok { - parts = append(parts, placeholder) - } else { - parts = append(parts, cleaned) - } - } - for _, file := range files { - kind := filedata.Kind(file.Path) - placeholder := m.nextInputAttachmentPlaceholder(kind) - m.inputAttachments = append(m.inputAttachments, chatInputAttachment{ - placeholder: placeholder, - kind: kind, - data: file.Data, - }) - parts = append(parts, placeholder) - } - m.insertInputRunes([]rune(strings.Join(parts, " "))) - return true -} - -func (m *chatModel) insertPastedTextPlaceholder(input string) bool { - placeholder, ok := m.pastedTextPlaceholder(input) - if !ok { - return false - } - m.insertInputRunes([]rune(placeholder)) - return true -} - -func (m *chatModel) pastedTextPlaceholder(input string) (string, bool) { - if !shouldCollapsePastedText(input) { - return "", false - } - placeholder := m.nextInputPastedTextPlaceholder(input) - m.inputPastedTexts = append(m.inputPastedTexts, chatInputPastedText{ - placeholder: placeholder, - content: input, - }) - return placeholder, true -} - -func shouldCollapsePastedText(input string) bool { - trimmed := strings.TrimSpace(input) - if trimmed == "" { - return false - } - return len([]rune(trimmed)) >= pastedTextPlaceholderMinRunes || pastedTextLineCount(trimmed) >= pastedTextPlaceholderMinLines -} - -func pastedTextLineCount(input string) int { - if input == "" { - return 0 - } - return strings.Count(input, "\n") + 1 -} - -func (m *chatModel) insertInputRunes(runes []rune) { - runes = normalizeInputRunes(runes) - if len(runes) == 0 { - return - } - m.resetPromptHistoryCursor() - m.disarmQuit() - cursor := m.normalizedInputCursor() - next := make([]rune, 0, len(m.input)+len(runes)) - next = append(next, m.input[:cursor]...) - next = append(next, runes...) - next = append(next, m.input[cursor:]...) - m.input = next - m.inputCursor = cursor + len(runes) - m.inputCursorSet = true - m.complete = 0 -} - -func normalizeInputRunes(runes []rune) []rune { - for _, r := range runes { - if r == '\r' { - return normalizeInputRunesSlow(runes) - } - } - return runes -} - -func normalizeInputRunesSlow(runes []rune) []rune { - out := make([]rune, 0, len(runes)) - for i := 0; i < len(runes); i++ { - if runes[i] != '\r' { - out = append(out, runes[i]) - continue - } - out = append(out, '\n') - if i+1 < len(runes) && runes[i+1] == '\n' { - i++ - } - } - return out -} - -func (m *chatModel) deleteInputBackward() { - cursor := m.normalizedInputCursor() - if cursor <= 0 { - return - } - start, end, ok := m.placeholderRangeForBackspace(cursor) - if !ok { - start, end = cursor-1, cursor - } - m.deleteInputRange(start, end) -} - -func (m *chatModel) deleteInputWordBackward() { - cursor := m.normalizedInputCursor() - if cursor <= 0 { - return - } - start, end, ok := m.placeholderRangeForWordDelete(cursor) - if !ok { - start, end = previousInputWordStart(m.input, cursor), cursor - } - m.deleteInputRange(start, end) -} - -func (m *chatModel) deleteInputForward() { - cursor := m.normalizedInputCursor() - if cursor >= len(m.input) { - return - } - start, end, ok := m.placeholderRangeForForwardDelete(cursor) - if !ok { - start, end = cursor, cursor+1 - } - m.deleteInputRange(start, end) -} - -func (m *chatModel) deleteInputRange(start, end int) { - start = clamp(start, 0, len(m.input)) - end = clamp(end, start, len(m.input)) - m.input = append(slices.Clone(m.input[:start]), m.input[end:]...) - m.inputCursor = start - m.inputCursorSet = true - m.complete = 0 - m.syncInputPlaceholders() -} - -func (m chatModel) placeholderRangeForBackspace(cursor int) (int, int, bool) { - cursor = clamp(cursor, 0, len(m.input)) - input := string(m.input) - for _, placeholder := range m.inputPlaceholders() { - if placeholder == "" { - continue - } - start, end, ok := inputPlaceholderRuneRange(input, placeholder) - if ok && cursor > start && cursor <= end { - return start, end, true - } - } - return 0, 0, false -} - -func (m chatModel) placeholderRangeForForwardDelete(cursor int) (int, int, bool) { - cursor = clamp(cursor, 0, len(m.input)) - input := string(m.input) - for _, placeholder := range m.inputPlaceholders() { - if placeholder == "" { - continue - } - start, end, ok := inputPlaceholderRuneRange(input, placeholder) - if ok && cursor >= start && cursor < end { - return start, end, true - } - } - return 0, 0, false -} - -func (m chatModel) placeholderRangeForWordDelete(cursor int) (int, int, bool) { - cursor = clamp(cursor, 0, len(m.input)) - end := cursor - for end > 0 && unicode.IsSpace(m.input[end-1]) { - end-- - } - input := string(m.input) - for _, placeholder := range m.inputPlaceholders() { - if placeholder == "" { - continue - } - start, placeholderEnd, ok := inputPlaceholderRuneRange(input, placeholder) - if ok && end > start && end <= placeholderEnd { - return start, cursor, true - } - } - return 0, 0, false -} - -func (m chatModel) inputPlaceholders() []string { - placeholders := make([]string, 0, len(m.inputAttachments)+len(m.inputPastedTexts)) - for _, attachment := range m.inputAttachments { - placeholders = append(placeholders, attachment.placeholder) - } - for _, pastedText := range m.inputPastedTexts { - placeholders = append(placeholders, pastedText.placeholder) - } - return placeholders -} - -func inputPlaceholderRuneRange(input, placeholder string) (int, int, bool) { - byteStart := strings.Index(input, placeholder) - if byteStart < 0 { - return 0, 0, false - } - start := len([]rune(input[:byteStart])) - end := start + len([]rune(placeholder)) - return start, end, true -} - -func previousInputWordStart(input []rune, cursor int) int { - end := clamp(cursor, 0, len(input)) - for end > 0 && unicode.IsSpace(input[end-1]) { - end-- - } - start := end - for start > 0 && !unicode.IsSpace(input[start-1]) { - start-- - } - return start -} - -func (m *chatModel) syncInputPlaceholders() { - m.inputAttachments = m.activeInputAttachmentsFor(string(m.input)) - m.inputPastedTexts = m.activeInputPastedTextsFor(string(m.input)) -} - -func (m chatModel) activeInputAttachmentsFor(input string) []chatInputAttachment { - if len(m.inputAttachments) == 0 { - return nil - } - active := make([]chatInputAttachment, 0, len(m.inputAttachments)) - for _, attachment := range m.inputAttachments { - if strings.Contains(input, attachment.placeholder) { - active = append(active, attachment) - } - } - return active -} - -func cloneInputAttachments(in []chatInputAttachment) []chatInputAttachment { - return slices.Clone(in) -} - -type chatInputPastedText struct { - placeholder string - content string -} - -func (m chatModel) activeInputPastedTextsFor(input string) []chatInputPastedText { - if len(m.inputPastedTexts) == 0 { - return nil - } - active := make([]chatInputPastedText, 0, len(m.inputPastedTexts)) - for _, pastedText := range m.inputPastedTexts { - if strings.Contains(input, pastedText.placeholder) { - active = append(active, pastedText) - } - } - return active -} - -func cloneInputPastedTexts(in []chatInputPastedText) []chatInputPastedText { - return slices.Clone(in) -} - -func (m chatModel) expandPastedTextPlaceholders(input string) string { - for _, pastedText := range m.activeInputPastedTextsFor(input) { - input = strings.ReplaceAll(input, pastedText.placeholder, pastedText.content) - } - return input -} - -func (m *chatModel) nextInputPastedTextPlaceholder(content string) string { - if m.nextPastedTextID <= 0 { - m.nextPastedTextID = 1 - } - id := m.nextPastedTextID - m.nextPastedTextID++ - return fmt.Sprintf("[Pasted text #%d +%d lines]", id, pastedTextLineCount(strings.TrimSpace(content))) -} - -func (m *chatModel) nextInputAttachmentPlaceholder(kind string) string { - label := inputAttachmentLabel(kind) - switch kind { - case "audio": - id := m.nextAudioID - m.nextAudioID++ - return fmt.Sprintf("[%s #%d]", label, id) - default: - id := m.nextImageID - m.nextImageID++ - return fmt.Sprintf("[%s #%d]", label, id) - } -} - -func inputAttachmentLabel(kind string) string { - switch kind { - case "audio": - return "Audio" - default: - return "Image" - } -} - -var ( - inputAttachmentPlaceholderPattern = regexp.MustCompile(`\[(Image|Audio) #([0-9]+)\]`) - inputPastedTextPlaceholderPattern = regexp.MustCompile(`\[Pasted text #([0-9]+) \+[0-9]+ lines?\]`) -) - -func nextInputAttachmentIDsFromMessages(messages []api.Message) (imageID int, audioID int) { - for _, msg := range messages { - for _, match := range inputAttachmentPlaceholderPattern.FindAllStringSubmatch(msg.Content, -1) { - if len(match) != 3 { - continue - } - id, err := strconv.Atoi(match[2]) - if err != nil { - continue - } - switch match[1] { - case "Image": - imageID = max(imageID, id+1) - case "Audio": - audioID = max(audioID, id+1) - } - } - } - return imageID, audioID -} - -func nextInputPastedTextIDFromMessages(messages []api.Message) int { - nextID := 1 - for _, msg := range messages { - for _, match := range inputPastedTextPlaceholderPattern.FindAllStringSubmatch(msg.Content, -1) { - if len(match) != 2 { - continue - } - id, err := strconv.Atoi(match[1]) - if err != nil { - continue - } - nextID = max(nextID, id+1) - } - } - return nextID -} - -func (m *chatModel) moveInputCursorHorizontal(delta int) bool { - if delta == 0 { - return false - } - cursor := clamp(m.normalizedInputCursor()+delta, 0, len(m.input)) - if cursor == m.normalizedInputCursor() { - return false - } - m.inputCursor = cursor - m.inputCursorSet = true - m.resetPromptHistoryCursor() - m.complete = 0 - return true -} - -func (m *chatModel) moveInputCursorToLineStart() bool { - cursor := m.normalizedInputCursor() - start, _ := inputLineBounds(m.input, cursor) - if start == cursor { - return false - } - m.inputCursor = start - m.inputCursorSet = true - m.resetPromptHistoryCursor() - m.complete = 0 - return true -} - -func (m *chatModel) moveInputCursorToLineEnd() bool { - cursor := m.normalizedInputCursor() - _, end := inputLineBounds(m.input, cursor) - if end == cursor { - return false - } - m.inputCursor = end - m.inputCursorSet = true - m.resetPromptHistoryCursor() - m.complete = 0 - return true -} - -func (m *chatModel) moveInputCursorWord(delta int) bool { - if delta == 0 || len(m.input) == 0 { - return false - } - cursor := m.normalizedInputCursor() - target := cursor - if delta < 0 { - for target > 0 && unicode.IsSpace(m.input[target-1]) { - target-- - } - for target > 0 && !unicode.IsSpace(m.input[target-1]) { - target-- - } - } else { - for target < len(m.input) && !unicode.IsSpace(m.input[target]) { - target++ - } - for target < len(m.input) && unicode.IsSpace(m.input[target]) { - target++ - } - } - if target == cursor { - return false - } - m.inputCursor = target - m.inputCursorSet = true - m.resetPromptHistoryCursor() - m.complete = 0 - return true -} - -func (m *chatModel) handleInputAltRunes(runes []rune) bool { - if len(runes) != 1 { - return false - } - switch runes[0] { - case 'b', 'B': - m.moveInputCursorWord(-1) - return true - case 'f', 'F': - m.moveInputCursorWord(1) - return true - default: - return false - } -} - -func (m *chatModel) moveInputCursorVertical(delta int) bool { - if delta == 0 || len(m.input) == 0 { - return false - } - cursor := m.normalizedInputCursor() - lineStart, lineEnd := inputLineBounds(m.input, cursor) - column := cursor - lineStart - var targetStart, targetEnd int - if delta < 0 { - if lineStart == 0 { - return false - } - targetEnd = lineStart - 1 - targetStart, _ = inputLineBounds(m.input, targetEnd) - } else { - if lineEnd >= len(m.input) { - return false - } - targetStart = lineEnd + 1 - _, targetEnd = inputLineBounds(m.input, targetStart) - } - target := min(targetStart+column, targetEnd) - m.inputCursor = target - m.inputCursorSet = true - m.resetPromptHistoryCursor() - m.complete = 0 - return true -} - -func inputLineBounds(input []rune, cursor int) (int, int) { - cursor = clamp(cursor, 0, len(input)) - start := cursor - for start > 0 && input[start-1] != '\n' { - start-- - } - end := cursor - for end < len(input) && input[end] != '\n' { - end++ - } - return start, end -} - -func (m chatModel) normalizedInputCursor() int { - if !m.inputCursorSet { - return len(m.input) - } - return clamp(m.inputCursor, 0, len(m.input)) -} - -func inputWithCursor(input []rune, cursor int) string { - cursor = clamp(cursor, 0, len(input)) - next := make([]rune, 0, len(input)+1) - next = append(next, input[:cursor]...) - next = append(next, []rune(inputCursorMarker)...) - next = append(next, input[cursor:]...) - return string(next) -} - -func isShiftEnterCSI(msg tea.Msg) bool { - switch fmt.Sprint(msg) { - case "?CSI[49 51 59 50 117]?", // \x1b[13;2u - "?CSI[49 51 59 50 126]?", // \x1b[13;2~ - "?CSI[50 55 59 50 59 49 51 126]?": // \x1b[27;2;13~ - return true - default: - return false - } -} - -func (m chatModel) emptyInputPlaceholder() string { - if m.promptDebug != nil || m.modelPicker != nil || m.thinkPicker != nil || m.approvalPrompt != nil || m.cloudAuthPrompt != nil { - return "" - } - if len(m.entries) > 0 || len(m.messages) > 0 || len(m.input) > 0 || len(m.inputAttachments) > 0 || len(m.inputPastedTexts) > 0 { - return "" - } - return m.emptyChatHint() -} - -func renderInputBoxLines(input string, cursor int, width, maxBodyLines int, placeholder string) []string { - if width < 12 { - width = 12 - } - if maxBodyLines < 1 { - maxBodyLines = 1 - } - - prefix := chatPromptPrefix - continuationPrefix := strings.Repeat(" ", lipgloss.Width(prefix)) - prefixWidth := lipgloss.Width(prefix) - innerWidth := max(1, width-2) - contentWidth := max(1, innerWidth-inputBoxHorizontalPadding*2) - bodyWidth := max(1, contentWidth-prefixWidth) - - var raw []string - placeholderMode := input == "" && strings.TrimSpace(placeholder) != "" - if input == "" && strings.TrimSpace(placeholder) != "" { - raw = renderInputPromptRawLines(inputWithCursor([]rune(placeholder), 0), prefix, continuationPrefix, bodyWidth) - } else { - if cursor >= 0 { - input = inputWithCursor([]rune(input), cursor) - } - raw = renderInputPromptRawLines(input, prefix, continuationPrefix, bodyWidth) - } - if len(raw) > maxBodyLines { - raw = slices.Clone(raw[len(raw)-maxBodyLines:]) - raw[0] = truncateInputLine(continuationPrefix+trimInputPromptPrefix(raw[0]), contentWidth) - } - - lines := make([]string, 0, len(raw)+2) - lines = append(lines, chatInputBorderStyle.Render(inputBoxTopBorderLine(width))) - for i, line := range raw { - rendered := line - if placeholderMode { - if i == 0 && strings.HasPrefix(line, prefix) { - rest := strings.TrimPrefix(line, prefix) - if strings.HasPrefix(rest, inputCursorMarker) { - rendered = chatUserStyle.Render(prefix) + renderInputTextWithCursorStyle(rest, chatInputPlaceholderStyle) - lines = append(lines, renderInputBoxBodyLine(rendered, contentWidth)) - continue - } - rendered = chatUserStyle.Render(prefix) + chatInputPlaceholderStyle.Render(rest) - lines = append(lines, renderInputBoxBodyLine(rendered, contentWidth)) - continue - } - if strings.HasPrefix(line, continuationPrefix) { - rendered = chatInputPlaceholderStyle.Render(continuationPrefix + strings.TrimPrefix(line, continuationPrefix)) - lines = append(lines, renderInputBoxBodyLine(rendered, contentWidth)) - continue - } - rendered = chatInputPlaceholderStyle.Render(line) - lines = append(lines, renderInputBoxBodyLine(rendered, contentWidth)) - continue - } - lines = append(lines, renderInputBoxBodyLine(renderInputTextWithCursor(rendered), contentWidth)) - } - lines = append(lines, chatInputBorderStyle.Render(inputBoxBottomBorderLine(width))) - return lines -} - -func renderInputTextWithCursor(line string) string { - return renderInputTextWithCursorStyle(line, chatUserStyle) -} - -func renderInputTextWithCursorStyle(line string, style lipgloss.Style) string { - before, after, ok := strings.Cut(line, inputCursorMarker) - if !ok { - return style.Render(line) - } - cell, rest := inputCursorCell(after) - return style.Render(before) + renderInputCursorCell(cell) + style.Render(rest) -} - -func inputCursorCell(after string) (string, string) { - if after == "" { - return "", "" - } - r, size := utf8.DecodeRuneInString(after) - if r == '\n' { - return "", after - } - return after[:size], after[size:] -} - -func renderInputCursorCell(cell string) string { - if cell == "" { - return chatBlankCursorStyle.Render(inputCursorGlyph) - } - return chatCursorStyle.Render(cell) -} - -func renderInputPromptRawLines(text, prefix, continuationPrefix string, width int) []string { - if width <= 0 { - width = 1 - } - text = string(normalizeInputRunes([]rune(text))) - body := wrapChatText(text, width) - for i, line := range body { - if i == 0 { - body[i] = prefix + line - continue - } - body[i] = continuationPrefix + line - } - if len(body) == 0 { - return []string{prefix} - } - return body -} - -func trimInputPromptPrefix(line string) string { - for _, prefix := range []string{chatPromptPrefix, strings.Repeat(" ", lipgloss.Width(chatPromptPrefix))} { - if prefix == "" { - continue - } - if strings.HasPrefix(line, prefix) { - return strings.TrimPrefix(line, prefix) - } - } - return strings.TrimSpace(line) -} - -func inputBoxTopBorderLine(width int) string { - return inputBoxBorderLine(width, "╭", "╮") -} - -func inputBoxBottomBorderLine(width int) string { - return inputBoxBorderLine(width, "╰", "╯") -} - -func inputBoxBorderLine(width int, left, right string) string { - if width < 4 { - width = 4 - } - return left + strings.Repeat("─", max(0, width-2)) + right -} - -func renderInputBoxBodyLine(line string, width int) string { - padding := strings.Repeat(" ", inputBoxHorizontalPadding) - return chatInputBorderStyle.Render("│") + padding + padRenderedLine(line, width) + padding + chatInputBorderStyle.Render("│") -} - -func padRenderedLine(line string, width int) string { - if width <= 0 { - return line - } - if renderedWidth := lipgloss.Width(line); renderedWidth < width { - return line + strings.Repeat(" ", width-renderedWidth) - } - return line -} - -func truncateInputLine(line string, width int) string { - if width <= 0 { - return line - } - if runewidth.StringWidth(line) <= width { - return line - } - return runewidth.Truncate(line, width, "") -} - -func (m chatModel) slashCommandLines(width int) []string { - return m.renderCompletions(m.slashCompletions(), width) -} - -func (m chatModel) completionLines(width int) []string { - return m.renderCompletions(m.completions(), width) -} - -func (m chatModel) renderCompletions(completions []chatCompletion, width int) []string { - if len(completions) == 0 { - return nil - } - selected := clamp(m.complete, 0, len(completions)-1) - start, end := completionWindow(len(completions), selected, m.completionVisibleLimit(len(completions))) - completions = completions[start:end] - - nameWidth := 0 - for _, completion := range completions { - nameWidth = max(nameWidth, lipgloss.Width(completion.label)) - } - - lines := make([]string, 0, len(completions)) - for i, completion := range completions { - marker := " " - if start+i == selected { - marker = "› " - } - name := chatCommandNameStyle.Render(completion.label) - padding := strings.Repeat(" ", max(1, nameWidth-lipgloss.Width(completion.label)+2)) - line := marker + name + padding + chatMetaStyle.Render(completion.description) - lines = append(lines, truncateRenderedLine(line, width)) - } - return lines -} - -func (m chatModel) completionVisibleLimit(total int) int { - if strings.HasPrefix(strings.TrimSpace(string(m.input)), "/") { - return min(maxSlashCompletions, total) - } - return total -} - -func completionWindow(total, selected, limit int) (int, int) { - if total <= 0 || limit <= 0 || limit >= total { - return 0, total - } - selected = clamp(selected, 0, total-1) - start := selected - limit + 1 - if start < 0 { - start = 0 - } - end := start + limit - if end > total { - end = total - start = max(0, end-limit) - } - return start, end -} - -func (m chatModel) completions() []chatCompletion { - if completions := m.slashCompletions(); len(completions) > 0 { - return completions - } - return m.mentionCompletions() -} - -func (m chatModel) slashCompletions() []chatCompletion { - rawInput := string(m.input) - input := strings.TrimLeftFunc(rawInput, unicode.IsSpace) - if !strings.HasPrefix(input, "/") { - return nil - } - if argument, ok := systemCommandArgument(rawInput); ok { - return systemCommandCompletions(argument) - } - if m.skillSlashPromptStarted(rawInput) { - return nil - } - if completions := matchingSkillsImportCompletions(input); completions != nil { - return completions - } - - commands := matchingSlashCommands(input) - completions := make([]chatCompletion, 0, len(commands)) - for _, command := range commands { - completions = append(completions, chatCompletion{ - value: command.name, - label: command.name, - description: command.description, - }) - } - if strings.EqualFold(input, "/skills") { - completions = append(completions, chatCompletion{ - value: "/skills import", - label: "/skills import", - description: "import skills from Codex, Claude, or Pi", - }) - } - // Each catalog skill is also invocable as "/"; surface them as - // completions so they are discoverable by typing. - if m.opts.Skills != nil { - prefix := strings.ToLower(input) - for _, skill := range m.opts.Skills.List() { - name := "/" + skill.Name - if !strings.HasPrefix(name, prefix) { - continue - } - if _, _, known := slashCommandInvocation(name); known { - continue // built-in command wins; don't shadow it - } - description := skill.Description - if description == "" { - description = "No description provided." - } - completions = append(completions, chatCompletion{ - value: name, - label: name, - description: description, - }) - } - } - if len(completions) == 0 { - return []chatCompletion{{label: "No matching commands"}} - } - return completions -} - -func matchingSkillsImportCompletions(input string) []chatCompletion { - const importCommand = "/skills import" - lower := strings.ToLower(input) - if lower == "/skills" { - return nil // Preserve Enter on /skills as the listing command. - } - if !strings.HasPrefix(lower, "/skills ") { - return nil - } - if strings.HasPrefix(importCommand, lower) { - return []chatCompletion{{ - value: importCommand, - label: importCommand, - description: "import skills from Codex, Claude, or Pi", - }} - } - if !strings.HasPrefix(lower, importCommand) { - return nil - } - prefix := strings.TrimSpace(strings.TrimPrefix(lower, importCommand)) - completions := make([]chatCompletion, 0, len(skillsImportCompletions)) - for _, completion := range skillsImportCompletions { - if strings.HasPrefix(strings.TrimPrefix(completion.value, importCommand+" "), prefix) { - completions = append(completions, completion) - } - } - if len(completions) == 0 { - return []chatCompletion{{label: "No matching skill sources"}} - } - return completions -} - -func hasSystemCommandArgument(input string) bool { - _, ok := systemCommandArgument(input) - return ok -} - -func systemCommandArgument(input string) (string, bool) { - input = strings.TrimLeftFunc(input, unicode.IsSpace) - end := strings.IndexFunc(input, unicode.IsSpace) - if end < 0 { - return "", false - } - command, _, known := slashCommandInvocation(input[:end]) - if !known || command != "/system" { - return "", false - } - return strings.TrimSpace(input[end:]), true -} - -func systemCommandCompletions(argument string) []chatCompletion { - argument = strings.ToLower(argument) - options := []chatCompletion{ - {value: "/system on", label: "on", description: "enable the built-in system prompt"}, - {value: "/system off", label: "off", description: "disable the built-in system prompt"}, - } - completions := make([]chatCompletion, 0, len(options)) - for _, option := range options { - if strings.HasPrefix(option.label, argument) { - completions = append(completions, option) - } - } - if len(completions) == 0 { - return []chatCompletion{{label: "No matching options"}} - } - return completions -} - -func (m chatModel) skillSlashPromptStarted(input string) bool { - input = strings.TrimLeftFunc(input, unicode.IsSpace) - end := strings.IndexFunc(input, unicode.IsSpace) - if end < 0 { - return false - } - _, _, ok := m.skillSlashInvocation(input[:end]) - return ok -} - -func matchingSlashCommands(input string) []chatSlashCommand { - prefix := strings.ToLower(strings.TrimSpace(input)) - if prefix == "" { - return nil - } - - var commands []chatSlashCommand - for _, command := range chatSlashCommands { - if command.hidden { - continue - } - if command.matchesPrefix(prefix) { - commands = append(commands, command) - } - } - return commands -} - -func (c chatSlashCommand) matchesPrefix(prefix string) bool { - if strings.HasPrefix(c.name, prefix) { - return true - } - for _, alias := range c.aliases { - if strings.HasPrefix(alias, prefix) { - return true - } - } - return false -} - -func slashCommandInvocation(input string) (string, string, bool) { - input = strings.TrimSpace(input) - if !strings.HasPrefix(input, "/") { - return "", "", false - } - token, args, _ := strings.Cut(input, " ") - token = strings.ToLower(token) - for _, command := range chatSlashCommands { - if command.name == token || slices.Contains(command.aliases, token) { - return command.name, strings.TrimSpace(args), true - } - } - return "", "", false -} - -func (m chatModel) mentionCompletions() []chatCompletion { - _, query, ok := activeMentionToken(m.input, m.normalizedInputCursor()) - if !ok { - return nil - } - - workingDir := m.currentWorkingDir() - if strings.TrimSpace(workingDir) == "" { - var err error - workingDir, err = os.Getwd() - if err != nil { - return []chatCompletion{{label: "No working directory"}} - } - } - - dirPart, prefix := splitMentionQuery(query) - dir, err := resolveCompletionDir(workingDir, dirPart) - if err != nil { - return []chatCompletion{{label: "No matching files"}} - } - entries, err := os.ReadDir(dir) - if err != nil { - return []chatCompletion{{label: "No matching files"}} - } - sort.SliceStable(entries, func(i, j int) bool { - if entries[i].IsDir() != entries[j].IsDir() { - return entries[i].IsDir() - } - return strings.ToLower(entries[i].Name()) < strings.ToLower(entries[j].Name()) - }) - - includeHidden := strings.HasPrefix(prefix, ".") - completions := make([]chatCompletion, 0, 8) - for _, entry := range entries { - name := entry.Name() - if !includeHidden && strings.HasPrefix(name, ".") { - continue - } - if !strings.HasPrefix(strings.ToLower(name), strings.ToLower(prefix)) { - continue - } - value := filepath.ToSlash(filepath.Join(dirPart, name)) - label := "@" + value - description := "file" - if entry.IsDir() { - value += "/" - label += "/" - description = "directory" - } - completions = append(completions, chatCompletion{ - value: value, - label: label, - description: description, - directory: entry.IsDir(), - }) - if len(completions) >= 8 { - break - } - } - if len(completions) == 0 { - return []chatCompletion{{label: "No matching files"}} - } - return completions -} - -func activeMentionToken(input []rune, cursor int) (int, string, bool) { - cursor = clamp(cursor, 0, len(input)) - start := cursor - for start > 0 && !unicode.IsSpace(input[start-1]) { - start-- - } - token := string(input[start:cursor]) - if !strings.HasPrefix(token, "@") { - return 0, "", false - } - return start, token[1:], true -} - -func splitMentionQuery(query string) (string, string) { - query = filepath.ToSlash(query) - index := strings.LastIndex(query, "/") - if index < 0 { - return ".", query - } - return query[:index+1], query[index+1:] -} - -func resolveCompletionDir(workingDir, dir string) (string, error) { - if filepath.IsAbs(dir) { - return "", fmt.Errorf("absolute paths are not allowed") - } - base := workingDir - if base == "" { - base = "." - } - base, err := filepath.Abs(base) - if err != nil { - return "", err - } - resolved := filepath.Clean(filepath.Join(base, dir)) - rel, err := filepath.Rel(base, resolved) - if err != nil { - return "", err - } - if rel == ".." || strings.HasPrefix(rel, ".."+string(filepath.Separator)) { - return "", fmt.Errorf("path escapes working directory") - } - return resolved, nil -} - -func (m *chatModel) moveCompletion(delta int) bool { - completions := m.completions() - if len(completions) == 0 || !completionIsSelectable(completions) { - return false - } - m.complete = (m.complete + delta) % len(completions) - if m.complete < 0 { - m.complete += len(completions) - } - return true -} - -func (m *chatModel) applyCompletion() bool { - completions := m.completions() - if len(completions) == 0 || !completionIsSelectable(completions) { - return false - } - m.resetPromptHistoryCursor() - selected := completions[clamp(m.complete, 0, len(completions)-1)] - cursor := m.normalizedInputCursor() - input := string(m.input) - if strings.HasPrefix(strings.TrimSpace(input), "/") { - m.input = []rune(selected.value) - m.inputCursor = len(m.input) - m.inputCursorSet = true - m.complete = 0 - return true - } - - start, _, ok := activeMentionToken(m.input, cursor) - if !ok { - return false - } - completed := []rune("@" + selected.value) - if !selected.directory && (cursor == len(m.input) || !unicode.IsSpace(m.input[cursor])) { - completed = append(completed, ' ') - } - next := make([]rune, 0, len(m.input)-cursor+start+len(completed)) - next = append(next, m.input[:start]...) - next = append(next, completed...) - next = append(next, m.input[cursor:]...) - m.input = next - m.inputCursor = start + len(completed) - if !selected.directory && m.inputCursor < len(m.input) && unicode.IsSpace(m.input[m.inputCursor]) { - m.inputCursor++ - } - m.inputCursorSet = true - m.complete = 0 - return true -} - -func (m *chatModel) applyMentionCompletion() bool { - if strings.HasPrefix(strings.TrimSpace(string(m.input)), "/") { - return false - } - return m.applyCompletion() -} - -func completionIsSelectable(completions []chatCompletion) bool { - return len(completions) > 0 && completions[0].value != "" -} - -func (m chatModel) helpSummary() string { - lines := []string{ - "**Commands**", - "", - } - for _, command := range chatSlashCommands { - if command.hidden || strings.TrimSpace(command.description) == "" { - continue - } - usage := command.name - if command.usage != "" { - usage = command.usage - } - lines = append(lines, fmt.Sprintf("- `%s`: %s", usage, command.description)) - } - lines = append(lines, - "", - "**Shortcuts**", - "", - "- `shift+enter`: insert a newline", - "- `shift+tab`: toggle permission mode", - "- `ctrl+o`: toggle transcript details", - "- `↑/↓`: previous or next prompt", - "- `ctrl+a/e`: move to line start or end", - ) - return strings.Join(lines, "\n") -} - -func (m chatModel) systemPrompt(extra string) string { - var parts []string - if !m.systemPromptDisabled && strings.TrimSpace(m.opts.SystemPrompt) != "" { - parts = append(parts, strings.TrimSpace(m.opts.SystemPrompt)) - } - if strings.TrimSpace(extra) != "" { - parts = append(parts, strings.TrimSpace(extra)) - } - return strings.Join(parts, "\n\n") -} diff --git a/cmd/tui/chat/input_test.go b/cmd/tui/chat/input_test.go deleted file mode 100644 index 70e9d6fabae..00000000000 --- a/cmd/tui/chat/input_test.go +++ /dev/null @@ -1,1351 +0,0 @@ -package chat - -import ( - "context" - "os" - "path/filepath" - "strings" - "testing" - - tea "github.com/charmbracelet/bubbletea" - "github.com/charmbracelet/lipgloss" - - coreagent "github.com/ollama/ollama/agent" - "github.com/ollama/ollama/api" -) - -func TestChatHelpCommandShowsV1Commands(t *testing.T) { - m := chatModel{input: []rune("/help")} - - updated, cmd := m.handleSubmit() - if cmd != nil { - t.Fatal("help command should not return a command") - } - - fm := updated.(chatModel) - if len(fm.entries) != 1 { - t.Fatalf("entries = %d, want 1", len(fm.entries)) - } - for _, want := range []string{ - "**Commands**", - "- `/model`: switch models", - "- `/think`: set thinking mode", - "- `/system [on|off]`: show or set the built-in system prompt", - "- `/compact`: summarize older context", - "- `/help`: show commands", - "- `/bye`: exit", - "- `/prompt`: show full prompt, tools, and messages", - "- `/save `: save request JSON; saved as .json", - "**Shortcuts**", - "- `shift+enter`: insert a newline", - "- `shift+tab`: toggle permission mode", - "- `ctrl+o`: toggle transcript details", - } { - if !strings.Contains(fm.entries[0].content, want) { - t.Fatalf("help output missing %q:\n%s", want, fm.entries[0].content) - } - } - for _, removed := range []string{"/history", "/load", "/raw", "/resume", "/set", "/show", "/verbose"} { - if strings.Contains(fm.entries[0].content, removed) { - t.Fatalf("removed command %q should stay hidden from help:\n%s", removed, fm.entries[0].content) - } - } -} - -func TestChatNewCommandRepaintsFromTop(t *testing.T) { - m := chatModel{ - input: []rune("/new"), - flowPrintedLines: 4, - entries: []chatEntry{{role: "assistant", content: "old transcript"}}, - messages: []api.Message{{Role: "user", Content: "old prompt"}}, - approvalState: testApprovalState(true, map[string]bool{"edit": true}), - opts: Options{AllowAllTools: true}, - permissionNotice: "full access enabled", - } - - updated, cmd := m.handleSubmit() - if cmd == nil { - t.Fatal("/new should return a repaint command") - } - m = updated.(chatModel) - if m.status != "new chat" { - t.Fatalf("status = %q, want new chat", m.status) - } - if len(m.entries) != 0 || len(m.messages) != 0 { - t.Fatalf("chat was not reset: entries=%#v messages=%#v", m.entries, m.messages) - } - if m.flowPrintedLines != 0 { - t.Fatalf("flowPrintedLines = %d, want 0", m.flowPrintedLines) - } - if m.approvalState.AllGranted() || m.opts.AllowAllTools || m.approvalState.Allows("edit") || m.permissionNotice != "" { - t.Fatalf("permissions were not reset: allowAll=%v opts=%v editAllowed=%v notice=%q", m.approvalState.AllGranted(), m.opts.AllowAllTools, m.approvalState.Allows("edit"), m.permissionNotice) - } - if msg := cmd(); msg == nil { - t.Fatal("repaint command returned nil") - } -} - -func TestChatNewCommandPreservesLaunchFullAccessDefault(t *testing.T) { - m := chatModel{ - input: []rune("/new"), - defaultAllowAll: true, - approvalState: testApprovalState(false, map[string]bool{"edit": true}), - } - - updated, _ := m.handleSubmit() - fm := updated.(chatModel) - if !fm.approvalState.AllGranted() || !fm.opts.AllowAllTools { - t.Fatalf("full access default was not restored: allowAll=%v opts=%v", fm.approvalState.AllGranted(), fm.opts.AllowAllTools) - } - fm.approvalState.Set(false, nil) - if fm.approvalState.Allows("edit") { - t.Fatal("edit scope should be cleared") - } -} - -func TestChatSaveCommandRequiresFilename(t *testing.T) { - m := chatModel{ - input: []rune("/save"), - opts: Options{Model: "llama3.2"}, - } - - updated, cmd := m.handleSubmit() - if cmd != nil { - t.Fatal("/save should not start a command") - } - fm := updated.(chatModel) - if fm.status != "error" { - t.Fatalf("status = %q, want error", fm.status) - } - if len(fm.entries) != 1 || !strings.Contains(fm.entries[0].content, "usage: /save ") { - t.Fatalf("entries = %#v, want usage error", fm.entries) - } -} - -func TestChatSaveCommandWritesRequestJSON(t *testing.T) { - dir := t.TempDir() - m := chatModel{ - input: []rune("/save request"), - workingDir: dir, - opts: Options{ - Model: "llama3.2", - SystemPrompt: "You are Ollama.", - }, - messages: []api.Message{{Role: "user", Content: "hello"}}, - } - - updated, cmd := m.handleSubmit() - if cmd != nil { - t.Fatal("/save redirect should not start a command") - } - fm := updated.(chatModel) - if fm.status != "saved" { - t.Fatalf("status = %q, want saved", fm.status) - } - - data, err := os.ReadFile(filepath.Join(dir, "request.json")) - if err != nil { - t.Fatal(err) - } - raw := string(data) - for _, want := range []string{ - `"model": "llama3.2"`, - `"role": "system"`, - `"content": "You are Ollama."`, - `"role": "user"`, - `"content": "hello"`, - } { - if !strings.Contains(raw, want) { - t.Fatalf("saved request missing %q:\n%s", want, raw) - } - } - if got := fm.entries[len(fm.entries)-1].content; got != "saved as request.json" { - t.Fatalf("save entry = %q, want saved filename", got) - } -} - -func TestChatSaveCommandRejectsPath(t *testing.T) { - dir := t.TempDir() - m := chatModel{ - input: []rune("/save ../request"), - workingDir: dir, - opts: Options{Model: "llama3.2"}, - } - - updated, cmd := m.handleSubmit() - if cmd != nil { - t.Fatal("invalid /save should not start a command") - } - fm := updated.(chatModel) - if fm.status != "error" { - t.Fatalf("status = %q, want error", fm.status) - } - if _, err := os.Stat(filepath.Join(filepath.Dir(dir), "request.json")); !os.IsNotExist(err) { - t.Fatalf("/save wrote outside working dir, stat err = %v", err) - } -} - -func TestChatPromptCommandOpensPromptDebugScreen(t *testing.T) { - registry := &coreagent.Registry{} - registry.Register(chatTestTool{}) - toolArgs := api.NewToolCallFunctionArguments() - toolArgs.Set("query", "show me everything") - m := chatModel{ - input: []rune("/prompt"), - width: 100, - height: 20, - opts: Options{ - Model: "llama3.2", - SystemPrompt: "You are Ollama.", - Tools: registry, - ContextWindowTokens: 1024, - }, - messages: []api.Message{ - {Role: "user", Content: "hello\nsecond line"}, - { - Role: "assistant", - Content: "I'll call a tool.", - Thinking: "I should call the fake tool first.", - ToolCalls: []api.ToolCall{{ - ID: "call-1", - Function: api.ToolCallFunction{ - Name: "fake_tool", - Arguments: toolArgs, - }, - }}, - }, - {Role: "tool", ToolName: "fake_tool", ToolCallID: "call-1", Content: "tool result line 1\ntool result line 2"}, - }, - } - - updated, cmd := m.handleSubmit() - if cmd == nil { - t.Fatal("/prompt should enable managed prompt screen") - } - fm := updated.(chatModel) - if fm.promptDebug == nil { - t.Fatal("/prompt should open prompt debug screen") - } - if len(fm.entries) != 0 { - t.Fatalf("/prompt should not append a transcript entry: %#v", fm.entries) - } - out := stripANSI(fm.View()) - body := stripANSI(strings.Join(fm.promptDebugLines(160), "\n")) - for _, want := range []string{ - "Prompt", - "full request preview", - } { - if !strings.Contains(out, want) { - t.Fatalf("/prompt view missing %q:\n%s", want, out) - } - } - for _, want := range []string{ - "model: llama3.2", - "estimated prompt:", - "/ 1024 tokens", - "messages: 4", - "tools: 1", - "Tools", - "1. fake_tool", - "description:", - "does test work", - "parameters: object", - "1. system", - "You are Ollama.", - "2. user", - "hello", - "second line", - "3. assistant", - "thinking:", - "I should call the fake tool first.", - "content:", - "I'll call a tool.", - "tool call 1: fake_tool", - "id: call-1", - "arguments:", - "query: show me everything", - "4. tool:fake_tool", - "tool_name: fake_tool", - "tool_call_id: call-1", - "tool result", - "tool result line 1", - "tool result line 2", - } { - if !strings.Contains(body, want) { - t.Fatalf("/prompt output missing %q:\n%s", want, body) - } - } - assistantStart := strings.Index(body, "3. assistant") - if assistantStart < 0 { - t.Fatalf("/prompt output missing assistant message:\n%s", body) - } - assistantBody := body[assistantStart:] - thinkingIndex := strings.Index(assistantBody, "thinking:") - contentIndex := strings.Index(assistantBody, "content:") - if thinkingIndex < 0 || contentIndex < 0 { - t.Fatalf("/prompt assistant output missing thinking/content labels:\n%s", body) - } - if thinkingIndex > contentIndex { - t.Fatalf("/prompt assistant thinking should render before content:\n%s", body) - } - for _, unwanted := range []string{`"query":`, `"name": "fake_tool"`} { - if strings.Contains(body, unwanted) { - t.Fatalf("/prompt output should be rendered, not raw JSON; found %q:\n%s", unwanted, body) - } - } - messagesIndex := strings.Index(body, "Messages") - toolsIndex := strings.Index(body, "Tools") - if messagesIndex < 0 || toolsIndex < 0 || toolsIndex < messagesIndex { - t.Fatalf("/prompt should render tools after messages:\n%s", body) - } - - updated, cmd = fm.Update(tea.KeyMsg{Type: tea.KeyEsc}) - if cmd == nil { - t.Fatal("closing prompt debug should disable prompt mouse mode") - } - fm = updated.(chatModel) - if fm.promptDebug != nil { - t.Fatal("esc should close prompt debug screen") - } -} - -func TestChatPromptDebugCapsToolResultPreview(t *testing.T) { - longResult := strings.Repeat("x", maxPromptDebugToolResultRunes+25) + "tail-marker" - m := chatModel{ - width: 120, - height: 20, - promptDebug: &chatPromptDebug{ - request: api.ChatRequest{ - Model: "llama3.2", - Messages: []api.Message{{ - Role: "tool", - ToolName: "bash", - ToolCallID: "call-1", - Content: longResult, - }}, - }, - }, - } - - body := stripANSI(strings.Join(m.promptDebugLines(120), "\n")) - if strings.Contains(body, "tail-marker") { - t.Fatalf("/prompt should cap rendered tool results:\n%s", body) - } - if !strings.Contains(body, "...") { - t.Fatalf("/prompt capped tool result should show ellipsis:\n%s", body) - } - if got := strings.Count(body, "x"); got != maxPromptDebugToolResultRunes-3 { - t.Fatalf("rendered tool result x count = %d, want %d:\n%s", got, maxPromptDebugToolResultRunes-3, body) - } - if got := m.promptDebug.request.Messages[0].Content; got != longResult { - t.Fatal("/prompt rendering should not mutate the request") - } -} - -func TestChatPromptDebugMouseWheelScrolls(t *testing.T) { - m := chatModel{ - width: 80, - height: 8, - opts: Options{Model: "llama3.2", ContextWindowTokens: 1024}, - promptDebug: &chatPromptDebug{ - request: api.ChatRequest{Model: "llama3.2"}, - tokens: 10, - }, - } - for range 30 { - m.promptDebug.request.Messages = append(m.promptDebug.request.Messages, api.Message{Role: "user", Content: "line"}) - } - if m.promptDebugMaxScroll() == 0 { - t.Fatal("test setup should produce a scrollable prompt debug screen") - } - - updated, _ := m.Update(tea.MouseMsg{Type: tea.MouseWheelDown}) - m = updated.(chatModel) - if m.promptDebug.scroll == 0 { - t.Fatal("mouse wheel down should scroll prompt debug screen") - } - - updated, _ = m.Update(tea.MouseMsg{Type: tea.MouseWheelUp}) - m = updated.(chatModel) - if m.promptDebug.scroll != 0 { - t.Fatalf("mouse wheel up should return prompt debug screen to top, got scroll %d", m.promptDebug.scroll) - } -} - -func TestChatPromptDebugCachesLinesByWidth(t *testing.T) { - m := chatModel{ - promptDebug: &chatPromptDebug{ - request: api.ChatRequest{ - Model: "llama3.2", - Messages: []api.Message{{ - Role: "user", - Content: strings.Repeat("a long prompt line ", 20), - }}, - }, - }, - } - - first := m.promptDebugLines(80) - if len(first) == 0 || m.promptDebug.linesWidth != 80 { - t.Fatalf("prompt cache = %#v, want lines cached at width 80", m.promptDebug) - } - if &first[0] != &m.promptDebugLines(80)[0] { - t.Fatal("prompt debug should reuse cached lines at the same width") - } - - resized := m.promptDebugLines(120) - if m.promptDebug.linesWidth != 120 { - t.Fatalf("prompt cache width = %d, want 120", m.promptDebug.linesWidth) - } - if &first[0] == &resized[0] { - t.Fatal("prompt debug should rebuild lines after a width change") - } -} - -func TestTruncateInputLineUsesDisplayWidth(t *testing.T) { - line := truncateInputLine(strings.Repeat("界", 10), 10) - if got := lipgloss.Width(line); got > 10 { - t.Fatalf("line %q width = %d, want <= 10", line, got) - } -} - -func TestRenderInputBoxTruncationUsesSingleContinuationMarker(t *testing.T) { - lines := renderInputBoxLines("one two three four five six seven", len("one two three four five six seven"), 16, 1, "") - rendered := strings.Join(lines, "\n") - if strings.Contains(rendered, "... ...") { - t.Fatalf("input rendered duplicate continuation marker: %q", rendered) - } - if strings.Contains(rendered, "one two") { - t.Fatalf("input should keep the latest truncated line: %q", rendered) - } -} - -type shiftEnterCSITestMsg string - -func (m shiftEnterCSITestMsg) String() string { - return string(m) -} - -func TestChatInputHandlesShiftEnterCSIMessage(t *testing.T) { - m := chatModel{input: []rune("line one")} - - updated, _ := m.Update(shiftEnterCSITestMsg("?CSI[49 51 59 50 117]?")) - m = updated.(chatModel) - if got := string(m.input); got != "line one\n" { - t.Fatalf("input = %q, want newline inserted", got) - } -} - -func TestChatInputAcceptsSpace(t *testing.T) { - m := chatModel{} - - updated, _ := m.Update(tea.KeyMsg{Type: tea.KeyRunes, Runes: []rune("hello")}) - m = updated.(chatModel) - updated, _ = m.Update(tea.KeyMsg{Type: tea.KeySpace, Runes: []rune(" ")}) - m = updated.(chatModel) - updated, _ = m.Update(tea.KeyMsg{Type: tea.KeyRunes, Runes: []rune("world")}) - m = updated.(chatModel) - - if got := string(m.input); got != "hello world" { - t.Fatalf("input = %q, want hello world", got) - } -} - -func TestChatCloudModelDefaultToolRoundsAreUnlimited(t *testing.T) { - const formerDefaultLimit = 100 - client := &chatToolLoopClient{toolRounds: formerDefaultLimit + 1} - registry := &coreagent.Registry{} - registry.Register(chatTestTool{}) - m := chatModel{ - ctx: context.Background(), - opts: Options{ - Model: "test:cloud", - Client: client, - Tools: registry, - AllowAllTools: true, - }, - } - - updated, cmd := m.startRun("keep going") - m = updated.(chatModel) - if cmd == nil { - t.Fatal("startRun should start a cloud model run") - } - done := waitForRunDone(t, m.events) - if done.err != nil { - t.Fatalf("cloud run returned error: %v", done.err) - } - if client.calls != formerDefaultLimit+2 { - t.Fatalf("client calls = %d, want %d", client.calls, formerDefaultLimit+2) - } -} - -func TestChatLargePasteUsesPlaceholderAndExpandsOnSubmit(t *testing.T) { - pasted := strings.Repeat("line\n", pastedTextPlaceholderMinLines-1) + "line" - m := chatModel{ - ctx: context.Background(), - opts: Options{ - Model: "test", - Client: chatTestClient{}, - }, - } - - updated, _ := m.Update(tea.KeyMsg{Type: tea.KeyRunes, Runes: []rune(pasted), Paste: true}) - m = updated.(chatModel) - if got, want := string(m.input), "[Pasted text #1 +8 lines]"; got != want { - t.Fatalf("input = %q, want %q", got, want) - } - - updated, cmd := m.Update(tea.KeyMsg{Type: tea.KeyEnter}) - m = updated.(chatModel) - if cmd == nil { - t.Fatal("enter should start a run") - } - done := waitForRunDone(t, m.events) - if done.err != nil { - t.Fatal(done.err) - } - if done.result == nil || len(done.result.Messages) < 1 || done.result.Messages[0].Content != pasted { - t.Fatalf("messages = %#v, want expanded pasted text", done.result) - } -} - -func TestChatBackspaceDeletesWholePastedTextPlaceholder(t *testing.T) { - m := chatModel{ - input: []rune("use [Pasted text #1 +8 lines]"), - inputPastedTexts: []chatInputPastedText{{ - placeholder: "[Pasted text #1 +8 lines]", - content: "hidden", - }}, - } - - updated, _ := m.Update(tea.KeyMsg{Type: tea.KeyBackspace}) - m = updated.(chatModel) - if got := string(m.input); got != "use " { - t.Fatalf("input after backspace = %q, want pasted text placeholder removed", got) - } - if got := len(m.inputPastedTexts); got != 0 { - t.Fatalf("pasted texts after backspace = %d, want 0", got) - } -} - -func TestInitialPromptHistoryLoadsFromMessages(t *testing.T) { - history := initialPromptHistory(context.Background(), Options{ - Messages: []api.Message{ - {Role: "user", Content: "old prompt"}, - {Role: "assistant", Content: "answer"}, - {Role: "user", Content: "new prompt"}, - }, - }) - - if got, want := strings.Join(history, "|"), "old prompt|new prompt"; got != want { - t.Fatalf("history = %#v, want %s", history, want) - } -} - -func TestSkillCommandsListAndPersistSyntheticToolCall(t *testing.T) { - dir := t.TempDir() - skillDir := filepath.Join(dir, "release-notes") - if err := os.MkdirAll(skillDir, 0o755); err != nil { - t.Fatal(err) - } - if err := os.WriteFile(filepath.Join(skillDir, "SKILL.md"), []byte("---\nname: release-notes\ndescription: Draft release notes.\n---\nUse concise bullets."), 0o644); err != nil { - t.Fatal(err) - } - catalog, err := coreagent.DiscoverSkills(dir) - if err != nil { - t.Fatal(err) - } - - m := chatModel{opts: Options{Skills: catalog}, input: []rune("/skills")} - updated, cmd := m.handleSubmit() - if cmd != nil { - t.Fatal("/skills should not start a model run") - } - m = updated.(chatModel) - if len(m.entries) != 1 || !strings.Contains(m.entries[0].content, "release-notes") { - t.Fatalf("/skills entries = %#v", m.entries) - } - - m = chatModel{ctx: context.Background(), opts: Options{Model: "test", Skills: catalog, Client: chatTestClient{}}, input: []rune("/release-notes")} - updated, cmd = m.handleSubmit() - if cmd == nil { - t.Fatal("/ should continue the chat with the loaded instructions") - } - m = updated.(chatModel) - events := m.events - for { - msg, ok := <-events - if !ok { - t.Fatal("skill run closed before it finished") - } - updated, _ = m.Update(msg) - m = updated.(chatModel) - if _, ok := msg.(chatRunDoneMsg); ok { - break - } - } - if len(m.messages) != 4 { - t.Fatalf("synthetic messages = %#v", m.messages) - } - call := m.messages[1] - result := m.messages[2] - if call.Role != "assistant" || len(call.ToolCalls) != 1 || call.ToolCalls[0].Function.Name != "skill" || !strings.HasPrefix(call.ToolCalls[0].ID, "call_skill_") { - t.Fatalf("synthetic call = %#v", call) - } - if result.Role != "tool" || result.ToolCallID != call.ToolCalls[0].ID || !strings.Contains(result.Content, "Use concise bullets.") { - t.Fatalf("synthetic result = %#v", result) - } - entries := entriesFromMessages(m.messages) - if len(entries) != 3 || entries[1].toolID != call.ToolCalls[0].ID || entries[1].detail != "skill" || entries[1].args["name"] != "release-notes" { - t.Fatalf("round-trip entries = %#v", entries) - } -} - -func TestSkillsImportReloadsCatalogRegistryAndSystemPrompt(t *testing.T) { - before := writeTestSkillCatalog(t) - dir := t.TempDir() - if err := os.Mkdir(filepath.Join(dir, "from-codex"), 0o755); err != nil { - t.Fatal(err) - } - if err := os.WriteFile(filepath.Join(dir, "from-codex", "SKILL.md"), []byte("---\nname: from-codex\ndescription: Imported skill.\n---\nImported instructions."), 0o644); err != nil { - t.Fatal(err) - } - after, err := coreagent.DiscoverSkills(dir) - if err != nil { - t.Fatal(err) - } - registry := &coreagent.Registry{} - var reloaded, rebuilt, prompted bool - m := chatModel{ - ctx: context.Background(), - opts: Options{ - Model: "test", - Skills: before, - ImportSkills: func(source string) (coreagent.SkillImportResult, error) { - if source != "codex" { - t.Fatalf("source = %q", source) - } - return coreagent.SkillImportResult{Source: source, SourceDir: "/source", Imported: []string{"from-codex"}}, nil - }, - ReloadSkills: func() (*coreagent.SkillCatalog, error) { - reloaded = true - return after, nil - }, - ToolRegistryForModel: func(context.Context, string) *coreagent.Registry { - rebuilt = true - return registry - }, - SystemPromptForModel: func(_ context.Context, _ string, got *coreagent.Registry, _ bool) string { - prompted = got == registry - return after.SystemContext() - }, - }, - input: []rune("/skills import codex"), - } - - updated, cmd := m.handleSubmit() - if cmd != nil { - t.Fatal("skills import should not start a model run") - } - m = updated.(chatModel) - if !reloaded || !rebuilt || !prompted { - t.Fatalf("reload=%v rebuilt=%v prompted=%v", reloaded, rebuilt, prompted) - } - if m.opts.Skills != after || m.opts.Tools != registry || !strings.Contains(m.opts.SystemPrompt, "from-codex") { - t.Fatalf("reloaded options = %#v", m.opts) - } - if m.status != "skills reloaded" || len(m.entries) != 1 || !strings.Contains(m.entries[0].content, "Imported 1 skill") { - t.Fatalf("import result = status %q entries %#v", m.status, m.entries) - } -} - -func TestSkillsImportUsage(t *testing.T) { - m := chatModel{input: []rune("/skills import")} - updated, cmd := m.handleSubmit() - if cmd != nil { - t.Fatal("invalid skills import should not start a model run") - } - m = updated.(chatModel) - if len(m.entries) != 1 || m.entries[0].role != "error" || !strings.Contains(m.entries[0].content, "usage: /skills [import codex|claude|pi]") { - t.Fatalf("entries = %#v", m.entries) - } -} - -func TestSkillSlashCommandPromptBecomesUserMessage(t *testing.T) { - catalog := writeTestSkillCatalog(t) - m := chatModel{ctx: context.Background(), opts: Options{Model: "test", Skills: catalog, Client: chatTestClient{}}, input: []rune("/release-notes draft the v1.2 notes")} - updated, cmd := m.handleSubmit() - if cmd == nil { - t.Fatal("/ should start a run") - } - m = updated.(chatModel) - for { - msg, ok := <-m.events - if !ok { - t.Fatal("skill run closed before it finished") - } - updated, _ = m.Update(msg) - m = updated.(chatModel) - if _, ok := msg.(chatRunDoneMsg); ok { - break - } - } - if len(m.messages) < 1 || m.messages[0].Role != "user" || m.messages[0].Content != "draft the v1.2 notes" { - t.Fatalf("user message = %#v, want the prompt", m.messages[0]) - } - // The skill still loads as a synthetic tool call right after the user turn. - if len(m.messages) < 3 || m.messages[1].Role != "assistant" || len(m.messages[1].ToolCalls) != 1 || m.messages[1].ToolCalls[0].Function.Name != "skill" { - t.Fatalf("synthetic skill call missing: %#v", m.messages) - } -} - -func TestChatSkillSubmitWhileActiveRunKeepsActiveState(t *testing.T) { - catalog := writeTestSkillCatalog(t) - for _, state := range []struct { - name string - running bool - compacting bool - }{ - {name: "running", running: true}, - {name: "compacting", compacting: true}, - } { - t.Run(state.name, func(t *testing.T) { - events := make(chan tea.Msg) - cancel := func() {} - m := chatModel{ - opts: Options{Skills: catalog}, - input: []rune("/release-notes draft notes"), - running: state.running, - compacting: state.compacting, - events: events, - cancel: cancel, - } - - updated, cmd := m.handleSubmit() - if cmd != nil { - t.Fatal("skill submit should not start another run while active") - } - got := updated.(chatModel) - if got.events != events || got.cancel == nil || got.running != state.running || got.compacting != state.compacting { - t.Fatalf("active run state changed: %#v", got) - } - if string(got.input) != "/release-notes draft notes" { - t.Fatalf("input = %q, want skill invocation preserved", got.input) - } - if got.status != "wait for current response" { - t.Fatalf("status = %q", got.status) - } - }) - } -} - -func writeTestSkillCatalog(t *testing.T) *coreagent.SkillCatalog { - t.Helper() - dir := t.TempDir() - skillDir := filepath.Join(dir, "release-notes") - if err := os.MkdirAll(skillDir, 0o755); err != nil { - t.Fatal(err) - } - if err := os.WriteFile(filepath.Join(skillDir, "SKILL.md"), []byte("---\nname: release-notes\ndescription: Draft release notes.\n---\nUse concise bullets."), 0o644); err != nil { - t.Fatal(err) - } - catalog, err := coreagent.DiscoverSkills(dir) - if err != nil { - t.Fatal(err) - } - return catalog -} - -func TestSkillSlashCommandAppearsInCompletions(t *testing.T) { - catalog := writeTestSkillCatalog(t) - m := chatModel{opts: Options{Skills: catalog}, input: []rune("/re")} - - lines := stripANSI(strings.Join(m.slashCommandLines(80), "\n")) - if !strings.Contains(lines, "/release-notes") || !strings.Contains(lines, "Draft release notes.") { - t.Fatalf("suggestions missing /release-notes: %q", lines) - } -} - -func TestSkillsImportSlashCompletions(t *testing.T) { - for _, test := range []struct { - input string - want []string - }{ - {input: "/skills", want: []string{"/skills", "/skills import"}}, - {input: "/skills impo", want: []string{"/skills import"}}, - {input: "/skills import ", want: []string{"/skills import codex", "/skills import claude", "/skills import pi"}}, - {input: "/skills import c", want: []string{"/skills import codex", "/skills import claude"}}, - {input: "/skills import pi", want: []string{"/skills import pi"}}, - } { - t.Run(test.input, func(t *testing.T) { - m := chatModel{input: []rune(test.input)} - completions := m.slashCompletions() - got := make([]string, 0, len(completions)) - for _, completion := range completions { - got = append(got, completion.value) - } - if strings.Join(got, "\n") != strings.Join(test.want, "\n") { - t.Fatalf("completions = %#v, want %#v", got, test.want) - } - }) - } -} - -func TestSkillSlashPromptHidesCommandCompletions(t *testing.T) { - catalog := writeTestSkillCatalog(t) - for _, input := range []string{"/release-notes ", "/release-notes draft the release notes"} { - t.Run(input, func(t *testing.T) { - m := chatModel{opts: Options{Skills: catalog}, input: []rune(input)} - if lines := m.completionLines(80); len(lines) != 0 { - t.Fatalf("completion lines = %#v, want none", lines) - } - }) - } -} - -func TestSkillSlashNameResolvesAndRejectsArgsAndUnknown(t *testing.T) { - catalog := writeTestSkillCatalog(t) - m := &chatModel{opts: Options{Skills: catalog}} - - if name, _, ok := m.skillSlashInvocation("/release-notes"); !ok || name != "release-notes" { - t.Fatalf("/release-notes = %q %v, want release-notes true", name, ok) - } - if name, prompt, ok := m.skillSlashInvocation("/release-notes draft notes"); !ok || name != "release-notes" || prompt != "draft notes" { - t.Fatalf("/release-notes draft notes = %q %q %v, want release-notes / draft notes / true", name, prompt, ok) - } - if _, _, ok := m.skillSlashInvocation("/no-such-skill"); ok { - t.Fatal("unknown skill should not resolve") - } - // A built-in command sharing a prefix must not be claimed as a skill. - if _, _, ok := m.skillSlashInvocation("/skills"); ok { - t.Fatal("/skills should resolve to the built-in, not a skill") - } - - // Unknown slash input that is not a skill stays an unknown command. - m2 := chatModel{opts: Options{Skills: catalog}, input: []rune("/no-such-skill")} - updated, cmd := m2.handleSubmit() - if cmd != nil { - t.Fatal("unknown slash command should not start a run") - } - m2 = updated.(chatModel) - if len(m2.entries) != 1 || m2.entries[0].role != "error" || !strings.Contains(m2.entries[0].content, "Unknown command") { - t.Fatalf("entries = %#v, want unknown command", m2.entries) - } -} - -func TestChatDeletedSlashCommandsAreUnknown(t *testing.T) { - for _, command := range []string{"/clear", "/copy", "/copy-all", "/launch", "/history", "/load", "/raw", "/resume", "/set", "/show", "/verbose"} { - t.Run(command, func(t *testing.T) { - m := chatModel{input: []rune(command)} - - updated, cmd := m.handleSubmit() - if cmd != nil { - t.Fatal("deleted slash command should not return a command") - } - m = updated.(chatModel) - if len(m.entries) != 1 || m.entries[0].role != "error" || !strings.Contains(m.entries[0].content, "Unknown command") { - t.Fatalf("entries = %#v, want unknown command error", m.entries) - } - }) - } -} - -func TestChatViewRendersSlashCommandSuggestions(t *testing.T) { - m := chatModel{ - input: []rune("/"), - width: 80, - height: 18, - } - - view := stripANSI(m.View()) - for _, want := range []string{"/model", "/new", "/think", "/tools", "/system"} { - if !strings.Contains(view, want) { - t.Fatalf("view missing %s suggestion: %q", want, view) - } - } - for _, removed := range []string{"/clear", "/copy", "/copy-all", "/history", "/load", "/raw", "/resume", "/set", "/show", "/verbose"} { - if strings.Contains(view, removed) { - t.Fatalf("bare slash should hide removed command %s: %q", removed, view) - } - } - if got := len(m.slashCommandLines(80)); got != maxSlashCompletions { - t.Fatalf("slash suggestions = %d, want %d", got, maxSlashCompletions) - } -} - -func TestChatToolsCommandTogglesToolRegistry(t *testing.T) { - registry := &coreagent.Registry{} - registry.Register(chatTestTool{}) - calls := 0 - m := chatModel{ - ctx: context.Background(), - input: []rune("/tools"), - opts: Options{ - Model: "llama3.2", - Tools: registry, - ToolRegistryForModel: func(context.Context, string) *coreagent.Registry { - calls++ - return registry - }, - }, - } - - updated, cmd := m.handleSubmit() - if cmd != nil { - t.Fatal("/tools should not start a command") - } - m = updated.(chatModel) - if !m.opts.ToolsDisabled || m.opts.Tools == nil { - t.Fatalf("tools off state = disabled:%v tools:%#v", m.opts.ToolsDisabled, m.opts.Tools) - } - if len(m.entries) != 0 { - t.Fatalf("/tools should only update action space, got entries:%#v", m.entries) - } - req, _ := m.requestPreview() - if got := len(req.Tools); got != 0 { - t.Fatalf("request preview tools = %d, want 0", got) - } - - m.input = []rune("/tools") - updated, cmd = m.handleSubmit() - if cmd != nil { - t.Fatal("/tools should not start a command") - } - m = updated.(chatModel) - if m.opts.ToolsDisabled || m.opts.Tools == nil { - t.Fatalf("tools on state = disabled:%v tools:%#v", m.opts.ToolsDisabled, m.opts.Tools) - } - if calls != 1 { - t.Fatalf("tool registry calls = %d, want 1", calls) - } - if len(m.entries) != 0 { - t.Fatalf("/tools should only update action space, got entries:%#v", m.entries) - } - req, _ = m.requestPreview() - if got := len(req.Tools); got != 1 { - t.Fatalf("request preview tools = %d, want 1", got) - } -} - -func TestChatToolsCommandRefreshesCapabilityAwareSystemPrompt(t *testing.T) { - registry := &coreagent.Registry{} - registry.Register(chatTestTool{}) - m := chatModel{ - ctx: context.Background(), - opts: Options{ - Model: "test", - Tools: registry, - SystemPromptForModel: func(_ context.Context, _ string, _ *coreagent.Registry, disabled bool) string { - if disabled { - return "tools disabled" - } - return "tools enabled" - }, - }, - } - - updated, _ := m.handleToolsCommand("") - m = updated.(chatModel) - if m.opts.SystemPrompt != "tools disabled" { - t.Fatalf("system prompt = %q, want disabled prompt", m.opts.SystemPrompt) - } - updated, _ = m.handleToolsCommand("") - m = updated.(chatModel) - if m.opts.SystemPrompt != "tools enabled" { - t.Fatalf("system prompt = %q, want enabled prompt", m.opts.SystemPrompt) - } -} - -func TestChatToolsCommandUsage(t *testing.T) { - m := chatModel{input: []rune("/tools off")} - - updated, cmd := m.handleSubmit() - if cmd != nil { - t.Fatal("invalid /tools should not start a command") - } - m = updated.(chatModel) - if m.status != "error" || len(m.entries) != 1 || !strings.Contains(m.entries[0].content, "usage: /tools") { - t.Fatalf("invalid /tools result = status:%q entries:%#v", m.status, m.entries) - } -} - -func TestChatSystemCommandControlsBuiltInSystemPrompt(t *testing.T) { - client := &chatCaptureClient{} - m := chatModel{ - ctx: context.Background(), - input: []rune("/system"), - opts: Options{ - Model: "test", - Client: client, - SystemPrompt: "canonical agent prompt", - }, - } - - updated, cmd := m.handleSubmit() - if cmd != nil { - t.Fatal("/system should not start a run") - } - m = updated.(chatModel) - if len(m.entries) != 1 || m.entries[0].role != "slash" || m.entries[0].content != "Built-in system prompt is on.\n\ncanonical agent prompt\n\nWarning: Changing the system prompt during a session breaks the prompt cache." { - t.Fatalf("/system entry = %#v", m.entries) - } - - m.input = []rune("/system off") - updated, _ = m.handleSubmit() - m = updated.(chatModel) - if !m.systemPromptDisabled || m.status != "system prompt off" { - t.Fatalf("/system off state = disabled:%v status:%q", m.systemPromptDisabled, m.status) - } - m.input = []rune("/system") - updated, _ = m.handleSubmit() - m = updated.(chatModel) - if got := m.entries[len(m.entries)-1].content; got != "Built-in system prompt is off.\n\ncanonical agent prompt\n\nWarning: Changing the system prompt during a session breaks the prompt cache." { - t.Fatalf("/system off entry = %q", got) - } - updated, cmd = m.startRun("hello") - if cmd == nil { - t.Fatal("run after /system off should start") - } - m = updated.(chatModel) - if done := waitForRunDone(t, m.events); done.err != nil { - t.Fatalf("run after /system off: %v", done.err) - } - if len(client.requests) != 1 || len(client.requests[0].Messages) != 1 || client.requests[0].Messages[0].Role != "user" { - t.Fatalf("request after /system off = %#v", client.requests) - } - - m.input = []rune("/system ON") - updated, _ = m.handleSubmit() - m = updated.(chatModel) - if m.systemPromptDisabled || m.status != "system prompt on" { - t.Fatalf("/system on state = disabled:%v status:%q", m.systemPromptDisabled, m.status) - } - updated, cmd = m.startRun("hello again") - if cmd == nil { - t.Fatal("run after /system on should start") - } - m = updated.(chatModel) - if done := waitForRunDone(t, m.events); done.err != nil { - t.Fatalf("run after /system on: %v", done.err) - } - if len(client.requests) != 2 { - t.Fatalf("client requests = %d, want 2", len(client.requests)) - } - request := client.requests[1] - if len(request.Messages) != 2 || request.Messages[0].Role != "system" || request.Messages[0].Content != "canonical agent prompt" { - t.Fatalf("request after /system on = %#v", request.Messages) - } - - m.input = []rune("/system sometimes") - updated, _ = m.handleSubmit() - m = updated.(chatModel) - if m.status != "error" || len(m.entries) == 0 || m.entries[len(m.entries)-1].content != "usage: /system [on|off]" { - t.Fatalf("invalid /system result = status:%q entries:%#v", m.status, m.entries) - } -} - -func TestChatSystemCommandArgumentCompletions(t *testing.T) { - for _, tt := range []struct { - input string - want []string - }{ - {input: "/system ", want: []string{"/system on", "/system off"}}, - {input: "/system o", want: []string{"/system on", "/system off"}}, - {input: "/system on", want: []string{"/system on"}}, - } { - t.Run(tt.input, func(t *testing.T) { - m := chatModel{input: []rune(tt.input)} - completions := m.slashCompletions() - if len(completions) != len(tt.want) { - t.Fatalf("completions = %#v, want %d", completions, len(tt.want)) - } - for i, want := range tt.want { - if completions[i].value != want { - t.Fatalf("completion %d = %q, want %q", i, completions[i].value, want) - } - } - }) - } - - m := chatModel{input: []rune("/system ")} - lines := stripANSI(strings.Join(m.slashCommandLines(80), "\n")) - for _, want := range []string{"on", "enable the built-in system prompt", "off", "disable the built-in system prompt"} { - if !strings.Contains(lines, want) { - t.Fatalf("/system option suggestions missing %q: %q", want, lines) - } - } - updated, cmd := m.Update(tea.KeyMsg{Type: tea.KeyEnter}) - if cmd != nil { - t.Fatal("selecting /system on should not submit the command") - } - m = updated.(chatModel) - if got := string(m.input); got != "/system on" { - t.Fatalf("input = %q, want /system on", got) - } - - m.input = []rune("/system maybe") - completions := m.slashCompletions() - if len(completions) != 1 || completions[0].label != "No matching options" { - t.Fatalf("invalid argument completions = %#v", completions) - } -} - -func TestChatSlashCommandSuggestionsIncludePromptAndSave(t *testing.T) { - for _, tt := range []struct { - input string - command string - description string - }{ - {input: "/pr", command: "/prompt", description: "show full prompt, tools, and messages"}, - {input: "/sa", command: "/save", description: "save request JSON; saved as .json"}, - } { - t.Run(tt.command, func(t *testing.T) { - m := chatModel{input: []rune(tt.input)} - - lines := stripANSI(strings.Join(m.slashCommandLines(80), "\n")) - if !strings.Contains(lines, tt.command) || !strings.Contains(lines, tt.description) { - t.Fatalf("suggestions missing %s: %q", tt.command, lines) - } - }) - } -} - -func TestChatSlashCommandSuggestionsIncludeThink(t *testing.T) { - m := chatModel{input: []rune("/th")} - - lines := stripANSI(strings.Join(m.slashCommandLines(80), "\n")) - if !strings.Contains(lines, "/think") || !strings.Contains(lines, "set thinking mode") { - t.Fatalf("suggestions missing /think: %q", lines) - } -} - -func TestChatEnterFillsSelectedSlashCommandBeforeSubmitting(t *testing.T) { - m := chatModel{input: []rune("/th")} - - updated, cmd := m.Update(tea.KeyMsg{Type: tea.KeyEnter}) - m = updated.(chatModel) - if cmd != nil { - t.Fatal("filling a slash command should not return a command") - } - if got := string(m.input); got != "/think" { - t.Fatalf("input = %q, want completed command", got) - } - if m.thinkPicker != nil { - t.Fatal("filling a slash command should not open its picker") - } - - updated, cmd = m.Update(tea.KeyMsg{Type: tea.KeyEnter}) - m = updated.(chatModel) - if cmd != nil { - t.Fatal("think command should not return a command") - } - if m.thinkPicker == nil { - t.Fatal("second enter should submit the completed /think command") - } -} - -func TestChatEnterSubmitsExactSlashCommandAliases(t *testing.T) { - t.Run("help", func(t *testing.T) { - m := chatModel{input: []rune("/?")} - - updated, cmd := m.Update(tea.KeyMsg{Type: tea.KeyEnter}) - if cmd != nil { - t.Fatal("help alias should not return a command") - } - m = updated.(chatModel) - if len(m.entries) != 1 || m.entries[0].role != "slash" { - t.Fatalf("entries = %#v, want help output", m.entries) - } - if got := string(m.input); got != "" { - t.Fatalf("input = %q, want cleared after submitting alias", got) - } - }) - - t.Run("exit", func(t *testing.T) { - m := chatModel{input: []rune("/exit")} - - updated, cmd := m.Update(tea.KeyMsg{Type: tea.KeyEnter}) - if cmd == nil { - t.Fatal("exit alias should return the quit command") - } - m = updated.(chatModel) - if !m.quitting { - t.Fatal("exit alias should quit without filling /bye first") - } - if got := string(m.input); got != "" { - t.Fatalf("input = %q, want cleared after submitting alias", got) - } - }) -} - -func TestChatSlashCommandsRunWhileModelResponds(t *testing.T) { - m := chatModel{running: true, input: []rune("/help")} - - updated, cmd := m.handleSubmit() - if cmd != nil { - t.Fatal("help command should not return a command") - } - m = updated.(chatModel) - if len(m.entries) != 1 || m.entries[0].role != "slash" { - t.Fatalf("entries = %#v, want immediate slash output", m.entries) - } -} - -func TestChatMessageSubmitWhileRunningPreservesDraft(t *testing.T) { - attachment := chatInputAttachment{placeholder: "[image 1]", kind: "image"} - pasted := chatInputPastedText{placeholder: "[pasted 1]", content: "long paste"} - m := chatModel{ - running: true, - input: []rune("next prompt [image 1] [pasted 1]"), - inputCursor: 4, - inputCursorSet: true, - inputAttachments: []chatInputAttachment{attachment}, - inputPastedTexts: []chatInputPastedText{pasted}, - } - - updated, cmd := m.handleSubmit() - if cmd != nil { - t.Fatal("busy submit should not start a command") - } - m = updated.(chatModel) - if got := string(m.input); got != "next prompt [image 1] [pasted 1]" { - t.Fatalf("input = %q, want draft preserved", got) - } - if m.inputCursor != 4 || !m.inputCursorSet { - t.Fatalf("cursor not preserved: cursor=%d set=%v", m.inputCursor, m.inputCursorSet) - } - if len(m.inputAttachments) != 1 || m.inputAttachments[0].placeholder != attachment.placeholder || m.inputAttachments[0].kind != attachment.kind { - t.Fatalf("attachments not preserved: %#v", m.inputAttachments) - } - if len(m.inputPastedTexts) != 1 || m.inputPastedTexts[0] != pasted { - t.Fatalf("pasted text not preserved: %#v", m.inputPastedTexts) - } - if len(m.entries) != 0 { - t.Fatalf("busy submit should not add entries: %#v", m.entries) - } -} - -func TestChatThinkCommandOpensPicker(t *testing.T) { - m := chatModel{input: []rune("/think")} - - updated, cmd := m.handleSubmit() - if cmd != nil { - t.Fatal("think command should not return a command") - } - m = updated.(chatModel) - if m.thinkPicker == nil { - t.Fatal("think picker should open") - } - if view := stripANSI(m.renderThinkPicker(80)); !strings.Contains(view, "Thinking mode") || !strings.Contains(view, "high") { - t.Fatalf("think picker view missing options: %q", view) - } -} - -func TestChatThinkCommandSetsModes(t *testing.T) { - for _, tt := range []struct { - input string - want any - }{ - {input: "/think on", want: true}, - {input: "/think off", want: false}, - {input: "/think high", want: "high"}, - } { - t.Run(tt.input, func(t *testing.T) { - m := chatModel{input: []rune(tt.input)} - updated, cmd := m.handleSubmit() - if cmd != nil { - t.Fatal("think command should not return a command") - } - m = updated.(chatModel) - if m.opts.Think == nil || m.opts.Think.Value != tt.want { - t.Fatalf("think = %#v, want %#v", m.opts.Think, tt.want) - } - }) - } -} - -func TestChatViewRendersFileMentionSuggestions(t *testing.T) { - dir := t.TempDir() - if err := os.Mkdir(filepath.Join(dir, "cmd"), 0o755); err != nil { - t.Fatal(err) - } - if err := os.WriteFile(filepath.Join(dir, "README.md"), []byte("hi"), 0o644); err != nil { - t.Fatal(err) - } - m := chatModel{ - workingDir: dir, - input: []rune("open @"), - width: 80, - height: 18, - } - - view := stripANSI(m.View()) - if !strings.Contains(view, "@cmd/") || !strings.Contains(view, "@README.md") { - t.Fatalf("file mention suggestions missing: %q", view) - } -} - -func TestChatFileMentionSuggestionsFilterAndComplete(t *testing.T) { - dir := t.TempDir() - if err := os.Mkdir(filepath.Join(dir, "cmd"), 0o755); err != nil { - t.Fatal(err) - } - if err := os.WriteFile(filepath.Join(dir, "README.md"), []byte("hi"), 0o644); err != nil { - t.Fatal(err) - } - m := chatModel{ - workingDir: dir, - input: []rune("open @REA"), - } - - lines := stripANSI(strings.Join(m.completionLines(80), "\n")) - if !strings.Contains(lines, "@README.md") || strings.Contains(lines, "@cmd/") { - t.Fatalf("filtered file suggestions = %q", lines) - } - m.applyCompletion() - if got := string(m.input); got != "open @README.md " { - t.Fatalf("completed input = %q", got) - } -} - -func TestChatEnterCompletesHighlightedFileMentionWithoutSubmitting(t *testing.T) { - dir := t.TempDir() - if err := os.WriteFile(filepath.Join(dir, "alpha.md"), []byte("hi"), 0o644); err != nil { - t.Fatal(err) - } - if err := os.WriteFile(filepath.Join(dir, "target.md"), []byte("hi"), 0o644); err != nil { - t.Fatal(err) - } - - m := chatModel{ - workingDir: dir, - input: []rune("review @ after this"), - inputCursor: len([]rune("review @")), - inputCursorSet: true, - } - - updated, _ := m.Update(tea.KeyMsg{Type: tea.KeyDown}) - m = updated.(chatModel) - if got, want := m.complete, 1; got != want { - t.Fatalf("selected completion = %d, want %d", got, want) - } - - updated, cmd := m.Update(tea.KeyMsg{Type: tea.KeyEnter}) - if cmd != nil { - t.Fatal("selecting a file mention should not submit the prompt") - } - m = updated.(chatModel) - if got, want := string(m.input), "review @target.md after this"; got != want { - t.Fatalf("input = %q, want %q", got, want) - } - if got, want := m.inputCursor, len([]rune("review @target.md ")); got != want || !m.inputCursorSet { - t.Fatalf("cursor = %d (set=%v), want %d after the inserted mention", got, m.inputCursorSet, want) - } - if len(m.entries) != 0 || len(m.messages) != 0 { - t.Fatalf("selecting a file mention submitted the prompt: entries=%#v messages=%#v", m.entries, m.messages) - } - if completions := m.mentionCompletions(); completions != nil { - t.Fatalf("mention selector remained visible after selection: %#v", completions) - } -} diff --git a/cmd/tui/chat/markdown.go b/cmd/tui/chat/markdown.go deleted file mode 100644 index 49dec98d64b..00000000000 --- a/cmd/tui/chat/markdown.go +++ /dev/null @@ -1,425 +0,0 @@ -package chat - -import ( - "strings" - "unicode" - "unicode/utf8" - - "github.com/charmbracelet/lipgloss" - "github.com/mattn/go-runewidth" -) - -func renderMarkdownForView(markdown string, width int) string { - if width < 20 { - width = 20 - } - - source := strings.Split(strings.TrimRight(markdown, "\n"), "\n") - var rendered []string - inCodeBlock := false - for i := 0; i < len(source); i++ { - line := strings.TrimRight(source[i], "\r") - trimmed := strings.TrimSpace(line) - - if strings.HasPrefix(trimmed, "```") { - inCodeBlock = !inCodeBlock - continue - } - if inCodeBlock { - rendered = append(rendered, renderMarkdownCodeLine(line, width)...) - continue - } - - if table, consumed := renderMarkdownTable(source[i:], width); consumed > 0 { - rendered = append(rendered, table...) - i += consumed - 1 - continue - } - - if heading, ok := markdownHeading(trimmed); ok { - rendered = append(rendered, chatHeaderStyle.Render(renderMarkdownRunes(parseMarkdownInline(heading)))) - continue - } - - if trimmed == "" { - rendered = append(rendered, "") - continue - } - rendered = append(rendered, wrapMarkdownInline(line, width)...) - } - return strings.Join(rendered, "\n") -} - -func splitRenderedBody(body string) []string { - body = strings.TrimRight(body, "\n") - if body == "" { - return []string{""} - } - return strings.Split(body, "\n") -} - -func markdownHeading(line string) (string, bool) { - if !strings.HasPrefix(line, "#") { - return "", false - } - level := 0 - for level < len(line) && line[level] == '#' { - level++ - } - if level == 0 || level > 6 || level >= len(line) || line[level] != ' ' { - return "", false - } - return strings.TrimSpace(line[level:]), true -} - -type markdownInlineStyle uint8 - -const ( - markdownPlain markdownInlineStyle = iota - markdownStrong - markdownCode -) - -type markdownInlineRune struct { - r rune - style markdownInlineStyle -} - -// wrapMarkdownInline parses a complete source line before wrapping it. That -// keeps emphasis intact when its opening and closing delimiters land on -// different visual lines. -func wrapMarkdownInline(line string, width int) []string { - return wrapInlineRunes(parseMarkdownInline(line), width) -} - -func wrapInlineRunes(runes []markdownInlineRune, width int) []string { - if len(runes) == 0 { - return []string{""} - } - - var rendered []string - for len(runes) > 0 { - hardCut, spaceCut, currentWidth := 0, 0, 0 - for i, item := range runes { - nextWidth := currentWidth + runewidth.RuneWidth(item.r) - if nextWidth > width { - break - } - currentWidth = nextWidth - hardCut = i + 1 - if unicode.IsSpace(item.r) && currentWidth > width/2 { - spaceCut = i - } - } - cut := hardCut - if spaceCut > 0 { - cut = spaceCut - } - if cut == 0 { - cut = 1 - } - - lineRunes := trimMarkdownSpace(runes[:cut]) - rendered = append(rendered, renderMarkdownRunes(lineRunes)) - runes = trimMarkdownSpace(runes[cut:]) - } - return rendered -} - -func parseMarkdownInline(line string) []markdownInlineRune { - var out []markdownInlineRune - for len(line) > 0 { - if strings.HasPrefix(line, "`") { - if end := strings.Index(line[1:], "`"); end >= 0 { - out = appendMarkdownRunes(out, line[1:end+1], markdownCode) - line = line[end+2:] - continue - } - } - if (strings.HasPrefix(line, "**") || strings.HasPrefix(line, "__")) && canOpenMarkdownStrong(out) { - delimiter := line[:2] - if end := strings.Index(line[2:], delimiter); end >= 0 { - out = appendMarkdownRunes(out, line[2:end+2], markdownStrong) - line = line[end+4:] - continue - } - } - - r, size := utf8.DecodeRuneInString(line) - out = append(out, markdownInlineRune{r: r, style: markdownPlain}) - line = line[size:] - } - return out -} - -// canOpenMarkdownStrong keeps delimiter-like text in bare URLs and identifiers -// literal, only treating ** / __ as strong emphasis at the common -// whitespace- or punctuation-delimited form. -func canOpenMarkdownStrong(out []markdownInlineRune) bool { - if len(out) == 0 { - return true - } - previous := out[len(out)-1].r - return (unicode.IsSpace(previous) || unicode.IsPunct(previous)) && !markdownStrongInURL(out) -} - -func markdownStrongInURL(out []markdownInlineRune) bool { - start := len(out) - for start > 0 && !unicode.IsSpace(out[start-1].r) { - start-- - } - - var token strings.Builder - for _, item := range out[start:] { - token.WriteRune(item.r) - } - return strings.Contains(token.String(), "://") -} - -func appendMarkdownRunes(out []markdownInlineRune, text string, style markdownInlineStyle) []markdownInlineRune { - for _, r := range text { - out = append(out, markdownInlineRune{r: r, style: style}) - } - return out -} - -func trimMarkdownSpace(runes []markdownInlineRune) []markdownInlineRune { - start, end := 0, len(runes) - for start < end && unicode.IsSpace(runes[start].r) { - start++ - } - for end > start && unicode.IsSpace(runes[end-1].r) { - end-- - } - return runes[start:end] -} - -func renderMarkdownRunes(runes []markdownInlineRune) string { - var b strings.Builder - for start := 0; start < len(runes); { - end := start + 1 - for end < len(runes) && runes[end].style == runes[start].style { - end++ - } - var text strings.Builder - for _, item := range runes[start:end] { - text.WriteRune(item.r) - } - switch runes[start].style { - case markdownStrong: - b.WriteString(chatStrongStyle.Render(text.String())) - case markdownCode: - b.WriteString(chatInlineCodeStyle.Render(text.String())) - default: - b.WriteString(text.String()) - } - start = end - } - return b.String() -} - -func renderMarkdownCodeLine(line string, width int) []string { - codeWidth := max(1, width-2) - lines := wrapChatText(line, codeWidth) - for i, wrapped := range lines { - lines[i] = " " + chatCodeBlockStyle.Render(wrapped) - } - return lines -} - -func renderMarkdownTable(lines []string, width int) ([]string, int) { - if len(lines) < 2 || !looksLikeMarkdownTableRow(lines[0]) || !isMarkdownTableSeparator(lines[1]) { - return nil, 0 - } - - var rows [][]string - consumed := 0 - for consumed < len(lines) && looksLikeMarkdownTableRow(lines[consumed]) { - if consumed == 1 && isMarkdownTableSeparator(lines[consumed]) { - consumed++ - continue - } - rows = append(rows, parseMarkdownTableRow(lines[consumed])) - consumed++ - } - if len(rows) == 0 { - return nil, 0 - } - - columnCount := 0 - for _, row := range rows { - columnCount = max(columnCount, len(row)) - } - naturalWidths := make([]int, columnCount) - for _, row := range rows { - for i := range columnCount { - cell := "" - if i < len(row) { - cell = row[i] - } - naturalWidths[i] = max(naturalWidths[i], markdownInlineWidth(cell)) - } - } - widths := markdownTableColumnWidths(naturalWidths, width) - - var rendered []string - for rowIndex, row := range rows { - wrappedCells := make([][]string, columnCount) - rowHeight := 1 - for i := range columnCount { - cell := "" - if i < len(row) { - cell = row[i] - } - wrappedCells[i] = wrapMarkdownTableCell(cell, widths[i]) - rowHeight = max(rowHeight, len(wrappedCells[i])) - } - for lineIndex := range rowHeight { - cells := make([]string, columnCount) - for i := range columnCount { - cellLine := "" - if lineIndex < len(wrappedCells[i]) { - cellLine = wrappedCells[i][lineIndex] - } - cells[i] = padPlainLine(cellLine, widths[i]) - } - line := strings.Join(cells, chatTableBorderStyle.Render(" | ")) - if rowIndex == 0 { - line = chatHeaderStyle.Render(stripANSIForWidth(line)) - } - rendered = append(rendered, line) - } - } - return rendered, consumed -} - -func markdownTableColumnWidths(naturalWidths []int, width int) []int { - if len(naturalWidths) == 0 { - return nil - } - separatorWidth := max(0, len(naturalWidths)-1) * lipglossWidth(" | ") - available := max(1, width-separatorWidth) - widths := make([]int, len(naturalWidths)) - minWidths := make([]int, len(naturalWidths)) - for i, natural := range naturalWidths { - widths[i] = max(1, natural) - minWidth := min(widths[i], 12) - if i == 0 { - minWidth = min(widths[i], 4) - } - minWidths[i] = max(1, minWidth) - } - - for sumInts(widths) > available { - index := widestShrinkableColumn(widths, minWidths) - if index < 0 { - break - } - widths[index]-- - } - for sumInts(widths) > available { - index := widestColumn(widths) - if index < 0 || widths[index] <= 1 { - break - } - widths[index]-- - } - return widths -} - -func widestShrinkableColumn(widths, minWidths []int) int { - index := -1 - for i, width := range widths { - if width <= minWidths[i] { - continue - } - if index < 0 || width > widths[index] { - index = i - } - } - return index -} - -func widestColumn(widths []int) int { - index := -1 - for i, width := range widths { - if index < 0 || width > widths[index] { - index = i - } - } - return index -} - -func sumInts(values []int) int { - sum := 0 - for _, value := range values { - sum += value - } - return sum -} - -func wrapMarkdownTableCell(cell string, width int) []string { - lines := wrapInlineRunes(parseMarkdownInline(cell), max(1, width)) - if len(lines) == 0 { - return []string{""} - } - return lines -} - -// markdownInlineWidth reports the visible width of a cell once Markdown -// delimiters are parsed away, so columns size to rendered content. -func markdownInlineWidth(cell string) int { - width := 0 - for _, item := range parseMarkdownInline(cell) { - width += runewidth.RuneWidth(item.r) - } - return width -} - -func looksLikeMarkdownTableRow(line string) bool { - line = strings.TrimSpace(line) - return strings.Contains(line, "|") && strings.Count(line, "|") >= 1 -} - -func isMarkdownTableSeparator(line string) bool { - cells := parseMarkdownTableRow(line) - if len(cells) == 0 { - return false - } - for _, cell := range cells { - cell = strings.TrimSpace(cell) - cell = strings.TrimPrefix(cell, ":") - cell = strings.TrimSuffix(cell, ":") - if cell == "" || strings.Trim(cell, "-") != "" { - return false - } - } - return true -} - -func parseMarkdownTableRow(line string) []string { - line = strings.TrimSpace(line) - line = strings.TrimPrefix(line, "|") - line = strings.TrimSuffix(line, "|") - raw := strings.Split(line, "|") - cells := make([]string, 0, len(raw)) - for _, cell := range raw { - cells = append(cells, strings.TrimSpace(cell)) - } - return cells -} - -func padPlainLine(line string, width int) string { - if extra := width - lipglossWidth(line); extra > 0 { - return line + strings.Repeat(" ", extra) - } - return line -} - -func stripANSIForWidth(line string) string { - return stripChatANSI(line) -} - -func lipglossWidth(line string) int { - return lipgloss.Width(line) -} diff --git a/cmd/tui/chat/modals.go b/cmd/tui/chat/modals.go deleted file mode 100644 index 7ee55dcb331..00000000000 --- a/cmd/tui/chat/modals.go +++ /dev/null @@ -1,263 +0,0 @@ -package chat - -import ( - "context" - "fmt" - "slices" - "strings" - - tea "github.com/charmbracelet/bubbletea" - - apptui "github.com/ollama/ollama/cmd/tui" -) - -type chatModelPicker = apptui.SelectorModel - -func (m *chatModel) openModelPicker(filter string) (tea.Model, tea.Cmd) { - if m.opts.ModelOptions == nil { - m.entries = append(m.entries, newChatEntry(chatEntry{role: "error", content: "Model picker is unavailable.", err: "Model picker is unavailable."})) - m.status = "error" - return *m, nil - } - - ctx := m.ctx - if ctx == nil { - ctx = context.Background() - } - models, err := m.opts.ModelOptions(ctx) - if err != nil { - m.entries = append(m.entries, newChatEntry(chatEntry{role: "error", content: fmt.Sprintf("Could not list models: %v", err), err: err.Error()})) - m.status = "error" - return *m, nil - } - models = normalizeModelOptions(models) - if len(models) == 0 { - m.entries = append(m.entries, newChatEntry(chatEntry{role: "system", content: "No models available."})) - m.status = "ready" - return *m, nil - } - - items := modelSelectorItems(models, m.opts.Model) - current := m.opts.Model - if !m.openModelOnInit { - items = compactModelSelectorItems(models, m.opts.Model) - current = "" - } - picker := apptui.NewModelSelectorModel("Select model", items, current, filter) - picker.SetHelpText("↑/↓ navigate • enter select • type search • esc cancel") - m.modelPicker = &picker - m.modelPickerModels = models - m.status = "model" - return *m, nil -} - -func normalizeModelOptions(models []ModelOption) []ModelOption { - seen := make(map[string]struct{}, len(models)) - out := make([]ModelOption, 0, len(models)) - for _, model := range models { - model.Name = strings.TrimSpace(model.Name) - model.Description = strings.TrimSpace(model.Description) - if model.Name == "" { - continue - } - key := strings.ToLower(model.Name) - if _, ok := seen[key]; ok { - continue - } - seen[key] = struct{}{} - out = append(out, model) - } - slices.SortStableFunc(out, func(a, b ModelOption) int { - if a.Recommended == b.Recommended { - return 0 - } - if a.Recommended { - return -1 - } - return 1 - }) - return out -} - -func modelSelectorItems(models []ModelOption, current string) []apptui.SelectItem { - return modelSelectorItemsWithCurrentPriority(models, current, true) -} - -func compactModelSelectorItems(models []ModelOption, current string) []apptui.SelectItem { - return modelSelectorItemsWithCurrentPriority(models, current, false) -} - -func modelSelectorItemsWithCurrentPriority(models []ModelOption, current string, pinCurrent bool) []apptui.SelectItem { - ordered := slices.Clone(models) - slices.SortStableFunc(ordered, func(a, b ModelOption) int { - if cmp := compareModelPickerGroup(modelPickerGroup(a, current, pinCurrent), modelPickerGroup(b, current, pinCurrent)); cmp != 0 { - return cmp - } - return 0 - }) - - items := make([]apptui.SelectItem, 0, len(ordered)) - for _, model := range ordered { - items = append(items, apptui.SelectItem{ - Name: model.Name, - Description: modelOptionMeta(model), - Recommended: model.Name == current || !model.Cloud || model.Recommended, - AvailabilityBadge: model.AvailabilityBadge, - }) - } - return items -} - -func modelPickerGroup(model ModelOption, current string, pinCurrent bool) int { - if pinCurrent && model.Name == current { - return 0 - } - if model.Recommended { - return 1 - } - if model.Name == current { - return 2 - } - if !model.Cloud { - return 3 - } - return 4 -} - -func compareModelPickerGroup(a, b int) int { - switch { - case a < b: - return -1 - case a > b: - return 1 - default: - return 0 - } -} - -func (m chatModel) updateModelPicker(msg tea.KeyMsg) (tea.Model, tea.Cmd) { - if m.modelPicker == nil { - return m, nil - } - switch msg.Type { - case tea.KeyCtrlC, tea.KeyEsc: - m.modelPicker = nil - m.modelPickerModels = nil - m.openModelOnInit = false - m.status = "ready" - return m, nil - case tea.KeyEnter: - return m.selectModel() - default: - m.modelPicker.UpdateNavigation(msg) - } - return m, nil -} - -func (m chatModel) selectModel() (tea.Model, tea.Cmd) { - if m.modelPicker == nil { - return m, nil - } - selectedItem, ok := m.modelPicker.SelectedItem() - if !ok { - return m, nil - } - selected, ok := m.modelOptionForSelection(selectedItem.Name) - if !ok { - return m, nil - } - - // Cloud models need auth + plan check before switching. If we already - // know the badge state from the model list, go directly to the right - // prompt — no "checking" spinner. - if selected.Cloud && m.opts.CheckCloudModel != nil { - switch selected.AvailabilityBadge { - case "Sign in required": - return m.startCloudAuthSignIn(selected.Name, selected.RequiredPlan, selected.SignInURL) - case "Upgrade required": - return m.startCloudAuthUpgrade(selected.Name, selected.RequiredPlan) - } - // Badge is empty — auth is satisfied (confirmed via Whoami when the - // list was built). Apply directly. - } - - m.modelPicker = nil - m.modelPickerModels = nil - m.openModelOnInit = false - if err := m.applyModelSelection(selected.Name, true); err != nil { - m.entries = append(m.entries, newChatEntry(chatEntry{role: "error", content: fmt.Sprintf("Could not switch model: %v", err), err: err.Error()})) - m.status = "error" - return m, nil - } - m.status = "ready" - return m, tea.Batch(m.startModelPreload(selected.Name), cloudModelPreflightCmd(m.ctx, m.opts, selected.Name, selected.RequiredPlan)) -} - -func (m chatModel) modelOptionForSelection(name string) (ModelOption, bool) { - for _, model := range m.modelPickerModels { - if model.Name == name { - return model, true - } - } - return ModelOption{}, false -} - -func (m *chatModel) applyModelSelection(modelName string, persist bool) error { - modelName = strings.TrimSpace(modelName) - if modelName == "" { - return nil - } - m.opts.Model = modelName - m.opts.ContextWindowTokens = 0 - if m.opts.ToolRegistryForModel != nil { - m.opts.Tools = m.opts.ToolRegistryForModel(m.ctx, modelName) - } - if m.opts.SystemPromptForModel != nil { - m.opts.SystemPrompt = m.opts.SystemPromptForModel(m.ctx, modelName, m.opts.Tools, m.opts.ToolsDisabled) - } - if m.opts.MultiModalForModel != nil { - ctx := m.ctx - if ctx == nil { - ctx = context.Background() - } - m.opts.MultiModal = m.opts.MultiModalForModel(ctx, modelName) - } - m.refreshContextWindowTokens(modelName) - m.contextTokens = m.estimatePromptTokens(m.messages, "") - m.contextEstimate = true - if persist && m.opts.OnModelSelected != nil { - ctx := m.ctx - if ctx == nil { - ctx = context.Background() - } - return m.opts.OnModelSelected(ctx, modelName) - } - return nil -} - -func (m *chatModel) startModelPreload(modelName string) tea.Cmd { - modelName = strings.TrimSpace(modelName) - if m == nil || modelName == "" || m.opts.PreloadModel == nil { - return nil - } - m.preloadingModel = modelName - m.spinner = 0 - return tea.Batch(preloadModelCmd(m.ctx, m.opts.PreloadModel, modelName, m.opts.Think), m.scheduleTick()) -} - -func (m chatModel) renderModelPicker(width int) string { - return m.modelPicker.RenderContent() -} - -func (m chatModel) renderInlineModelPicker(width int) []string { - rendered := m.modelPicker.RenderCompactContent(maxInlineModelPickerItems) - lines := strings.Split(strings.TrimRight(rendered, "\n"), "\n") - for i := range lines { - lines[i] = truncateRenderedLine(lines[i], width) - } - return lines -} - -func modelOptionMeta(model ModelOption) string { - return strings.TrimSpace(model.Description) -} diff --git a/cmd/tui/chat/modals_test.go b/cmd/tui/chat/modals_test.go deleted file mode 100644 index f51b6ed32eb..00000000000 --- a/cmd/tui/chat/modals_test.go +++ /dev/null @@ -1,441 +0,0 @@ -package chat - -import ( - "context" - "slices" - "strings" - "testing" - - tea "github.com/charmbracelet/bubbletea" - - coreagent "github.com/ollama/ollama/agent" - "github.com/ollama/ollama/api" - apptui "github.com/ollama/ollama/cmd/tui" -) - -func TestChatModelCommandOpensPicker(t *testing.T) { - m := chatModel{ - ctx: context.Background(), - input: []rune("/model"), - width: 100, - height: 20, - opts: Options{ - Model: "llama3.2", - ContextWindowTokens: 131072, - ModelOptions: func(context.Context) ([]ModelOption, error) { - return []ModelOption{ - {Name: "kimi-k2.6:cloud", Description: "cloud coding"}, - {Name: "llama3.2", Description: "local"}, - }, nil - }, - }, - } - - updated, cmd := m.handleSubmit() - if cmd != nil { - t.Fatal("model command should not return a command") - } - m = updated.(chatModel) - if m.modelPicker == nil { - t.Fatal("model picker was not opened") - } - view := stripANSI(m.View()) - if !strings.Contains(view, "Select model") || - !strings.Contains(view, "Type to filter") || - !strings.Contains(view, "kimi-k2.6:cloud") || - !strings.Contains(view, "llama3.2") { - t.Fatalf("model picker view missing content: %q", view) - } - if strings.Contains(view, "Search...") { - t.Fatalf("model picker should render inline without full search box: %q", view) - } - if strings.Contains(view, "local") || strings.Contains(view, "cloud coding") { - t.Fatalf("inline model picker should stay compact without descriptions: %q", view) - } - if !strings.Contains(view, "│ █") { - t.Fatalf("inline model picker should keep input box visible: %q", view) - } -} - -func TestChatModelCommandShowsRecommendedFirstWithoutSections(t *testing.T) { - m := chatModel{ - ctx: context.Background(), - input: []rune("/model"), - width: 100, - height: 20, - opts: Options{ - Model: "llama3.2", - ModelOptions: func(context.Context) ([]ModelOption, error) { - return []ModelOption{ - {Name: "llama3.2", Description: "selected local"}, - {Name: "gemma4", Description: "local"}, - {Name: "glm-5.2:cloud", Description: "recommended cloud", Recommended: true, Cloud: true}, - {Name: "kimi-k2.7-code:cloud", Description: "another recommended cloud", Recommended: true, Cloud: true}, - }, nil - }, - }, - } - - updated, cmd := m.handleSubmit() - if cmd != nil { - t.Fatal("model command should not return a command") - } - view := stripANSI(updated.(chatModel).View()) - for _, unwanted := range []string{"Recommended", "More", "recommended cloud", "selected local"} { - if strings.Contains(view, unwanted) { - t.Fatalf("compact model picker should be flat and description-free; found %q in %q", unwanted, view) - } - } - firstRecommended := strings.Index(view, "glm-5.2:cloud") - secondRecommended := strings.Index(view, "kimi-k2.7-code:cloud") - current := strings.Index(view, "llama3.2") - local := strings.Index(view, "gemma4") - if firstRecommended < 0 || secondRecommended < 0 || current < 0 || local < 0 { - t.Fatalf("compact model picker missing expected models: %q", view) - } - if !(firstRecommended < current && secondRecommended < current && current < local) { - t.Fatalf("compact model picker order should be recommended, current, local: %q", view) - } -} - -func TestChatModelCommandOpensSmallPicker(t *testing.T) { - m := chatModel{ - ctx: context.Background(), - input: []rune("/model"), - width: 100, - height: 24, - opts: Options{ - Model: "model-1", - ModelOptions: func(context.Context) ([]ModelOption, error) { - return []ModelOption{ - {Name: "model-1"}, - {Name: "model-2"}, - {Name: "model-3"}, - {Name: "model-4"}, - {Name: "model-5"}, - {Name: "model-6"}, - {Name: "model-7"}, - }, nil - }, - }, - } - - updated, cmd := m.handleSubmit() - if cmd != nil { - t.Fatal("model command should not return a command") - } - m = updated.(chatModel) - view := stripANSI(m.View()) - for _, want := range []string{"model-1", "model-5", "... and 2 more"} { - if !strings.Contains(view, want) { - t.Fatalf("small model picker missing %q: %q", want, view) - } - } - if strings.Contains(view, "model-6") || strings.Contains(view, "model-7") { - t.Fatalf("small model picker rendered too many items: %q", view) - } -} - -func TestChatModelPickerStaysInlineWhenSmall(t *testing.T) { - m := chatModel{ - ctx: context.Background(), - input: []rune("/model"), - width: 44, - height: 10, - opts: Options{ - Model: "llama3.2", - ModelOptions: func(context.Context) ([]ModelOption, error) { - return []ModelOption{ - {Name: "kimi-k2.6:cloud", Description: "cloud coding"}, - {Name: "llama3.2", Description: "local"}, - }, nil - }, - }, - } - - updated, cmd := m.handleSubmit() - if cmd != nil { - t.Fatal("model command should not return a command") - } - m = updated.(chatModel) - view := stripANSI(m.View()) - if !strings.Contains(view, "Select model") || !strings.Contains(view, "Type to filter") { - t.Fatalf("small model picker should stay inline: %q", view) - } - if strings.Contains(view, "Search...") { - t.Fatalf("small model picker should not use bespoke full-frame search: %q", view) - } -} - -func TestChatModelPickerShowsRecommendedModelsFirst(t *testing.T) { - models := normalizeModelOptions([]ModelOption{ - {Name: "llama3.2", Description: "local"}, - {Name: "kimi-k2.6:cloud", Description: "cloud coding", Recommended: true}, - {Name: "qwen3.5:cloud", Description: "cloud reasoning", Recommended: true}, - {Name: "gemma4", Description: "local"}, - }) - got := make([]string, 0, len(models)) - for _, model := range models { - got = append(got, model.Name) - } - want := []string{"kimi-k2.6:cloud", "qwen3.5:cloud", "llama3.2", "gemma4"} - if !slices.Equal(got, want) { - t.Fatalf("model order = %#v, want %#v", got, want) - } -} - -func TestChatModelPickerPinsCurrentThenRecommendedModels(t *testing.T) { - models := normalizeModelOptions([]ModelOption{ - {Name: "llama3.2", Description: "local"}, - {Name: "glm-5.2:cloud", Description: "cloud selected", Recommended: true, Cloud: true}, - {Name: "kimi-k2.7-code:cloud", Description: "cloud coding", Recommended: true, Cloud: true}, - {Name: "gemma4", Description: "local"}, - }) - - items := modelSelectorItems(models, "glm-5.2:cloud") - got := make([]string, 0, len(items)) - for _, item := range items { - got = append(got, item.Name) - } - want := []string{"glm-5.2:cloud", "kimi-k2.7-code:cloud", "llama3.2", "gemma4"} - if !slices.Equal(got, want) { - t.Fatalf("selector item order = %#v, want %#v", got, want) - } - for _, item := range items[:3] { - if !item.Recommended { - t.Fatalf("%q should be pinned in the first picker section", item.Name) - } - } - if items[0].Description != "cloud selected" { - t.Fatalf("current model description = %q, want plain model description", items[0].Description) - } -} - -func TestInitialModelPickerRendersBeforeChatShell(t *testing.T) { - models := normalizeModelOptions([]ModelOption{ - {Name: "glm-5.2:cloud", Description: "cloud selected", Recommended: true, Cloud: true}, - {Name: "llama3.2", Description: "local"}, - }) - picker := apptui.NewModelSelectorModel("Select model", modelSelectorItems(models, "glm-5.2:cloud"), "glm-5.2:cloud", "") - m := chatModel{ - width: 100, - height: 20, - openModelOnInit: true, - modelPicker: &picker, - entries: []chatEntry{{role: "assistant", content: "old chat content"}}, - } - - view := stripANSI(m.View()) - if !strings.Contains(view, "Select model") || !strings.Contains(view, "llama3.2") { - t.Fatalf("initial picker view missing model content: %q", view) - } - if strings.Contains(view, "old chat content") || strings.Contains(view, "│ █") { - t.Fatalf("initial picker should render before chat shell: %q", view) - } -} - -func TestChatModelPickerRanksClosestFilteredModelFirst(t *testing.T) { - models := normalizeModelOptions([]ModelOption{ - {Name: "gemma3:27b", Description: "recommended but longer", Recommended: true}, - {Name: "llama3.2", Description: "mentions gemm in description"}, - {Name: "gemma4:27b", Description: "longer local"}, - {Name: "gemma4", Description: "short local"}, - }) - picker := apptui.NewModelSelectorModel("Select model", modelSelectorItems(models, ""), "", "gemm") - - filtered := picker.FilteredItems() - got := make([]string, 0, len(filtered)) - for _, model := range filtered { - got = append(got, model.Name) - } - want := []string{"gemma4", "gemma3:27b", "gemma4:27b", "llama3.2"} - if !slices.Equal(got, want) { - t.Fatalf("filtered model order = %#v, want %#v", got, want) - } -} - -func TestChatModelPickerFiltersAndSwitchesModel(t *testing.T) { - var savedModel string - originalMessages := []api.Message{{Role: "user", Content: "keep me"}} - m := chatModel{ - ctx: context.Background(), - chatID: "chat-1", - input: []rune("/model qwen"), - width: 100, - height: 20, - messages: slices.Clone(originalMessages), - opts: Options{ - Model: "llama3.2", - ModelOptions: func(context.Context) ([]ModelOption, error) { - return []ModelOption{ - {Name: "llama3.2", Description: "local"}, - {Name: "qwen3.5:cloud", Description: "cloud reasoning"}, - }, nil - }, - ToolRegistryForModel: func(ctx context.Context, model string) *coreagent.Registry { - if model != "qwen3.5:cloud" { - t.Fatalf("tool registry model = %q, want qwen3.5:cloud", model) - } - registry := &coreagent.Registry{} - registry.Register(chatTestTool{}) - return registry - }, - ContextWindowTokensForModel: func(ctx context.Context, model string, fallback int) int { - if model != "qwen3.5:cloud" { - t.Fatalf("context model = %q, want qwen3.5:cloud", model) - } - if fallback != 0 { - t.Fatalf("context fallback = %d, want 0 after model switch", fallback) - } - return 262144 - }, - SystemPromptForModel: func(ctx context.Context, model string, registry *coreagent.Registry, toolsDisabled bool) string { - if model != "qwen3.5:cloud" { - t.Fatalf("system prompt model = %q, want qwen3.5:cloud", model) - } - if registry == nil { - t.Fatalf("system prompt registry missing fake tool: %#v", registry) - } - if _, ok := registry.Get("fake_tool"); !ok { - t.Fatalf("system prompt registry missing fake tool: %#v", registry) - } - return "system for " + model - }, - OnModelSelected: func(ctx context.Context, model string) error { - savedModel = model - return nil - }, - }, - } - - updated, cmd := m.handleSubmit() - if cmd != nil { - t.Fatal("model command should not return a command") - } - m = updated.(chatModel) - if m.modelPicker == nil || m.modelPicker.Filter() != "qwen" { - t.Fatalf("model picker = %#v, want qwen filter", m.modelPicker) - } - if view := stripANSI(m.View()); !strings.Contains(view, "qwen3.5:cloud") || strings.Contains(view, "llama3.2") { - t.Fatalf("filtered model picker view = %q", view) - } - - updated, cmd = m.Update(tea.KeyMsg{Type: tea.KeyEnter}) - m = updated.(chatModel) - if cmd != nil { - t.Fatal("switching models should not start a command") - } - if m.modelPicker != nil { - t.Fatal("model picker should close after selection") - } - if m.status != "ready" || m.notificationLine() != "" { - t.Fatalf("model switch should not show action status, status=%q notification=%q", m.status, m.notificationLine()) - } - if m.opts.Model != "qwen3.5:cloud" { - t.Fatalf("model = %q, want qwen3.5:cloud", m.opts.Model) - } - if m.chatID != "chat-1" { - t.Fatalf("chatID = %q, want chat-1", m.chatID) - } - if len(m.messages) != len(originalMessages) || m.messages[0].Content != originalMessages[0].Content { - t.Fatalf("messages changed on model switch: %#v", m.messages) - } - if len(m.entries) != 0 { - t.Fatalf("model switch should not append transcript entries: %#v", m.entries) - } - if savedModel != "qwen3.5:cloud" { - t.Fatalf("saved model = %q, want qwen3.5:cloud", savedModel) - } - if m.opts.Tools == nil { - t.Fatalf("tools registry was not rebuilt for model: %#v", m.opts.Tools) - } - if _, ok := m.opts.Tools.Get("fake_tool"); !ok { - t.Fatalf("tools registry was not rebuilt for model: %#v", m.opts.Tools) - } - if m.opts.ContextWindowTokens != 262144 { - t.Fatalf("context window = %d, want 262144", m.opts.ContextWindowTokens) - } - if m.opts.SystemPrompt != "system for qwen3.5:cloud" { - t.Fatalf("system prompt = %q", m.opts.SystemPrompt) - } -} - -func TestChatModelSelectionStartsBackgroundPreload(t *testing.T) { - m := chatModel{ - ctx: context.Background(), - opts: Options{ - Model: "llama3.2", - PreloadModel: func(context.Context, string, *api.ThinkValue) (int, error) { - return 0, nil - }, - }, - } - - if err := m.applyModelSelection("qwen3", false); err != nil { - t.Fatal(err) - } - cmd := m.startModelPreload("qwen3") - if cmd == nil { - t.Fatal("model switch should start background preload when configured") - } - if m.preloadingModel != "qwen3" { - t.Fatalf("preloadingModel = %q, want qwen3", m.preloadingModel) - } -} - -func TestChatModelSwitchNextRunKeepsHistory(t *testing.T) { - client := &chatCaptureClient{} - history := []api.Message{ - {Role: "user", Content: "old question"}, - {Role: "assistant", Content: "old answer"}, - } - m := chatModel{ - ctx: context.Background(), - chatID: "chat-1", - messages: slices.Clone(history), - input: []rune("continue"), - opts: Options{ - Model: "llama3.2", - Client: client, - SystemPromptForModel: func(_ context.Context, model string, _ *coreagent.Registry, _ bool) string { - return "system for " + model - }, - }, - } - if err := m.applyModelSelection("qwen3", true); err != nil { - t.Fatal(err) - } - - updated, cmd := m.handleSubmit() - m = updated.(chatModel) - if cmd == nil { - t.Fatal("next prompt should start a model run") - } - done := waitForRunDone(t, m.events) - if done.err != nil { - t.Fatal(done.err) - } - - if len(client.requests) != 1 { - t.Fatalf("requests = %d, want 1", len(client.requests)) - } - req := client.requests[0] - if req.Model != "qwen3" { - t.Fatalf("request model = %q, want qwen3", req.Model) - } - if len(req.Messages) != 4 { - t.Fatalf("request messages = %#v, want system + 2 history + new user", req.Messages) - } - if req.Messages[0].Role != "system" || req.Messages[0].Content != "system for qwen3" { - t.Fatalf("system message = %#v", req.Messages[0]) - } - for i, want := range history { - got := req.Messages[i+1] - if got.Role != want.Role || got.Content != want.Content { - t.Fatalf("history message %d = %#v, want %#v", i, got, want) - } - } - if req.Messages[3].Role != "user" || req.Messages[3].Content != "continue" { - t.Fatalf("new user message = %#v", req.Messages[3]) - } -} diff --git a/cmd/tui/chat/multimodal_test.go b/cmd/tui/chat/multimodal_test.go deleted file mode 100644 index 159465ce261..00000000000 --- a/cmd/tui/chat/multimodal_test.go +++ /dev/null @@ -1,307 +0,0 @@ -package chat - -import ( - "context" - "net/url" - "os" - "path/filepath" - "strings" - "testing" - - tea "github.com/charmbracelet/bubbletea" -) - -func TestChatStartRunAttachesDroppedImagePath(t *testing.T) { - fp := writeTestPNG(t) - m := chatModel{ - ctx: context.Background(), - opts: Options{ - Model: "test", - Client: chatTestClient{}, - MultiModal: true, - }, - } - - updated, cmd := m.startRun("describe " + fp) - m = updated.(chatModel) - - if cmd == nil { - t.Fatal("startRun should return a command") - } - if len(m.liveMessages) != 1 { - t.Fatalf("liveMessages = %d, want 1", len(m.liveMessages)) - } - if got := m.liveMessages[0].Content; got != "describe" { - t.Fatalf("content = %q, want describe", got) - } - if got := len(m.liveMessages[0].Images); got != 1 { - t.Fatalf("images = %d, want 1", got) - } - if len(m.entries) == 0 { - t.Fatal("missing user transcript entry") - } - entry := m.entries[0].content - if strings.Contains(entry, fp) { - t.Fatalf("transcript entry should hide local file path: %q", entry) - } - if !strings.Contains(entry, "describe") || !strings.Contains(entry, "[attached 1 file]") { - t.Fatalf("transcript entry = %q, want prompt plus attachment note", entry) - } -} - -func TestChatStartRunAttachesDroppedFileURL(t *testing.T) { - fp := writeTestPNG(t) - fileURL := (&url.URL{Scheme: "file", Path: fp}).String() - m := chatModel{ - ctx: context.Background(), - opts: Options{ - Model: "test", - Client: chatTestClient{}, - MultiModal: true, - }, - } - - updated, _ := m.startRun(fileURL) - m = updated.(chatModel) - - if got := m.liveMessages[0].Content; got != "" { - t.Fatalf("content = %q, want empty prompt after extracting file URL", got) - } - if got := len(m.liveMessages[0].Images); got != 1 { - t.Fatalf("images = %d, want 1", got) - } - if got := m.entries[0].content; got != "[attached 1 file]" { - t.Fatalf("transcript entry = %q, want attachment-only note", got) - } -} - -func TestChatPasteImagePathAttachesOnSubmit(t *testing.T) { - fp := writeTestPNG(t) - m := chatModel{ - ctx: context.Background(), - opts: Options{ - Model: "test", - Client: chatTestClient{}, - MultiModal: true, - }, - } - - updated, _ := m.Update(tea.KeyMsg{Type: tea.KeyRunes, Runes: []rune("describe " + fp), Paste: true}) - m = updated.(chatModel) - if got := string(m.input); got != "describe [Image #0]" { - t.Fatalf("pasted path input = %q, want placeholder", got) - } - if got := m.notificationLine(); got != "" { - t.Fatalf("notification = %q, want no attachment notification", got) - } - if got := string(m.input); strings.Contains(got, fp) { - t.Fatalf("pasted path should be hidden behind placeholder, input = %q", got) - } - if completions := m.slashCompletions(); len(completions) != 0 { - t.Fatalf("placeholder input should not show slash completions: %#v", completions) - } - - updated, cmd := m.Update(tea.KeyMsg{Type: tea.KeyEnter}) - m = updated.(chatModel) - if cmd == nil { - t.Fatal("submit should start a run") - } - if got := m.liveMessages[0].Content; got != "describe [Image #0]" { - t.Fatalf("content = %q, want prompt with placeholder", got) - } - if got := len(m.liveMessages[0].Images); got != 1 { - t.Fatalf("images = %d, want 1", got) - } - if strings.Contains(m.entries[0].content, fp) { - t.Fatalf("transcript entry should hide pasted file path: %q", m.entries[0].content) - } - if !strings.Contains(m.entries[0].content, "[Image #0]") { - t.Fatalf("transcript entry should show placeholder: %q", m.entries[0].content) - } -} - -func TestChatPasteImagePathAfterSwitchingToMultimodalModel(t *testing.T) { - fp := writeTestPNG(t) - m := chatModel{ - ctx: context.Background(), - input: []rune("/model vision"), - opts: Options{ - Model: "text", - ModelOptions: func(context.Context) ([]ModelOption, error) { - return []ModelOption{ - {Name: "text", Description: "local"}, - {Name: "vision", Description: "local vision"}, - }, nil - }, - MultiModalForModel: func(_ context.Context, model string) bool { - return model == "vision" - }, - }, - } - - updated, cmd := m.handleSubmit() - if cmd != nil { - t.Fatal("model picker should not return a command") - } - m = updated.(chatModel) - - updated, cmd = m.Update(tea.KeyMsg{Type: tea.KeyEnter}) - if cmd != nil { - t.Fatal("model switch should not preload without a preload hook") - } - m = updated.(chatModel) - if m.opts.Model != "vision" { - t.Fatalf("model = %q, want vision", m.opts.Model) - } - if !m.opts.MultiModal { - t.Fatal("switching to a multimodal model should enable image paste handling") - } - - updated, _ = m.Update(tea.KeyMsg{Type: tea.KeyRunes, Runes: []rune("describe " + fp), Paste: true}) - m = updated.(chatModel) - if got := string(m.input); got != "describe [Image #0]" { - t.Fatalf("pasted path input = %q, want placeholder", got) - } - if strings.Contains(string(m.input), fp) { - t.Fatalf("pasted path should be hidden behind placeholder, input = %q", string(m.input)) - } -} - -func TestChatImagePlaceholdersUseSessionNumbers(t *testing.T) { - first := writeTestPNG(t) - second := writeTestPNG(t) - m := chatModel{ - ctx: context.Background(), - opts: Options{ - Model: "test", - Client: chatTestClient{}, - MultiModal: true, - }, - } - - updated, _ := m.Update(tea.KeyMsg{Type: tea.KeyRunes, Runes: []rune(first), Paste: true}) - m = updated.(chatModel) - if got := string(m.input); got != "[Image #0]" { - t.Fatalf("first placeholder = %q, want [Image #0]", got) - } - - updated, cmd := m.Update(tea.KeyMsg{Type: tea.KeyEnter}) - m = updated.(chatModel) - if cmd == nil { - t.Fatal("submit should start a run") - } - - updated, _ = m.Update(tea.KeyMsg{Type: tea.KeyRunes, Runes: []rune(second), Paste: true}) - m = updated.(chatModel) - if got := string(m.input); got != "[Image #1]" { - t.Fatalf("second placeholder = %q, want [Image #1]", got) - } -} - -func TestChatAbsoluteImagePathBypassesSlashCommandParsing(t *testing.T) { - fp := writeTestPNG(t) - m := chatModel{ - ctx: context.Background(), - opts: Options{ - Model: "test", - Client: chatTestClient{}, - MultiModal: true, - }, - } - - m.input = []rune(fp) - m.inputCursor = len(m.input) - m.inputCursorSet = true - updated, cmd := m.handleSubmit() - m = updated.(chatModel) - - if cmd == nil { - t.Fatal("absolute image path should start a run instead of being parsed as a slash command") - } - if got := len(m.liveMessages[0].Images); got != 1 { - t.Fatalf("images = %d, want 1", got) - } - if got := m.entries[0].role; got != "user" { - t.Fatalf("entry role = %q, want user", got) - } -} - -func TestChatDeletingImagePlaceholderRemovesAttachment(t *testing.T) { - fp := writeTestPNG(t) - m := chatModel{ - ctx: context.Background(), - opts: Options{ - Model: "test", - Client: chatTestClient{}, - MultiModal: true, - }, - } - - updated, _ := m.Update(tea.KeyMsg{Type: tea.KeyRunes, Runes: []rune("describe " + fp), Paste: true}) - m = updated.(chatModel) - if got := len(m.inputAttachments); got != 1 { - t.Fatalf("input attachments = %d, want 1", got) - } - - updated, _ = m.Update(tea.KeyMsg{Type: tea.KeyBackspace}) - m = updated.(chatModel) - if got := string(m.input); got != "describe " { - t.Fatalf("input after backspace = %q, want image placeholder removed", got) - } - if got := len(m.inputAttachments); got != 0 { - t.Fatalf("input attachments after editing placeholder = %d, want 0", got) - } - - m.input = []rune("describe") - m.inputCursor = len(m.input) - m.inputCursorSet = true - updated, cmd := m.Update(tea.KeyMsg{Type: tea.KeyEnter}) - m = updated.(chatModel) - if cmd == nil { - t.Fatal("submit should start a run") - } - if got := len(m.liveMessages[0].Images); got != 0 { - t.Fatalf("images = %d, want 0 after deleting placeholder", got) - } - if got := m.liveMessages[0].Content; got != "describe" { - t.Fatalf("content = %q, want describe", got) - } -} - -func TestChatWordDeletingImagePlaceholderRemovesAttachment(t *testing.T) { - fp := writeTestPNG(t) - m := chatModel{ - ctx: context.Background(), - opts: Options{ - Model: "test", - Client: chatTestClient{}, - MultiModal: true, - }, - } - - updated, _ := m.Update(tea.KeyMsg{Type: tea.KeyRunes, Runes: []rune("describe " + fp), Paste: true}) - m = updated.(chatModel) - updated, _ = m.Update(tea.KeyMsg{Type: tea.KeySpace}) - m = updated.(chatModel) - updated, _ = m.Update(tea.KeyMsg{Type: tea.KeyBackspace, Alt: true}) - m = updated.(chatModel) - - if got := string(m.input); got != "describe " { - t.Fatalf("input after word backspace = %q, want image placeholder removed", got) - } - if got := len(m.inputAttachments); got != 0 { - t.Fatalf("input attachments after word backspace = %d, want 0", got) - } -} - -func writeTestPNG(t *testing.T) string { - t.Helper() - dir := t.TempDir() - fp := filepath.Join(dir, "dragged image.png") - data := make([]byte, 600) - copy(data, []byte{0x89, 'P', 'N', 'G', '\r', '\n', 0x1a, '\n'}) - if err := os.WriteFile(fp, data, 0o600); err != nil { - t.Fatalf("failed to write test image: %v", err) - } - return fp -} diff --git a/cmd/tui/chat/render.go b/cmd/tui/chat/render.go deleted file mode 100644 index fb85688a6c1..00000000000 --- a/cmd/tui/chat/render.go +++ /dev/null @@ -1,2240 +0,0 @@ -package chat - -import ( - "context" - "encoding/json" - "fmt" - "slices" - "sort" - "strconv" - "strings" - "time" - "unicode" - "unicode/utf8" - - "github.com/charmbracelet/lipgloss" - "github.com/charmbracelet/x/ansi" - "github.com/mattn/go-runewidth" - - coreagent "github.com/ollama/ollama/agent" - "github.com/ollama/ollama/api" -) - -type chatEntry struct { - role string - content string - label string - detail string - status string - err string - toolID string - args map[string]any - expanded bool - startedAt time.Time - finishedAt time.Time - tools []chatEntry - metrics *api.Metrics - tokenCount int - - version int - renderKey chatEntryRenderKey - renderLines []string -} - -const ( - chatMessageIndent = " " - chatUserMessagePrefix = "" - maxCtrlOToolOutputRunes = 400 - maxLiveThinkingRunes = 4096 - - defaultViewWidth = 80 - defaultViewHeight = 24 -) - -type chatEntryRenderKey struct { - width int - version int -} - -func (m chatModel) findToolEntry(toolID string) int { - if toolID == "" { - return -1 - } - for i := len(m.entries) - 1; i >= 0; i-- { - if m.entries[i].role == "tool" && m.entries[i].toolID == toolID { - return i - } - } - return -1 -} - -func (m chatModel) findActiveToolEntry(toolID string) int { - idx := m.findToolEntry(toolID) - if idx < 0 || !isToolActiveStatus(m.entries[idx].status) { - return -1 - } - return idx -} - -func (m chatModel) toolStartedAt(toolID string) time.Time { - idx := m.findToolEntry(toolID) - if idx < 0 { - return time.Time{} - } - return m.entries[idx].startedAt -} - -func entryHasExpandableOutput(entry chatEntry) bool { - return (entry.role == "tool" && (len(entry.args) > 0 || strings.TrimSpace(entry.content) != "")) || - (entry.role == "tool_group" && len(entry.tools) > 0) || - (entry.role == "compaction_summary" && strings.TrimSpace(entry.content) != "") -} - -func entryHasToolOutputMode(entry chatEntry) bool { - return (entry.role == "tool" && (isToolActiveStatus(entry.status) || isToolResultStatus(entry.status) || entry.content != "")) || - (entry.role == "tool_group" && len(entry.tools) > 0) || - (entry.role == "compaction_summary" && strings.TrimSpace(entry.content) != "") -} - -func (m *chatModel) applyToolOutputMode() { - if !m.toolOutputMode { - return - } - for i := range m.entries { - m.applyToolOutputModeTo(i) - } -} - -func (m *chatModel) applyToolOutputModeTo(index int) { - if !m.toolOutputMode || index < 0 || index >= len(m.entries) { - return - } - if !entryHasToolOutputMode(m.entries[index]) { - return - } - if m.entries[index].expanded == m.toolOutputOpen { - return - } - m.entries[index].expanded = m.toolOutputOpen - m.markEntryDirty(index) -} - -func (m *chatModel) groupCompletedToolHistory() { - if m.hasPendingDetectedToolCalls() { - m.applyToolOutputMode() - return - } - m.entries = groupCompletedToolEntries(m.entries, m.detectedToolCalls...) - m.applyToolOutputMode() -} - -func (m chatModel) hasPendingDetectedToolCalls() bool { - if len(m.detectedToolCalls) == 0 { - return false - } - results := map[string]struct{}{} - for _, tool := range flattenToolHistory(m.entries) { - if tool.toolID == "" || !isToolResultStatus(tool.status) { - continue - } - results[tool.toolID] = struct{}{} - } - for _, tool := range m.detectedToolCalls { - if tool.toolID == "" { - continue - } - if _, ok := results[tool.toolID]; !ok { - return true - } - } - return false -} - -func (m *chatModel) ensureAssistantEntry() int { - if len(m.entries) > 0 && m.entries[len(m.entries)-1].role == "assistant" { - return len(m.entries) - 1 - } - m.entries = append(m.entries, newChatEntry(chatEntry{role: "assistant"})) - return len(m.entries) - 1 -} - -func (m chatModel) renderTranscript(width int) string { - var b strings.Builder - first := true - for index, entry := range m.entries { - prefix, body := m.renderEntry(entry) - prefixWidth := lipgloss.Width(prefix) - continuation := "" - if prefixWidth > 0 { - continuation = strings.Repeat(" ", prefixWidth) - } - lines := m.renderEntryLinesCached(index, entry, body, width-prefixWidth) - if len(lines) == 0 { - continue - } - if !first { - b.WriteByte('\n') - } - first = false - for i, line := range lines { - if i == 0 { - b.WriteString(prefix) - b.WriteString(line) - } else { - b.WriteString(continuation) - b.WriteString(line) - } - b.WriteByte('\n') - } - } - return b.String() -} - -// renderEntryLinesCached renders (and memoizes) the wrapped lines for an entry. -// Despite the value receiver, the cache write below mutates the caller's entries -// in place: m.entries is a slice, so the value-receiver copy shares its backing -// array with the caller. markEntryDirty invalidates a cached entry by bumping -// its version, which entryRenderKey compares against to decide whether to reuse -// or re-render. -func (m chatModel) renderEntryLinesCached(index int, entry chatEntry, body string, width int) []string { - key := entryRenderKey(entry, width) - if index >= 0 && index < len(m.entries) { - cached := m.entries[index] - if cached.renderKey == key && cached.renderLines != nil { - return cached.renderLines - } - } - - lines := m.renderEntryLines(entry, body, width) - if index >= 0 && index < len(m.entries) { - m.entries[index].renderKey = key - m.entries[index].renderLines = slices.Clone(lines) - } - return lines -} - -func (m chatModel) transcriptLines(width int) []string { - transcript := m.renderTranscript(width) - transcript = strings.TrimRight(transcript, "\n") - if transcript == "" { - return nil - } - return strings.Split(transcript, "\n") -} - -func (m chatModel) visibleTranscriptStartLine(width, available int) int { - return m.visibleTranscriptStartLineForLines(len(m.transcriptLines(width)), available) -} - -func (m chatModel) visibleTranscriptStartLineForLines(total, available int) int { - if total <= available || available <= 0 { - return 0 - } - maxScroll := total - available - scroll := clamp(m.scroll, 0, maxScroll) - return maxScroll - scroll -} - -func (m chatModel) viewWidth() int { - width, _ := m.viewSize() - return width -} - -// defaultSize clamps caller-supplied dimensions to the standard fallbacks -// (80x24) when unset. Free functions that receive width/height as arguments -// (rather than reading them off the model) use this directly; methods use -// viewSize, which reads m.width/m.height first. -func defaultSize(width, height int) (int, int) { - if width <= 0 { - width = defaultViewWidth - } - if height <= 0 { - height = defaultViewHeight - } - return width, height -} - -func (m chatModel) viewSize() (int, int) { - return defaultSize(m.width, m.height) -} - -func (m chatModel) transcriptHeight() int { - width, height := m.viewSize() - lineCount := len(m.transcriptLines(width)) - baseHeaderHeight := 2 - baseMaxHeight := height - baseHeaderHeight - baseBottomLines := m.bottomLines(width, baseMaxHeight) - baseHeight := max(0, baseMaxHeight-len(baseBottomLines)-transcriptInputGap(baseMaxHeight, len(baseBottomLines), lineCount)) - if lineCount <= baseHeight { - return lineCount - } - statusHeaderHeight := 3 - statusMaxHeight := height - statusHeaderHeight - statusBottomLines := m.bottomLines(width, statusMaxHeight) - return max(0, statusMaxHeight-len(statusBottomLines)-transcriptInputGap(statusMaxHeight, len(statusBottomLines), lineCount)) -} - -func (m chatModel) transcriptLayout() (top, height int) { - width, height := m.viewSize() - headerHeight := len(m.headerLines()) - bottomLines := m.bottomLines(width, height-headerHeight) - transcriptLineCount := len(m.transcriptLines(width)) - return headerHeight, max(0, height-headerHeight-len(bottomLines)-transcriptInputGap(height-headerHeight, len(bottomLines), transcriptLineCount)) -} - -func (m chatModel) maxScroll() int { - width, _ := m.viewSize() - return max(0, len(m.transcriptLines(width))-m.transcriptHeight()) -} - -func (m chatModel) bottomLines(width, maxHeight int) []string { - var lines []string - if m.modelPicker != nil { - lines = append(lines, m.renderInlineModelPicker(width)...) - } else { - lines = append(lines, m.completionLines(width)...) - } - - actionStatusLines := m.renderActionStatusLines(width) - approvalLines := m.renderApprovalPromptLines(width) - if maxHeight > 0 { - maxApprovalLines := max(0, maxHeight-len(lines)-len(actionStatusLines)-3) - if len(approvalLines) > maxApprovalLines { - approvalLines = approvalLines[:maxApprovalLines] - } - maxActionStatusLines := max(1, maxHeight-len(lines)-len(approvalLines)-3) - if len(actionStatusLines) > maxActionStatusLines { - actionStatusLines = actionStatusLines[:maxActionStatusLines] - } - } - - lines = append(lines, actionStatusLines...) - lines = append(lines, approvalLines...) - modelLines := m.renderModelStatusLines(width) - fixedLines := len(lines) + 2 - if len(modelLines) > 0 { - fixedLines += len(modelLines) - } - inputBodyLines := maxInputBoxBodyLines - if maxHeight > 0 { - inputBodyLines = min(inputBodyLines, max(1, maxHeight-fixedLines)) - } - inputCursor := m.normalizedInputCursor() - if m.approvalPrompt != nil || m.cloudAuthPrompt != nil { - inputCursor = -1 - } - lines = append(lines, renderInputBoxLines(string(m.input), inputCursor, width, inputBodyLines, m.emptyInputPlaceholder())...) - if len(modelLines) > 0 { - lines = append(lines, modelLines...) - } - return lines -} - -func (m chatModel) renderModelStatusLines(width int) []string { - if m.modelPicker != nil { - return nil - } - var parts []string - if model := strings.TrimSpace(m.opts.Model); model != "" { - parts = append(parts, model) - } - if contextStatus := m.contextStatus(); contextStatus != "" { - parts = append(parts, contextStatus) - } - if notice := m.permissionModeNotice(); notice != "" { - parts = append(parts, notice) - } - if len(parts) == 0 { - return nil - } - indent := inputBoxTextIndent() - lines := wrapChatText(strings.Join(parts, " "), max(20, width-lipgloss.Width(indent))) - for i := range lines { - lines[i] = renderFooterPlainLine(indent + lines[i]) - } - return lines -} - -func (m chatModel) renderActionStatusLines(width int) []string { - if activity := m.activityLine(); activity != "" { - return []string{chatMetaStyle.Render(inputBoxTextIndent() + activity)} - } - if notificationLines := m.renderNotificationLines(width); len(notificationLines) > 0 { - return notificationLines - } - return nil -} - -func transcriptInputGap(maxHeight, bottomLineCount, transcriptLineCount int) int { - const desiredGap = 1 - if transcriptLineCount == 0 { - return 0 - } - if maxHeight <= 0 { - return desiredGap - } - available := maxHeight - bottomLineCount - if available <= 1 { - return 0 - } - return min(desiredGap, available-1) -} - -func (m *chatModel) scrollBy(lines int) { - if lines == 0 { - return - } - m.scroll = clamp(m.scroll+lines, 0, m.maxScroll()) -} - -func stripChatANSI(s string) string { - s = ansi.Strip(s) - return strings.Map(func(r rune) rune { - if r == '\n' || r == '\t' { - return r - } - if unicode.IsControl(r) { - return -1 - } - return r - }, s) -} - -func (m chatModel) normalizedSelectionRange() (chatSelectionPoint, chatSelectionPoint, bool) { - return normalizedSelectionRangeFor(m.selection) -} - -func normalizedSelectionRangeFor(selection chatSelection) (chatSelectionPoint, chatSelectionPoint, bool) { - if !selection.active { - return chatSelectionPoint{}, chatSelectionPoint{}, false - } - start, end := selection.anchor, selection.cursor - if start.line > end.line || (start.line == end.line && start.col > end.col) { - start, end = end, start - } - if start.line == end.line && start.col == end.col { - return chatSelectionPoint{}, chatSelectionPoint{}, false - } - return start, end, true -} - -func displayColumnToRuneIndex(line string, col int) int { - if col <= 0 { - return 0 - } - width := 0 - for i, r := range []rune(line) { - next := width + runewidth.RuneWidth(r) - if col < next { - return i - } - width = next - } - return len([]rune(line)) -} - -func (m chatModel) selectedTranscriptText(width int) string { - start, end, ok := m.normalizedSelectionRange() - if !ok { - return "" - } - lines := m.transcriptLines(width) - if len(lines) == 0 { - return "" - } - start.line = clamp(start.line, 0, len(lines)-1) - end.line = clamp(end.line, 0, len(lines)-1) - var selected []string - for lineIndex := start.line; lineIndex <= end.line; lineIndex++ { - text := transcriptLineTextForSelection(stripChatANSI(lines[lineIndex])) - runes := []rune(text) - startCol, endCol := 0, len(runes) - if lineIndex == start.line { - startCol = displayColumnToRuneIndex(text, start.col) - } - if lineIndex == end.line { - endCol = displayColumnToRuneIndex(text, end.col) - } - if startCol > endCol { - startCol, endCol = endCol, startCol - } - selected = append(selected, string(runes[startCol:endCol])) - } - return strings.TrimRight(strings.Join(selected, "\n"), "\n") -} - -func transcriptLineTextForSelection(text string) string { - firstPrefix := chatMessageIndent + chatUserMessagePrefix - if strings.HasPrefix(text, firstPrefix) { - return chatMessageIndent + strings.TrimPrefix(text, firstPrefix) - } - continuationPrefix := chatMessageIndent + strings.Repeat(" ", lipgloss.Width(chatUserMessagePrefix)) - if strings.HasPrefix(text, continuationPrefix) { - return chatMessageIndent + strings.TrimPrefix(text, continuationPrefix) - } - return text -} - -func (m chatModel) renderEntry(entry chatEntry) (string, string) { - switch entry.role { - case "user": - return "", entry.content - case "assistant": - return "", entry.content - case "thinking": - return chatMetaStyle.Render("•") + " ", thinkingStatusLine(entry) - case "slash": - return chatMetaStyle.Render("•") + " ", entry.content - case "compaction_summary": - prefix := toolStatusStyle(entry.status).Render("•") + " " - return prefix, compactionSummaryStatusLine(entry) - case "tool": - prefix := toolStatusStyle(entry.status).Render("•") + " " - return prefix, toolStatusLine(entry) - case "tool_group": - prefix := toolGroupPrefixStyle(entry).Render("•") + " " - return prefix, toolGroupStatusLine(entry) - case "error": - return chatErrorStyle.Render("err ") + " ", entry.content - case "system": - return "", entry.content - default: - return "", entry.content - } -} - -func (m chatModel) renderEntryLines(entry chatEntry, body string, width int) []string { - if width < 20 { - width = 20 - } - switch entry.role { - case "assistant": - innerWidth := max(1, width-lipgloss.Width(chatMessageIndent)) - lines := indentLines(splitRenderedBody(renderMarkdownForView(body, innerWidth)), chatMessageIndent) - lines = append(lines, indentLines(renderMetricsLines(entry.metrics, innerWidth), chatMessageIndent)...) - return lines - case "thinking": - return renderThinkingLines(entry, width) - case "system", "slash": - return splitRenderedBody(renderMarkdownForView(body, width)) - case "user": - return renderUserMessageLines(body, width) - case "compaction_summary": - return renderCompactionSummaryLines(entry, width) - case "tool": - if entryHasExpandableOutput(entry) { - return renderToolResultLines(entry, width) - } - if isToolResultStatus(entry.status) { - return renderToolResultLines(entry, width) - } - return wrapChatText(body, width) - case "tool_group": - return renderToolGroupLines(entry, width) - default: - return wrapChatText(body, width) - } -} - -func renderUserMessageLines(content string, width int) []string { - if width < 20 { - width = 20 - } - firstPrefix := chatMessageIndent + chatUserMessagePrefix - continuationPrefix := chatMessageIndent + strings.Repeat(" ", lipgloss.Width(chatUserMessagePrefix)) - innerWidth := max(1, width-lipgloss.Width(firstPrefix)) - lines := wrapChatText(content, innerWidth) - for i, line := range lines { - prefix := continuationPrefix - if i == 0 { - prefix = firstPrefix - } - lines[i] = chatUserBlockStyle.Render(padRenderedLine(prefix+line, width)) - } - return lines -} - -func renderMetricsLines(metrics *api.Metrics, width int) []string { - summary := metricsSummaryLines(metrics) - if len(summary) == 0 { - return nil - } - var lines []string - for _, line := range summary { - for _, wrapped := range wrapChatText(line, width) { - lines = append(lines, chatMetaStyle.Render(wrapped)) - } - } - return lines -} - -func metricsSummaryLines(metrics *api.Metrics) []string { - if metrics == nil || metricsEmpty(*metrics) { - return nil - } - var lines []string - if metrics.TotalDuration > 0 { - lines = append(lines, fmt.Sprintf("total duration: %v", metrics.TotalDuration)) - } - if metrics.LoadDuration > 0 { - lines = append(lines, fmt.Sprintf("load duration: %v", metrics.LoadDuration)) - } - if metrics.PromptEvalCount > 0 { - lines = append(lines, fmt.Sprintf("prompt eval count: %d token(s)", metrics.PromptEvalCount)) - } - cached := 0 - if metrics.PromptEvalCachedCount != nil { - cached = *metrics.PromptEvalCachedCount - } - if cached > 0 { - lines = append(lines, fmt.Sprintf("prompt eval cached: %d token(s)", cached)) - } - if metrics.PromptEvalDuration > 0 { - lines = append(lines, fmt.Sprintf("prompt eval duration: %s", metrics.PromptEvalDuration)) - uncached := max(0, metrics.PromptEvalCount-cached) - lines = append(lines, fmt.Sprintf("prompt eval rate: %.2f tokens/s", float64(uncached)/metrics.PromptEvalDuration.Seconds())) - } - if metrics.EvalCount > 0 { - lines = append(lines, fmt.Sprintf("eval count: %d token(s)", metrics.EvalCount)) - } - if metrics.EvalDuration > 0 { - lines = append(lines, fmt.Sprintf("eval duration: %s", metrics.EvalDuration)) - lines = append(lines, fmt.Sprintf("eval rate: %.2f tokens/s", float64(metrics.EvalCount)/metrics.EvalDuration.Seconds())) - } - return lines -} - -func metricsEmpty(metrics api.Metrics) bool { - cached := 0 - if metrics.PromptEvalCachedCount != nil { - cached = *metrics.PromptEvalCachedCount - } - return metrics.TotalDuration <= 0 && - metrics.LoadDuration <= 0 && - metrics.PromptEvalCount <= 0 && - cached <= 0 && - metrics.PromptEvalDuration <= 0 && - metrics.EvalCount <= 0 && - metrics.EvalDuration <= 0 -} - -func historyRoleStyle(role string) lipgloss.Style { - switch role { - case "system": - return chatHistorySystemRoleStyle - case "user": - return chatHistoryUserRoleStyle - case "assistant": - return chatHistoryAssistantRoleStyle - case "tool": - return chatHistoryToolRoleStyle - default: - return chatHistoryTitleStyle - } -} - -func renderToolResultLines(entry chatEntry, width int) []string { - lines := wrapChatText(toolStatusLine(entry), width) - if !entry.expanded { - return lines - } - - detailLines := renderToolCallDetailLines(entry, width) - if len(detailLines) > 0 { - lines = append(lines, "") - lines = append(lines, detailLines...) - } - if strings.TrimSpace(entry.content) != "" { - lines = append(lines, "") - lines = append(lines, renderToolOutputLines(entry, entry.content, width)...) - } - return lines -} - -func renderToolGroupLines(entry chatEntry, width int) []string { - lines := wrapChatText(toolGroupStatusLine(entry), width) - if !entry.expanded { - return lines - } - - for i, tool := range entry.tools { - if i > 0 { - lines = append(lines, "") - } - lines = append(lines, " "+toolGroupChildStatusLine(tool)) - if detailLines := renderToolCallDetailLines(tool, width-4); len(detailLines) > 0 { - lines = append(lines, indentLines(detailLines, " ")...) - } - if strings.TrimSpace(tool.content) == "" { - continue - } - lines = append(lines, indentLines(renderToolOutputLines(tool, tool.content, width-4), " ")...) - } - return lines -} - -func renderCompactionSummaryLines(entry chatEntry, width int) []string { - lines := wrapChatText(compactionSummaryStatusLine(entry), width) - if !entry.expanded || strings.TrimSpace(entry.content) == "" { - return lines - } - lines = append(lines, "") - lines = append(lines, indentLines(splitRenderedBody(renderMarkdownForView(entry.content, width-2)), " ")...) - return lines -} - -func compactionSummaryStatusLine(entry chatEntry) string { - segment := toolStatusStyle(entry.status).Render(toolStatusLabel(entry)) - if segment == "" { - return "Compacted summary" - } - return fmt.Sprintf("Compacted summary %s", segment) -} - -func renderThinkingLines(entry chatEntry, width int) []string { - lines := wrapChatText(thinkingStatusLine(entry), width) - if !entry.expanded || entry.content == "" { - return lines - } - lines = append(lines, "") - if entry.status == "running" { - body := styleLines(renderLiveThinkingLines(entry.content, width), chatToolOutputStyle) - lines = append(lines, body...) - return lines - } - body := styleLines(splitRenderedBody(renderMarkdownForView(stripChatANSI(entry.content), width)), chatToolOutputStyle) - lines = append(lines, body...) - return lines -} - -// renderLiveThinkingLines keeps each streaming redraw bounded. Completed -// traces use the normal Markdown renderer when explicitly reopened, but a -// running trace is rendered as plain text so an ever-growing document is not -// reparsed for every delta. -func renderLiveThinkingLines(content string, width int) []string { - content, omitted := liveThinkingTail(content, maxLiveThinkingRunes) - content = stripChatANSI(content) - lines := wrapChatText(content, width) - if omitted { - lines = append([]string{"… earlier thinking omitted while streaming"}, lines...) - } - return lines -} - -func liveThinkingTail(content string, limit int) (string, bool) { - if limit <= 0 { - return content, false - } - start := len(content) - for range limit { - if start == 0 { - return content, false - } - _, size := utf8.DecodeLastRuneInString(content[:start]) - start -= size - } - tail := content[start:] - if newline := strings.IndexByte(tail, '\n'); newline >= 0 { - tail = tail[newline+1:] - } - return tail, true -} - -func thinkingStatusLine(entry chatEntry) string { - if entry.status != "running" { - return thoughtLabel(entry.tokenCount, entry.expanded) - } - - label := "Thinking" - if strings.TrimSpace(entry.label) != "" { - label = entry.label - } - return label -} - -func thoughtLabel(tokens int, expanded bool) string { - if !expanded || tokens <= 0 { - return "Thought" - } - return "Thought (" + formatTokenCount(tokens) + ")" -} - -func (m chatModel) thinkingLabel() string { - return thinkingActivityLabel(m.thinkingTokens) -} - -func thinkingActivityLabel(tokens int) string { - if tokens > 0 { - return "Thinking ↓ " + formatTokenCount(tokens) - } - return "Thinking" -} - -func (m *chatModel) syncThinkingEntry(content string) { - idx := -1 - if len(m.entries) > 0 && m.entries[len(m.entries)-1].role == "thinking" && m.entries[len(m.entries)-1].status == "running" { - idx = len(m.entries) - 1 - } - if idx < 0 { - if strings.TrimSpace(content) == "" { - return - } - m.entries = append(m.entries, newChatEntry(chatEntry{role: "thinking", status: "running"})) - idx = len(m.entries) - 1 - } - m.entries[idx].content = content - m.entries[idx].label = m.thinkingLabel() - m.entries[idx].status = "running" - m.entries[idx].tokenCount = m.thinkingTokens - m.entries[idx].expanded = true - m.markEntryDirty(idx) -} - -func (m *chatModel) applyThinkingDetails() { - for i := range m.entries { - entry := &m.entries[i] - if entry.role != "thinking" || entry.status == "running" || strings.TrimSpace(entry.content) == "" || entry.expanded == m.thinkingDetailsOpen { - continue - } - entry.expanded = m.thinkingDetailsOpen - m.markEntryDirty(i) - } -} - -func (m *chatModel) finishThinkingEntry() { - if len(m.entries) == 0 { - return - } - idx := len(m.entries) - 1 - if m.entries[idx].role != "thinking" || m.entries[idx].status != "running" { - return - } - m.entries[idx].status = "done" - m.entries[idx].label = m.thinkingLabel() - m.entries[idx].tokenCount = m.thinkingTokens - m.entries[idx].expanded = m.thinkingDetailsOpen - m.markEntryDirty(idx) -} - -func toolGroupChildStatusLine(entry chatEntry) string { - label := toolGroupChildStatusLabel(entry) - - segment := renderToolStatusSegment(entry) - if segment == "" { - return boldToolInvocationName(label) - } - return fmt.Sprintf("%s %s", boldToolInvocationName(label), segment) -} - -func toolGroupChildStatusLabel(entry chatEntry) string { - if strings.TrimSpace(entry.label) != "" { - return entry.label - } - if strings.TrimSpace(entry.detail) != "" { - return toolInvocationLabel(entry.detail, entry.args) - } - return toolEntryStatusLabel(entry) -} - -func boldToolInvocationName(label string) string { - name, rest, ok := strings.Cut(label, "(") - if !ok || name == "" { - return chatHeaderStyle.Render(label) - } - return chatHeaderStyle.Render(name) + "(" + rest -} - -func renderToolOutputLines(entry chatEntry, output string, width int) []string { - output = stripInternalToolTruncationMarkers(output) - output = truncateCtrlOToolOutput(output) - if looksLikeUnifiedDiff(output) { - return splitRenderedBody(renderDiffForView(output, width)) - } - if toolOutputUsesMarkdown(entry.detail) { - return styleLines(splitRenderedBody(renderMarkdownForView(output, width)), chatToolOutputStyle) - } - return styleLines(wrapChatText(output, width), chatToolOutputStyle) -} - -func styleLines(lines []string, style lipgloss.Style) []string { - for i := range lines { - lines[i] = style.Render(lines[i]) - } - return lines -} - -func truncateCtrlOToolOutput(output string) string { - runes := []rune(output) - if len(runes) <= maxCtrlOToolOutputRunes { - return output - } - return string(runes[:maxCtrlOToolOutputRunes-3]) + "..." -} - -func stripInternalToolTruncationMarkers(output string) string { - lines := strings.Split(output, "\n") - filtered := lines[:0] - for _, line := range lines { - trimmed := strings.TrimSpace(line) - if strings.HasPrefix(trimmed, "[tool output truncated: ") && strings.HasSuffix(trimmed, "]") { - continue - } - filtered = append(filtered, line) - } - return strings.TrimSpace(strings.Join(filtered, "\n")) -} - -func looksLikeUnifiedDiff(output string) bool { - lines := strings.Split(output, "\n") - hasOldFile := false - hasNewFile := false - hasHunk := false - for _, line := range lines { - switch { - case strings.HasPrefix(line, "diff --git "): - return true - case strings.HasPrefix(line, "--- "): - hasOldFile = true - case strings.HasPrefix(line, "+++ "): - hasNewFile = true - case strings.HasPrefix(line, "@@ "): - hasHunk = true - } - } - return hasHunk || (hasOldFile && hasNewFile) -} - -func renderDiffForView(diff string, width int) string { - if width < 20 { - width = 20 - } - lines := strings.Split(strings.TrimRight(diff, "\n"), "\n") - rendered := make([]string, 0, len(lines)) - for _, line := range lines { - style := diffLineStyle(line) - for _, wrapped := range wrapChatText(line, width) { - rendered = append(rendered, style.Render(wrapped)) - } - } - return strings.Join(rendered, "\n") -} - -func diffLineStyle(line string) lipgloss.Style { - switch { - case strings.HasPrefix(line, "diff --git "), - strings.HasPrefix(line, "--- "), - strings.HasPrefix(line, "+++ "): - return chatDiffFileStyle - case strings.HasPrefix(line, "@@ "): - return chatDiffHunkStyle - case strings.HasPrefix(line, "+"): - return chatDiffAddStyle - case strings.HasPrefix(line, "-"): - return chatDiffDeleteStyle - case strings.HasPrefix(line, "index "), - strings.HasPrefix(line, "new file "), - strings.HasPrefix(line, "deleted file "), - strings.HasPrefix(line, "similarity index "), - strings.HasPrefix(line, "rename from "), - strings.HasPrefix(line, "rename to "), - strings.HasPrefix(line, "\\ "): - return chatDiffMetaStyle - default: - return chatToolStyle - } -} - -func isToolActiveStatus(status string) bool { - return status == "queued" || status == "running" || status == "approval" -} - -func isToolResultStatus(status string) bool { - return status == "done" || status == "error" || status == "denied" || status == "disabled" -} - -// renderToolStatusSegment returns the styled status text for a tool entry, -// or "" when the entry has no status word to show (e.g. a completed or failed -// tool, where the colored dot alone conveys the result). -func renderToolStatusSegment(entry chatEntry) string { - s := toolStatusLabel(entry) - if s == "" { - return "" - } - return chatMetaStyle.Render(s) -} - -func toolStatusLine(entry chatEntry) string { - label := toolEntryStatusLabel(entry) - - segment := renderToolStatusSegment(entry) - if segment == "" { - return label - } - return fmt.Sprintf("%s %s", label, segment) -} - -func toolEntryStatusLabel(entry chatEntry) string { - if isShellToolName(entry.detail) { - switch entry.status { - case "approval": - if entry.label != "" { - return entry.label - } - return toolInvocationLabel(entry.detail, entry.args) - case "denied": - if entry.label != "" { - return entry.label + " denied" - } - return toolInvocationLabel(entry.detail, entry.args) + " denied" - case "queued", "running", "done", "error": - if entry.label != "" { - return entry.label - } - return toolInvocationLabel(entry.detail, entry.args) - } - } - if entry.status == "denied" { - if entry.label != "" { - return entry.label + " denied" - } - return toolDisplayName(entry.detail) + " denied" - } - if entry.label != "" { - return entry.label - } - return toolDisplayName(entry.detail) -} - -func toolGroupStatusLine(entry chatEntry) string { - label := entry.label - if label == "" || strings.HasPrefix(label, "Tool calls (") { - label = toolGroupSummary(entry.tools) - } - - segment := renderToolStatusSegment(entry) - if segment == "" { - return label - } - return fmt.Sprintf("%s %s", label, segment) -} - -func toolGroupSummary(tools []chatEntry) string { - if len(tools) == 0 { - return "Used tools" - } - - type actionCount struct { - action string - count int - } - - var counts []actionCount - indexes := map[string]int{} - for _, tool := range tools { - action := toolActionForEntry(tool) - if index, ok := indexes[action]; ok { - counts[index].count++ - continue - } - indexes[action] = len(counts) - counts = append(counts, actionCount{action: action, count: 1}) - } - - phrases := make([]string, 0, len(counts)) - for _, count := range counts { - phrases = append(phrases, toolActionPhrase(count.action, count.count)) - } - return joinToolActionPhrases(phrases) -} - -func toolActionForEntry(tool chatEntry) string { - if tool.status == "denied" { - if isShellToolName(tool.detail) || strings.Contains(strings.ToLower(tool.label), "bash(") || strings.Contains(strings.ToLower(tool.label), "powershell(") { - return "denied_command" - } - return "denied_tool" - } - if tool.status == "disabled" { - if isShellToolName(tool.detail) || strings.Contains(strings.ToLower(tool.label), "bash(") || strings.Contains(strings.ToLower(tool.label), "powershell(") { - return "disabled_command" - } - return "disabled_tool" - } - action := toolAction(tool.detail) - if action == "" { - action = toolAction(tool.label) - } - if action == "" { - action = "tool" - } - return action -} - -func toolAction(name string) string { - name = strings.TrimSpace(strings.ToLower(name)) - switch { - case isShellToolName(name): - return "command" - case strings.Contains(name, "bash") || strings.Contains(name, "powershell"): - return "command" - case strings.HasPrefix(name, "edit("): - return "edit" - case strings.HasPrefix(name, "read("): - return "read" - case strings.HasPrefix(name, "list("): - return "list" - case strings.HasPrefix(name, "web search("): - return "search" - case strings.HasPrefix(name, "web fetch("): - return "fetch" - case strings.HasPrefix(name, "skill("): - return "skill" - } - switch name { - case "edit": - return "edit" - case "read": - return "read" - case "list": - return "list" - case "web_search": - return "search" - case "web_fetch": - return "fetch" - case "skill": - return "skill" - default: - return "tool" - } -} - -func toolActionPhrase(action string, count int) string { - plural := count != 1 - switch action { - case "denied_command": - if plural { - return fmt.Sprintf("Denied %d commands", count) - } - return "Denied a command" - case "denied_tool": - if plural { - return fmt.Sprintf("Denied %d tools", count) - } - return "Denied a tool" - case "disabled_command": - if plural { - return fmt.Sprintf("Skipped %d commands", count) - } - return "Skipped a command" - case "disabled_tool": - if plural { - return fmt.Sprintf("Skipped %d tools", count) - } - return "Skipped a tool" - case "command": - if plural { - return fmt.Sprintf("Ran %d commands", count) - } - return "Ran 1 command" - case "edit": - if plural { - return fmt.Sprintf("Edited %d files", count) - } - return "Edited a file" - case "read": - if plural { - return fmt.Sprintf("Read %d files", count) - } - return "Read a file" - case "list": - if plural { - return fmt.Sprintf("Listed files %d times", count) - } - return "Listed files" - case "search": - if plural { - return fmt.Sprintf("Searched the web %d times", count) - } - return "Searched the web" - case "fetch": - if plural { - return fmt.Sprintf("Fetched %d URLs", count) - } - return "Fetched a URL" - case "skill": - if plural { - return fmt.Sprintf("Loaded %d skills", count) - } - return "Loaded a skill" - default: - if plural { - return fmt.Sprintf("Used %d tools", count) - } - return "Used a tool" - } -} - -func joinToolActionPhrases(phrases []string) string { - switch len(phrases) { - case 0: - return "Used tools" - case 1: - return phrases[0] - case 2: - return phrases[0] + " and " + lowerInitial(phrases[1]) - default: - for i := 1; i < len(phrases); i++ { - phrases[i] = lowerInitial(phrases[i]) - } - return strings.Join(phrases[:len(phrases)-1], ", ") + ", and " + phrases[len(phrases)-1] - } -} - -func lowerInitial(s string) string { - if s == "" { - return s - } - runes := []rune(s) - runes[0] = []rune(strings.ToLower(string(runes[0])))[0] - return string(runes) -} - -func toolGroupPrefixStyle(entry chatEntry) lipgloss.Style { - succeeded, failed, denied := toolGroupResultCounts(entry.tools) - switch { - case succeeded > 0 && (failed > 0 || denied > 0): - return chatToolMixedStyle - case succeeded > 0: - return chatToolDoneStyle - case denied > 0 && failed == 0: - return toolStatusStyle("denied") - default: - return toolStatusStyle(entry.status) - } -} - -func toolGroupResultCounts(tools []chatEntry) (succeeded int, failed int, denied int) { - for _, tool := range tools { - if tool.status == "denied" || tool.status == "disabled" { - denied++ - continue - } - if tool.err != "" || tool.status == "error" { - failed++ - continue - } - if tool.status == "done" { - succeeded++ - } - } - return succeeded, failed, denied -} - -func toolStatusLabel(entry chatEntry) string { - if entry.status == "approval" { - return "needs approval" - } - if entry.status == "running" || entry.status == "queued" { - return "" - } - return "" -} - -func toolStatusStyle(status string) lipgloss.Style { - switch status { - case "queued", "running", "approval": - return chatToolRunningStyle - case "done": - return chatToolDoneStyle - case "error": - return chatErrorStyle - case "denied", "disabled": - return chatToolMixedStyle - default: - return chatMetaStyle - } -} - -func toolInvocationLabel(name string, args map[string]any) string { - displayName := toolDisplayName(name) - for _, key := range []string{"query", "url", "command", "path", "name"} { - if value, ok := rawStringArg(args, key); ok { - if isShellToolName(name) && key == "command" { - value = truncateRunes(value, 100) - } - return fmt.Sprintf("%s(%s)", displayName, strconv.Quote(value)) - } - } - if len(args) == 0 { - return displayName - } - return fmt.Sprintf("%s(%s)", displayName, formatDisplayArgs(args)) -} - -func toolDisplayName(name string) string { - switch name { - case "web_search": - return "Web Search" - case "web_fetch": - return "Web Fetch" - case "bash": - return "Bash" - case "powershell": - return "PowerShell" - case "read": - return "Read" - case "list": - return "List" - case "edit": - return "Edit" - default: - if name == "" { - return "Tool" - } - return name - } -} - -func toolOutputUsesMarkdown(name string) bool { - switch name { - case "read", "skill", "web_search", "web_fetch": - return true - default: - return false - } -} - -func renderToolCallDetailLines(entry chatEntry, width int) []string { - if len(entry.args) == 0 { - return nil - } - if isShellToolName(entry.detail) { - if command, ok := rawStringArg(entry.args, "command"); ok { - return renderToolCallArgLine(shellPromptPrefix(entry.detail)+command, width) - } - } - switch entry.detail { - case "web_search": - if query, ok := rawStringArg(entry.args, "query"); ok { - return renderToolCallArgLine("query: "+query, width) - } - case "web_fetch": - if targetURL, ok := rawStringArg(entry.args, "url"); ok { - return renderToolCallArgLine("url: "+targetURL, width) - } - } - return renderToolCallArgs(entry.args, width) -} - -func isShellToolName(name string) bool { - return name == "bash" || name == "powershell" -} - -func formatDisplayArgs(args map[string]any) string { - keys := make([]string, 0, len(args)) - for key := range args { - keys = append(keys, key) - } - sort.Strings(keys) - - parts := make([]string, 0, len(keys)) - for _, key := range keys { - value := truncateRunes(fmt.Sprintf("%v", args[key]), 100) - parts = append(parts, fmt.Sprintf("%s=%s", key, strconv.Quote(value))) - } - return strings.Join(parts, ", ") -} - -func shellPromptPrefix(toolName string) string { - if toolName == "powershell" { - return "PS> " - } - return "$ " -} - -func renderToolCallArgs(args map[string]any, width int) []string { - keys := make([]string, 0, len(args)) - for key := range args { - keys = append(keys, key) - } - sort.Strings(keys) - - var lines []string - for _, key := range keys { - value := toolArgDisplayValue(args[key]) - if strings.Contains(value, "\n") { - lines = append(lines, chatMetaStyle.Render(key+":")) - lines = append(lines, indentLines(renderToolCallArgLine(value, max(20, width-2)), " ")...) - continue - } - lines = append(lines, renderToolCallArgLine(key+": "+value, width)...) - } - return lines -} - -func renderToolCallArgLine(line string, width int) []string { - wrapped := wrapChatText(line, width) - for i := range wrapped { - wrapped[i] = chatMetaStyle.Render(wrapped[i]) - } - return wrapped -} - -func toolArgDisplayValue(value any) string { - if value == nil { - return "null" - } - if text, ok := value.(string); ok { - return text - } - data, err := json.MarshalIndent(value, "", " ") - if err == nil { - return string(data) - } - return fmt.Sprint(value) -} - -func rawStringArg(args map[string]any, key string) (string, bool) { - value, ok := args[key].(string) - if !ok || strings.TrimSpace(value) == "" { - return "", false - } - return value, true -} - -func isDeniedToolResult(value string) bool { - value = strings.ToLower(strings.TrimSpace(value)) - return strings.Contains(value, "tool execution denied") || - strings.Contains(value, "tool approval canceled") -} - -func isDisabledToolResult(value string) bool { - return strings.Contains(strings.ToLower(strings.TrimSpace(value)), "tool execution disabled") -} - -func truncateRunes(value string, limit int) string { - runes := []rune(value) - if len(runes) <= limit { - return value - } - return string(runes[:limit]) + "..." -} - -func (m chatModel) notificationLine() string { - status := strings.TrimSpace(m.status) - if status == "" || status == "ready" { - return "" - } - if m.running || m.compacting || m.approvalPrompt != nil { - return "" - } - switch status { - case "running", "compacting", "approval required", "full access enabled", "review mode enabled": - return "" - default: - return status - } -} - -func (m chatModel) renderNotificationLines(width int) []string { - line := m.notificationLine() - if line == "" { - return nil - } - indent := inputBoxTextIndent() - lines := wrapChatText(line, max(20, width-lipgloss.Width(indent))) - for i, wrapped := range lines { - lines[i] = chatNotificationStyle.Render(indent + wrapped) - } - return lines -} - -func inputBoxTextIndent() string { - return strings.Repeat(" ", inputBoxHorizontalPadding+1) -} - -func renderFooterPlainLine(line string) string { - const fullAccess = "full access" - if !strings.Contains(line, fullAccess) { - return chatFooterStyle.Render(line) - } - - var b strings.Builder - for { - before, after, ok := strings.Cut(line, fullAccess) - if before != "" { - b.WriteString(chatFooterStyle.Render(before)) - } - if !ok { - break - } - b.WriteString(chatFullAccessStyle.Render(fullAccess)) - line = after - } - return b.String() -} - -func (m chatModel) permissionModeNotice() string { - if notice := strings.TrimSpace(m.permissionNotice); notice != "" { - return notice - } - switch strings.TrimSpace(m.status) { - case "full access enabled", "review mode enabled": - return strings.TrimSpace(m.status) - } - if m.allowAllToolsEnabled() && m.notificationLine() == "" { - return "full access enabled" - } - return "" -} - -func (m *chatModel) refreshContextWindowTokens(modelName string) { - if m == nil || m.opts.ContextWindowTokensForModel == nil { - return - } - modelName = strings.TrimSpace(modelName) - if modelName == "" { - return - } - ctx := m.ctx - if ctx == nil { - ctx = context.Background() - } - tokens := m.opts.ContextWindowTokensForModel(ctx, modelName, m.opts.ContextWindowTokens) - m.updateContextWindowTokens(tokens) -} - -func (m *chatModel) updateContextWindowTokens(tokens int) { - if tokens <= 0 || tokens == m.opts.ContextWindowTokens { - return - } - m.opts.ContextWindowTokens = tokens - if compactor, ok := m.opts.Compactor.(*coreagent.SimpleCompactor); ok && compactor != nil { - compactor.Options.ContextWindowTokens = tokens - } -} - -func (m chatModel) currentWorkingDir() string { - if strings.TrimSpace(m.workingDir) != "" { - return m.workingDir - } - return m.opts.WorkingDir -} - -func (m chatModel) activityLine() string { - if m.approvalPrompt != nil { - return "" - } - if m.preloadingModel != "" && !m.running && !m.compacting { - return "" - } - if !m.running && !m.compacting && m.preloadingModel == "" && m.approvalPrompt == nil { - return "" - } - if m.thinking && len(m.entries) > 0 { - entry := m.entries[len(m.entries)-1] - if entry.role == "thinking" && entry.status == "running" { - return "" - } - } - label := m.activityLabel() - if label == "" { - if m.awaitingToolStart() { - return statusWithSpinner(m.spinnerFrame(), "Working") - } - if m.awaitingModel { - return statusWithSpinner(m.spinnerFrame(), "Working") - } - if !m.waitingForModel() || m.spinner < idleWorkingDelayTicks { - return "" - } - if m.preloadingModel != "" { - return statusWithSpinner(m.spinnerFrame(), "Working") - } - return statusWithSpinner(m.spinnerFrame(), "Working") - } - if m.thinking { - return label - } - return statusWithSpinner(m.spinnerFrame(), label) -} - -func (m chatModel) activityLabel() string { - if m.status == "canceling" { - return "canceling" - } - if m.compacting { - if m.compactingTokens > 0 { - return "Compacting " + formatTokenCount(m.compactingTokens) - } - return "Compacting" - } - if m.thinking { - return thinkingActivityLabel(m.thinkingTokens) - } - start := m.currentTurnEntryStart() - for i := len(m.entries) - 1; i >= start; i-- { - entry := m.entries[i] - switch entry.role { - case "tool": - if isToolActiveStatus(entry.status) { - return "" - } - case "tool_group": - if entryHasActiveTool(entry) { - return "" - } - case "assistant": - if entry.content != "" { - return "" - } - } - } - return "" -} - -func (m chatModel) waitingForModel() bool { - if m.preloadingModel != "" && !m.compacting && m.approvalPrompt == nil && m.status != "canceling" { - return true - } - if !m.running || m.compacting || m.approvalPrompt != nil || m.thinking || m.status == "canceling" { - return false - } - if m.awaitingModel { - return true - } - start := m.currentTurnEntryStart() - for i := len(m.entries) - 1; i >= start; i-- { - entry := m.entries[i] - switch entry.role { - case "tool": - if isToolActiveStatus(entry.status) { - return false - } - case "tool_group": - if entryHasActiveTool(entry) { - return false - } - case "assistant": - return true - } - } - return true -} - -func (m chatModel) currentTurnEntryStart() int { - for i := len(m.entries) - 1; i >= 0; i-- { - if m.entries[i].role == "user" { - return i + 1 - } - } - return 0 -} - -func (m *chatModel) applyResponseMetrics(response *api.ChatResponse) { - if response == nil { - return - } - tokens := 0 - if response.PromptEvalCount > 0 { - tokens = response.PromptEvalCount - } - if response.EvalCount > 0 { - tokens += response.EvalCount - } - if tokens <= 0 { - return - } - if m.running { - if tokens > m.contextTokens { - m.contextTokens = tokens - } - return - } - if tokens > 0 { - m.contextTokens = tokens - m.contextEstimate = false - } -} - -func (m chatModel) estimatePromptTokens(messages []api.Message, systemPrompt string) int { - if strings.TrimSpace(systemPrompt) == "" { - systemPrompt = m.systemPrompt("") - } - var tools api.Tools - if m.opts.Tools != nil { - tools = m.opts.Tools.Tools() - } - return estimatePromptTokenCount(systemPrompt, messages, tools, m.opts.Format) -} - -func estimatePromptTokenCount(systemPrompt string, messages []api.Message, tools api.Tools, format string) int { - total := approximateTokenCount(systemPrompt) - for _, msg := range messages { - total += approximateTokenCount(msg.Role) - total += approximateTokenCount(msg.Content) - total += approximateTokenCount(msg.Thinking) - total += approximateTokenCount(msg.ToolName) - total += approximateTokenCount(msg.ToolCallID) - for _, call := range msg.ToolCalls { - total += approximateTokenCount(call.Function.Name) - total += approximateTokenCount(call.Function.Arguments.String()) - } - } - total += approximateTokenCount(tools.String()) - total += approximateTokenCount(format) - return total -} - -func approximateTokenCount(text string) int { - n := len([]rune(text)) - if n <= 0 { - return 0 - } - return max(1, (n+3)/4) -} - -func formatTokenCount(count int) string { - if count == 1 { - return "1 token" - } - return fmt.Sprintf("%d tokens", count) -} - -func (m chatModel) contextStatus() string { - window := m.displayContextWindowTokens() - if window <= 0 { - return "" - } - used := max(m.contextTokens, 0) - percent := 0 - if window > 0 { - percent = (used*100 + window/2) / window - } - - prefix := "" - if m.contextEstimate { - prefix = "~" - } - - threshold := coreagent.ResolveCompactionThreshold(m.opts.CompactionThreshold) - compactAt := int(float64(window)*threshold + 0.999999) - if compactAt <= 0 || compactAt > window { - compactAt = window - } - - if used >= compactAt { - return fmt.Sprintf("ctx %s%s / %s (%d%% used)", prefix, formatContextTokenCount(used), formatContextTokenCount(window), percent) - } - - noticeDistance := int(float64(window)*0.1 + 0.999999) - if noticeDistance < 1 { - noticeDistance = 1 - } - if compactAt-used <= noticeDistance { - return fmt.Sprintf("ctx %s%s / %s (%d%% used)", prefix, formatContextTokenCount(used), formatContextTokenCount(window), percent) - } - - if percent > 60 { - return fmt.Sprintf("ctx %s%s / %s (%d%% used)", prefix, formatContextTokenCount(used), formatContextTokenCount(window), percent) - } - - return "" -} - -func (m chatModel) displayContextWindowTokens() int { - if n := chatIntOption(m.opts.Options, "num_ctx"); n > 0 { - return n - } - return max(0, m.opts.ContextWindowTokens) -} - -func chatIntOption(options map[string]any, key string) int { - if options == nil { - return 0 - } - switch v := options[key].(type) { - case int: - return max(0, v) - case int32: - return max(0, int(v)) - case int64: - return max(0, int(v)) - case uint: - return int(v) - case uint32: - return int(v) - case uint64: - return int(v) - case float64: - if v == float64(int(v)) { - return max(0, int(v)) - } - case string: - n, err := strconv.Atoi(strings.TrimSpace(v)) - if err == nil { - return max(0, n) - } - } - return 0 -} - -func formatContextTokenCount(value int) string { - sign := "" - if value < 0 { - sign = "-" - value = -value - } - if value >= 950_000 { - return fmt.Sprintf("%s%dM", sign, int(float64(value)/1_000_000+0.5)) - } - if value >= 10_240 { - return fmt.Sprintf("%s%dk", sign, int(float64(value)/1024+0.5)) - } - return sign + formatInteger(value) -} - -func formatInteger(value int) string { - sign := "" - if value < 0 { - sign = "-" - value = -value - } - s := strconv.Itoa(value) - if len(s) <= 3 { - return sign + s - } - var b strings.Builder - b.WriteString(sign) - first := len(s) % 3 - if first == 0 { - first = 3 - } - b.WriteString(s[:first]) - for i := first; i < len(s); i += 3 { - b.WriteByte(',') - b.WriteString(s[i : i+3]) - } - return b.String() -} - -func (m chatModel) spinnerFrame() string { - if len(chatSpinnerFrames) == 0 { - return "" - } - return chatSpinnerFrames[m.spinner%len(chatSpinnerFrames)] -} - -func statusWithSpinner(frame, label string) string { - label = strings.TrimSpace(label) - if label == "" { - return strings.TrimSpace(frame) - } - return label + frame -} - -func renderFullFrame(content string, width, height int) string { - width, height = defaultSize(width, height) - rendered := lipgloss.NewStyle().MaxWidth(width).Render(content) - lines := strings.Split(strings.TrimRight(rendered, "\n"), "\n") - if len(lines) > height { - lines = lines[:height] - } - for len(lines) < height { - lines = append(lines, "") - } - return strings.Join(lines, "\n") -} - -func renderFrameLines(lines []string, width, height int) string { - width, height = defaultSize(width, height) - if len(lines) > height { - lines = lines[:height] - } - out := make([]string, 0, height) - for _, line := range lines { - out = append(out, padRenderedLine(clipRenderedLine(line, width), width)) - } - for len(out) < height { - out = append(out, strings.Repeat(" ", width)) - } - return strings.Join(out, "\n") -} - -func truncateRenderedLine(line string, width int) string { - if width <= 0 || lipgloss.Width(line) <= width { - return line - } - return lipgloss.NewStyle().MaxWidth(width).Render(line) -} - -func clipRenderedLine(line string, width int) string { - if width <= 0 { - return "" - } - line, _, _ = strings.Cut(line, "\n") - if lipgloss.Width(line) <= width { - return line - } - clipped := lipgloss.NewStyle().MaxWidth(width).Render(line) - clipped, _, _ = strings.Cut(clipped, "\n") - return clipped -} - -func clamp(value, minValue, maxValue int) int { - if value < minValue { - return minValue - } - if value > maxValue { - return maxValue - } - return value -} - -func newChatEntry(entry chatEntry) chatEntry { - if entry.version <= 0 { - entry.version = 1 - } - return entry -} - -func newSlashEntry(content string) chatEntry { - return newChatEntry(chatEntry{role: "slash", content: content}) -} - -func (m *chatModel) markEntryDirty(index int) { - if index < 0 || index >= len(m.entries) { - return - } - entry := &m.entries[index] - entry.version++ - if entry.version <= 0 { - entry.version = 1 - } - entry.renderKey = chatEntryRenderKey{} - entry.renderLines = nil -} - -func entryRenderKey(entry chatEntry, width int) chatEntryRenderKey { - return chatEntryRenderKey{width: width, version: entry.version} -} - -func entriesFromMessages(messages []api.Message) []chatEntry { - entries := make([]chatEntry, 0, len(messages)) - toolCalls := make(map[string]api.ToolCall) - for _, msg := range messages { - switch msg.Role { - case "user", "system": - if summary, ok := compactionSummaryContent(msg); ok { - entries = append(entries, newChatEntry(chatEntry{ - role: "compaction_summary", - content: summary, - status: "done", - })) - continue - } - entries = append(entries, newChatEntry(chatEntry{role: msg.Role, content: msg.Content})) - case "assistant": - if strings.TrimSpace(msg.Thinking) != "" { - entries = append(entries, newChatEntry(chatEntry{ - role: "thinking", - content: msg.Thinking, - label: "Thinking", - status: "done", - })) - } - for _, call := range msg.ToolCalls { - if call.ID != "" { - toolCalls[call.ID] = call - } - } - if strings.TrimSpace(msg.Content) != "" { - entries = append(entries, newChatEntry(chatEntry{role: "assistant", content: msg.Content})) - } - case "tool": - if summary, ok := compactionSummaryContent(msg); ok { - entries = append(entries, newChatEntry(chatEntry{ - role: "compaction_summary", - content: summary, - status: "done", - })) - continue - } - toolName := msg.ToolName - var args map[string]any - if call, ok := toolCalls[msg.ToolCallID]; ok { - if toolName == "" { - toolName = call.Function.Name - } - args = call.Function.Arguments.ToMap() - } - status := "done" - if isDisabledToolResult(msg.Content) { - status = "disabled" - } else if isDeniedToolResult(msg.Content) { - status = "denied" - } - entries = append(entries, newChatEntry(chatEntry{ - role: "tool", - content: msg.Content, - label: toolInvocationLabel(toolName, args), - detail: toolName, - status: status, - toolID: msg.ToolCallID, - args: args, - })) - } - } - return groupCompletedToolEntries(entries) -} - -func compactionSummaryContent(msg api.Message) (string, bool) { - return coreagent.CompactionSummaryContent(msg) -} - -func groupCompletedToolEntries(entries []chatEntry, detected ...chatEntry) []chatEntry { - grouped := make([]chatEntry, 0, len(entries)) - visibleToolIDs := visibleToolIDs(entries) - for i := 0; i < len(entries); { - if !isGroupableToolHistoryEntry(entries[i]) { - grouped = append(grouped, entries[i]) - i++ - continue - } - - start := i - var tools []chatEntry - for i < len(entries) { - if isGroupableToolHistoryEntry(entries[i]) { - tools = append(tools, flattenToolHistory([]chatEntry{entries[i]})...) - i++ - continue - } - if isInvisibleToolGroupingBoundary(entries[i]) && nextGroupableToolHistoryIndex(entries, i+1) >= 0 { - i++ - continue - } - break - } - - summaryTools := toolSummaryEntries(tools, detected, visibleToolIDs) - if len(summaryTools) <= 1 { - grouped = append(grouped, flattenToolHistory(entries[start:i])...) - continue - } - - group := chatEntry{ - role: "tool_group", - label: toolGroupSummary(summaryTools), - status: aggregateToolStatus(tools), - expanded: anyToolExpanded(tools), - startedAt: firstToolStartedAt(tools), - finishedAt: lastToolFinishedAt(tools), - tools: tools, - } - if group.status == "error" { - group.err = "one or more tool calls failed" - } - grouped = append(grouped, newChatEntry(group)) - } - return grouped -} - -func isInvisibleToolGroupingBoundary(entry chatEntry) bool { - switch entry.role { - case "assistant": - return strings.TrimSpace(entry.content) == "" && - strings.TrimSpace(entry.label) == "" && - strings.TrimSpace(entry.detail) == "" && - entry.metrics == nil - default: - return false - } -} - -func nextGroupableToolHistoryIndex(entries []chatEntry, index int) int { - for index < len(entries) && isInvisibleToolGroupingBoundary(entries[index]) { - index++ - } - if index < len(entries) && isGroupableToolHistoryEntry(entries[index]) { - return index - } - return -1 -} - -func anyToolExpanded(tools []chatEntry) bool { - for _, tool := range tools { - if tool.expanded { - return true - } - } - return false -} - -func isGroupableToolHistoryEntry(entry chatEntry) bool { - return (entry.role == "tool" && isToolResultStatus(entry.status)) || - (entry.role == "tool_group" && len(entry.tools) > 0) -} - -func visibleToolIDs(entries []chatEntry) map[string]struct{} { - ids := map[string]struct{}{} - for _, tool := range flattenToolHistory(entries) { - if tool.toolID != "" { - ids[tool.toolID] = struct{}{} - } - } - return ids -} - -func toolSummaryEntries(tools []chatEntry, detected []chatEntry, visible map[string]struct{}) []chatEntry { - if len(detected) == 0 { - return tools - } - seen := make(map[string]struct{}, len(visible)+len(tools)) - for id := range visible { - seen[id] = struct{}{} - } - summary := slices.Clone(tools) - for _, tool := range detected { - if tool.toolID != "" { - if _, ok := seen[tool.toolID]; ok { - continue - } - seen[tool.toolID] = struct{}{} - } - summary = append(summary, tool) - } - return summary -} - -func entryHasActiveTool(entry chatEntry) bool { - switch entry.role { - case "tool": - return isToolActiveStatus(entry.status) - case "tool_group": - for _, tool := range entry.tools { - if isToolActiveStatus(tool.status) { - return true - } - } - } - return false -} - -func flattenToolHistory(entries []chatEntry) []chatEntry { - var tools []chatEntry - for _, entry := range entries { - switch entry.role { - case "tool": - tools = append(tools, entry) - case "tool_group": - tools = append(tools, entry.tools...) - } - } - return tools -} - -func aggregateToolStatus(tools []chatEntry) string { - denied := false - active := false - for _, tool := range tools { - if tool.err != "" || tool.status == "error" { - return "error" - } - if tool.status == "denied" { - denied = true - } - if isToolActiveStatus(tool.status) { - active = true - } - } - if denied { - return "denied" - } - if active { - return "running" - } - return "done" -} - -func firstToolStartedAt(tools []chatEntry) time.Time { - for _, tool := range tools { - if !tool.startedAt.IsZero() { - return tool.startedAt - } - } - return time.Time{} -} - -func lastToolFinishedAt(tools []chatEntry) time.Time { - for i := len(tools) - 1; i >= 0; i-- { - if !tools[i].finishedAt.IsZero() { - return tools[i].finishedAt - } - } - return time.Time{} -} - -func indentLines(lines []string, prefix string) []string { - if len(lines) == 0 { - return nil - } - out := make([]string, len(lines)) - for i, line := range lines { - if line == "" { - out[i] = prefix - } else { - out[i] = prefix + line - } - } - return out -} - -func wrapChatText(text string, width int) []string { - if width < 20 { - width = 20 - } - var out []string - for _, rawLine := range strings.Split(text, "\n") { - line := strings.TrimRight(rawLine, "\r") - for runewidth.StringWidth(line) > width { - cut := chatDisplayWidthCut(line, width) - out = append(out, strings.TrimSpace(line[:cut])) - line = strings.TrimSpace(line[cut:]) - } - out = append(out, line) - } - if len(out) == 0 { - return []string{""} - } - return out -} - -func chatDisplayWidthCut(line string, width int) int { - hardCut := 0 - currentWidth := 0 - spaceCut := 0 - spaceWidth := 0 - for i := 0; i < len(line); { - r, size := utf8.DecodeRuneInString(line[i:]) - nextWidth := currentWidth + runewidth.RuneWidth(r) - if nextWidth > width { - break - } - currentWidth = nextWidth - hardCut = i + size - if (r == ' ' || r == '\t') && currentWidth > width/2 { - spaceCut = i - spaceWidth = currentWidth - } - i += size - } - if spaceCut > 0 && spaceWidth > 0 { - return spaceCut - } - if hardCut > 0 { - return hardCut - } - _, size := utf8.DecodeRuneInString(line) - return size -} diff --git a/cmd/tui/chat/render_test.go b/cmd/tui/chat/render_test.go deleted file mode 100644 index 0a31c7bbb7a..00000000000 --- a/cmd/tui/chat/render_test.go +++ /dev/null @@ -1,2282 +0,0 @@ -package chat - -import ( - "context" - "fmt" - "os" - "path/filepath" - "strings" - "testing" - "time" - "unicode/utf8" - - tea "github.com/charmbracelet/bubbletea" - "github.com/charmbracelet/lipgloss" - - coreagent "github.com/ollama/ollama/agent" - "github.com/ollama/ollama/api" -) - -func testIntPtr(v int) *int { - return &v -} - -func TestMetricsSummaryLinesCachedPromptTokens(t *testing.T) { - lines := metricsSummaryLines(&api.Metrics{ - PromptEvalCount: 10, - PromptEvalCachedCount: testIntPtr(4), - PromptEvalDuration: time.Second, - }) - got := strings.Join(lines, "\n") - for _, want := range []string{"prompt eval count: 10 token(s)", "prompt eval cached: 4 token(s)", "prompt eval rate: 6.00 tokens/s"} { - if !strings.Contains(got, want) { - t.Errorf("summary missing %q:\n%s", want, got) - } - } -} - -func TestChatAssistantEntryHasNoLabel(t *testing.T) { - m := chatModel{} - - prefix, _ := m.renderEntry(chatEntry{role: "assistant", content: "hello"}) - - if strings.Contains(prefix, "Ollama:") { - t.Fatalf("prefix should not include Ollama label: %q", prefix) - } - if prefix != "" { - t.Fatalf("prefix = %q, want empty", prefix) - } -} - -func TestChatViewRendersEmptyPromptHint(t *testing.T) { - m := chatModel{ - chatID: "chat-a", - width: 80, - height: 12, - } - - view := stripANSI(m.View()) - lines := strings.Split(view, "\n") - hintLine := lineIndexContaining(lines, "what changed on this branch?") - if hintLine < 0 { - t.Fatalf("empty chat view missing prompt hint: %q", view) - } - if strings.Contains(view, "Start a conversation. Use /help for commands.") { - t.Fatalf("empty chat view should use rotating prompt hint: %q", view) - } -} - -func TestChatUserEntryHasNoLabel(t *testing.T) { - m := chatModel{entries: []chatEntry{{role: "user", content: "hello"}}} - - prefix, body := m.renderEntry(m.entries[0]) - - if prefix != "" { - t.Fatalf("prefix = %q, want empty", prefix) - } - if body != "hello" { - t.Fatalf("body = %q, want hello", body) - } - - transcript := stripANSI(m.renderTranscript(80)) - if !strings.Contains(transcript, " hello") { - t.Fatalf("user transcript should render as user block: %q", transcript) - } -} - -func TestChatSystemEntryHasNoLabel(t *testing.T) { - m := chatModel{entries: []chatEntry{{role: "system", content: "Available tools"}}} - - prefix, body := m.renderEntry(m.entries[0]) - - if prefix != "" { - t.Fatalf("prefix = %q, want empty", prefix) - } - if body != "Available tools" { - t.Fatalf("body = %q, want Available tools", body) - } - if transcript := stripANSI(m.renderTranscript(80)); strings.Contains(transcript, "sys ") { - t.Fatalf("system transcript should not render sys prefix: %q", transcript) - } -} - -func TestChatMixedToolGroupUsesMixedPrefix(t *testing.T) { - m := chatModel{} - entry := newChatEntry(chatEntry{ - role: "tool_group", - label: "Tool calls (2)", - status: "error", - tools: []chatEntry{ - newChatEntry(chatEntry{role: "tool", status: "done"}), - newChatEntry(chatEntry{role: "tool", status: "error", err: "failed"}), - }, - }) - - prefix, body := m.renderEntry(entry) - - if prefix != chatToolMixedStyle.Render("•")+" " { - t.Fatalf("prefix = %q, want mixed-styled bullet", prefix) - } - body = stripANSI(body) - if strings.Contains(body, "failed") || strings.Contains(body, "succeeded") || strings.Contains(body, "done") { - t.Fatalf("body = %q, should not show status words for mixed group", body) - } -} - -func TestChatToolStatusLineDoesNotUseDisclosureGlyph(t *testing.T) { - startedAt := time.Date(2026, 6, 22, 13, 0, 0, 0, time.UTC) - entry := newChatEntry(chatEntry{ - role: "tool", - label: `Web Search("who is parth sareen")`, - status: "done", - content: "hidden result", - startedAt: startedAt, - finishedAt: startedAt.Add(812 * time.Millisecond), - }) - - line := stripANSI(toolStatusLine(entry)) - if strings.Contains(line, "▸") || strings.Contains(line, "▾") { - t.Fatalf("tool status line should not include disclosure glyph: %q", line) - } - if line != `Web Search("who is parth sareen")` { - t.Fatalf("tool status line = %q, want label only", line) - } - for _, word := range []string{"done", "in 812ms", "812ms"} { - if strings.Contains(line, word) { - t.Fatalf("tool status line should not include %q: %q", word, line) - } - } -} - -func TestChatCompletedToolStatusLineUsesResultStyle(t *testing.T) { - entry := newChatEntry(chatEntry{ - role: "tool", - detail: "bash", - label: `Bash("pwd")`, - status: "done", - }) - - if line := toolStatusLine(entry); line != `Bash("pwd")` { - t.Fatalf("tool status line = %q, want command label", line) - } - prefix, _ := (chatModel{}).renderEntry(entry) - if prefix != chatToolDoneStyle.Render("•")+" " { - t.Fatalf("tool prefix = %q, want done-styled marker", prefix) - } -} - -func TestChatToolStatusMarkerUsesStateColors(t *testing.T) { - if got, want := chatToolRunningStyle.GetForeground(), lipgloss.Color(chatAnsiYellow); got != want { - t.Fatalf("running tool foreground = %v, want %v", got, want) - } - if got, want := chatToolDoneStyle.GetForeground(), lipgloss.Color(chatAnsiGreen); got != want { - t.Fatalf("done tool foreground = %v, want %v", got, want) - } - - running := newChatEntry(chatEntry{ - role: "tool", - detail: "bash", - label: `Bash("pwd")`, - status: "running", - }) - if line := toolStatusLine(running); line != `Bash("pwd")` { - t.Fatalf("running tool status line = %q, want command label", line) - } - prefix, _ := (chatModel{}).renderEntry(running) - if prefix != chatToolRunningStyle.Render("•")+" " { - t.Fatalf("running tool prefix = %q, want running-styled marker", prefix) - } - - child := newChatEntry(chatEntry{ - role: "tool", - detail: "bash", - label: `Bash("pwd")`, - status: "done", - }) - if line := toolGroupChildStatusLine(child); line != boldToolInvocationName(`Bash("pwd")`) { - t.Fatalf("tool group child line = %q, want neutral invocation", line) - } -} - -func TestChatViewRendersInputBox(t *testing.T) { - m := chatModel{ - input: []rune("hello"), - width: 40, - height: 12, - } - - view := stripANSI(m.View()) - if !strings.Contains(view, inputBoxTopBorderLine(40)) || !strings.Contains(view, inputBoxBottomBorderLine(40)) { - t.Fatalf("prompt input should render box borders: %q", view) - } - if !strings.Contains(view, "│ hello") { - t.Fatalf("view missing prompt input row: %q", view) - } -} - -func TestChatViewRendersSentUserPromptWithoutPrefix(t *testing.T) { - m := chatModel{ - width: 40, - height: 12, - entries: []chatEntry{ - {role: "user", content: "hello"}, - }, - } - - view := stripANSI(m.View()) - if !strings.Contains(view, "hello") { - t.Fatalf("submitted user message should render: %q", view) - } - if strings.Contains(view, "> hello") { - t.Fatalf("submitted user message should not include prompt prefix: %q", view) - } - if strings.Contains(view, "│ >") { - t.Fatalf("active input should not include prompt prefix: %q", view) - } -} - -func TestChatFlowViewStartsAtInputWhenEmpty(t *testing.T) { - m := chatModel{ - input: []rune("hello"), - width: 40, - } - - lines := strings.Split(stripANSI(m.View()), "\n") - if len(lines) == 0 || !strings.Contains(lines[0], inputBoxTopBorderLine(40)) { - t.Fatalf("empty flow view should start at input box:\n%s", strings.Join(lines, "\n")) - } -} - -func TestFlowTranscriptChangedPrefixStartDetectsToolGrouping(t *testing.T) { - firstArgs := map[string]any{"command": "pwd"} - secondArgs := map[string]any{"command": "ls"} - beforeModel := chatModel{ - width: 120, - entries: []chatEntry{ - { - role: "tool", - label: `Bash("pwd")`, - detail: "bash", - status: "done", - toolID: "call-1", - args: firstArgs, - }, - }, - } - before := beforeModel.transcriptLines(120) - - afterModel := beforeModel - afterModel.entries = groupCompletedToolEntries([]chatEntry{ - beforeModel.entries[0], - { - role: "tool", - label: `Bash("ls")`, - detail: "bash", - status: "done", - toolID: "call-2", - args: secondArgs, - }, - }) - - after := afterModel.transcriptLines(120) - if start := flowTranscriptChangedPrefixStart(before, after, len(before)); start != 0 { - t.Fatalf("changed prefix start = %d, want 0; before=%q after=%q", start, before, after) - } - sequence := flowTranscriptRewriteSequence(len(before), after) - if !strings.HasPrefix(sequence, "\x1b[1A\r\x1b[J") { - t.Fatalf("rewrite sequence = %q, want cursor-up clear-below prefix", sequence) - } - if !strings.Contains(sequence, "Ran 2 commands") { - t.Fatalf("grouping should invalidate already printed tool row; before=%q after=%q", before, afterModel.transcriptLines(120)) - } -} - -func TestFlowTranscriptChangedPrefixStartIgnoresAppendedLines(t *testing.T) { - beforeModel := chatModel{ - width: 120, - entries: []chatEntry{ - {role: "assistant", content: "hello"}, - }, - } - before := beforeModel.transcriptLines(120) - afterModel := beforeModel - afterModel.entries = append(afterModel.entries, chatEntry{role: "assistant", content: "world"}) - - if start := flowTranscriptChangedPrefixStart(before, afterModel.transcriptLines(120), len(before)); start >= 0 { - t.Fatalf("appended transcript lines should not invalidate already printed prefix") - } -} - -func TestChatViewRendersCursorWithEmptyPlaceholder(t *testing.T) { - m := chatModel{ - chatID: "placeholder-chat", - width: 80, - height: 12, - } - - view := stripANSI(m.View()) - if !strings.Contains(view, "summarize this file and suggest edits") { - t.Fatalf("empty placeholder should show hint: %q", view) - } -} - -func TestChatViewRendersModelUnderInputBox(t *testing.T) { - m := chatModel{ - input: []rune("hello"), - width: 48, - height: 12, - opts: Options{ - Model: "kimi-k2.7-code:cloud", - }, - } - - lines := strings.Split(stripANSI(m.View()), "\n") - inputLine := lineIndexContaining(lines, "│ hello") - modelLine := lineIndexContaining(lines, "kimi-k2.7-code:cloud") - if inputLine < 0 || modelLine < 0 { - t.Fatalf("view missing input or model line:\n%s", strings.Join(lines, "\n")) - } - if modelLine != inputLine+2 { - t.Fatalf("model line should sit directly under input: input=%d model=%d\n%s", inputLine, modelLine, strings.Join(lines, "\n")) - } - if strings.Contains(lines[modelLine], "model ") { - t.Fatalf("model line should not include a label:\n%s", strings.Join(lines, "\n")) - } -} - -func TestChatViewExpandsInputBoxForLongPrompt(t *testing.T) { - m := chatModel{ - input: []rune(strings.Repeat("long prompt ", 8)), - width: 32, - height: 14, - } - - view := stripANSI(m.View()) - if got := inputPromptLineCount(t, view); got < 2 { - t.Fatalf("input body lines = %d, want wrapped prompt:\n%s", got, view) - } - if !strings.Contains(view, "█") { - t.Fatalf("view missing cursor: %q", view) - } -} - -func TestChatViewCapsTallInputBox(t *testing.T) { - m := chatModel{ - input: []rune(strings.Repeat("pasted text ", 80)), - width: 32, - height: 12, - } - - view := stripANSI(m.View()) - if got := inputPromptLineCount(t, view); got > maxInputBoxBodyLines { - t.Fatalf("input body lines = %d, want <= %d:\n%s", got, maxInputBoxBodyLines, view) - } - if strings.Contains(view, "... ... ") { - t.Fatalf("truncated pasted prompt should not duplicate omission markers:\n%s", view) - } -} - -func TestChatViewWrapsNotificationWhenNarrow(t *testing.T) { - m := chatModel{ - input: []rune("hello"), - width: 28, - height: 14, - status: "cache will break by turning system prompt off", - approvalState: testApprovalState(true, nil), - opts: Options{ - ContextWindowTokens: 262144, - }, - contextTokens: 12345, - contextEstimate: true, - } - - view := stripANSI(m.View()) - for _, want := range []string{ - "cache will break by", - "system prompt off", - } { - if !strings.Contains(view, want) { - t.Fatalf("wrapped notification missing %q:\n%s", want, view) - } - } - for _, hidden := range []string{"enter", "send", "/model", "full", "access"} { - if strings.Contains(view, hidden) { - t.Fatalf("view should not render footer chrome %q:\n%s", hidden, view) - } - } - if strings.Contains(view, "ctx") { - t.Fatalf("view should hide distant context pressure:\n%s", view) - } - if strings.Contains(view, "ctrl+g") { - t.Fatalf("view should not include ctrl+g hint:\n%s", view) - } - for _, line := range strings.Split(view, "\n") { - if len([]rune(line)) > 28 { - t.Fatalf("line width = %d, want <= 28: %q\n%s", len([]rune(line)), line, view) - } - } -} - -func TestPromptTokenTextDoesNotUseCompactionFallback(t *testing.T) { - m := chatModel{} - if got := m.promptTokenText(702); got != "702 tokens" { - t.Fatalf("promptTokenText without model context = %q, want bare token count", got) - } - - m.opts.ContextWindowTokens = 262144 - if got := m.promptTokenText(702); got != "702 / 256k tokens" { - t.Fatalf("promptTokenText with model context = %q", got) - } - - m.opts.Options = map[string]any{"num_ctx": 131072} - if got := m.promptTokenText(702); got != "702 / 128k tokens" { - t.Fatalf("promptTokenText with num_ctx = %q", got) - } - - m.opts.Options = map[string]any{"num_ctx": 1000000} - if got := m.promptTokenText(702); got != "702 / 1M tokens" { - t.Fatalf("promptTokenText with large num_ctx = %q", got) - } - - m.opts.Options = map[string]any{"num_ctx": 99999} - if got := m.promptTokenText(702); got != "702 / 99999 tokens" { - t.Fatalf("promptTokenText below compact threshold = %q", got) - } -} - -func TestPreloadDoneUpdatesContextWindowTokens(t *testing.T) { - compactor := &coreagent.SimpleCompactor{} - m := chatModel{ - preloadingModel: "ornith", - opts: Options{ - ContextWindowTokens: 32768, - Compactor: compactor, - }, - } - - next, _ := m.Update(chatModelPreloadDoneMsg{model: "ornith", contextWindowTokens: 262144}) - got := next.(chatModel) - if got.opts.ContextWindowTokens != 262144 { - t.Fatalf("context window = %d, want 262144", got.opts.ContextWindowTokens) - } - if compactor.Options.ContextWindowTokens != 262144 { - t.Fatalf("compactor context window = %d, want 262144", compactor.Options.ContextWindowTokens) - } -} - -func TestChatViewRendersNotificationAboveInput(t *testing.T) { - m := chatModel{ - input: []rune("hello"), - width: 40, - height: 12, - status: "copied latest output", - } - - view := stripANSI(m.View()) - lines := strings.Split(view, "\n") - borderLine := lineIndexContaining(lines, inputBoxTopBorderLine(40)) - if borderLine < 0 { - t.Fatalf("view missing input box:\n%s", view) - } - if borderLine < 1 || !strings.Contains(lines[borderLine-1], "copied latest output") { - t.Fatalf("notification should sit directly above input box:\n%s", view) - } -} - -func TestChatNotificationsUseSecondaryStyle(t *testing.T) { - if !chatNotificationStyle.GetFaint() { - t.Fatal("notification style should use secondary/faint styling") - } -} - -func TestChatSubmittedPromptUsesThemeSecondaryGrey(t *testing.T) { - if got, want := chatUserBlockStyle.GetForeground(), lipgloss.TerminalColor(lipgloss.AdaptiveColor{Light: "#777777", Dark: "#8a8a8a"}); got != want { - t.Fatalf("submitted prompt foreground = %v, want %v", got, want) - } - if chatUserBlockStyle.GetFaint() { - t.Fatal("submitted prompt should use secondary grey, not faint styling") - } -} - -func TestChatToolOutputUsesDistinctSecondaryGrey(t *testing.T) { - if got, want := chatToolOutputStyle.GetForeground(), lipgloss.TerminalColor(lipgloss.AdaptiveColor{Light: "#666666", Dark: "#a0a0a0"}); got != want { - t.Fatalf("tool output foreground = %v, want %v", got, want) - } -} - -func TestChatInlineCodeDoesNotLookSelected(t *testing.T) { - if chatInlineCodeStyle.GetReverse() { - t.Fatal("inline code should not use reverse-video styling") - } - if !chatInlineCodeStyle.GetBold() { - t.Fatal("inline code should still have lightweight emphasis") - } -} - -func inputPromptLineCount(t *testing.T, view string) int { - t.Helper() - count := 0 - inInputBox := false - for _, line := range strings.Split(view, "\n") { - trimmed := strings.TrimSpace(line) - if strings.HasPrefix(trimmed, "╭") { - inInputBox = true - continue - } - if strings.HasPrefix(trimmed, "╰") { - inInputBox = false - continue - } - if inInputBox && strings.Contains(trimmed, "│") { - count++ - } - } - if count == 0 { - t.Fatalf("input prompt lines not found:\n%s", view) - } - return count -} - -func TestChatViewKeepsInputBoxWhileRunning(t *testing.T) { - m := chatModel{ - input: []rune("next"), - running: true, - thinking: true, - thinkingTokens: 42, - width: 40, - height: 12, - } - - view := stripANSI(m.View()) - if !strings.Contains(view, "│ next") { - t.Fatalf("running view should keep input row: %q", view) - } - if strings.Contains(view, "↑/↓ scroll") || strings.Contains(view, "/new chat") || strings.Contains(view, "/clear reset") { - t.Fatalf("footer should not include scroll/new/clear hints: %q", view) - } - lines := strings.Split(view, "\n") - borderLine := lineIndexContaining(lines, inputBoxTopBorderLine(40)) - if borderLine < 0 { - t.Fatalf("view missing input box: %q", view) - } - if borderLine < 1 || !strings.Contains(lines[borderLine-1], "Thinking ↓ 42 tokens") { - t.Fatalf("thinking line should sit directly above input box:\n%s", view) - } - if strings.Contains(lines[borderLine-1], "...") { - t.Fatalf("thinking line should not show spinner dots:\n%s", view) - } -} - -func TestChatViewShowsActiveThinkingStatusOnce(t *testing.T) { - m := chatModel{ - running: true, - thinking: true, - thinkingTokens: 42, - width: 60, - height: 16, - entries: []chatEntry{ - {role: "thinking", label: "Thinking ↓ 42 tokens", status: "running", content: "streamed trace", expanded: true}, - }, - } - - view := stripANSI(m.View()) - if count := strings.Count(view, "Thinking ↓ 42 tokens"); count != 1 { - t.Fatalf("active thinking status appears %d times, want once:\n%s", count, view) - } -} - -func TestChatToolFinishedUpdatesLiveWorkingDirOnly(t *testing.T) { - root := t.TempDir() - subdir := filepath.Join(root, "sub") - if err := os.Mkdir(subdir, 0o755); err != nil { - t.Fatal(err) - } - m := chatModel{ - width: 140, - height: 12, - opts: Options{ - RootDir: root, - WorkingDir: root, - }, - workingDir: root, - } - - m.applyAgentEvent(coreagent.Event{ - Type: coreagent.EventToolFinished, - ToolCallID: "call-1", - ToolName: "bash", - WorkingDir: subdir, - }) - - if m.workingDir != subdir { - t.Fatalf("workingDir = %q, want %q", m.workingDir, subdir) - } - if m.opts.WorkingDir != root { - t.Fatalf("opts.WorkingDir mutated to %q, want %q", m.opts.WorkingDir, root) - } -} - -func TestChatBoundedFramePinsInputAfterResize(t *testing.T) { - m := chatModel{ - input: []rune("next"), - width: 60, - height: 14, - } - for i := range 20 { - m.entries = append(m.entries, chatEntry{role: "assistant", content: fmt.Sprintf("line %02d", i)}) - } - inputLine := renderedInputLine(m.View()) - if inputLine < m.height-4 { - t.Fatalf("bounded frame should pin input near bottom, input line=%d height=%d\n%s", inputLine, m.height, stripANSI(m.View())) - } -} - -func TestChatViewSeparatesActionStatusFromTranscript(t *testing.T) { - m := chatModel{ - input: []rune("next"), - width: 72, - height: 14, - running: true, - thinking: true, - entries: []chatEntry{ - {role: "assistant", content: "Working through it."}, - }, - } - - lines := strings.Split(stripANSI(m.View()), "\n") - assistantLine := lineIndexContaining(lines, "Working through it.") - activityLine := lineIndexContaining(lines, "Thinking") - if assistantLine < 0 || activityLine < 0 { - t.Fatalf("view missing assistant/activity lines:\n%s", strings.Join(lines, "\n")) - } - if gap := activityLine - assistantLine - 1; gap < 1 { - t.Fatalf("gap between transcript and action status = %d, want at least 1:\n%s", gap, strings.Join(lines, "\n")) - } -} - -func TestChatViewDoesNotReserveIdleActionSpacerAfterResponse(t *testing.T) { - m := chatModel{ - input: []rune("next"), - width: 72, - height: 12, - entries: []chatEntry{ - {role: "user", content: "hi"}, - {role: "assistant", content: "Hello."}, - }, - } - - lines := strings.Split(stripANSI(m.View()), "\n") - assistantLine := lineIndexContaining(lines, "Hello.") - inputLine := lineIndexContaining(lines, "│ next") - if assistantLine < 0 || inputLine < 0 { - t.Fatalf("view missing assistant/input lines:\n%s", strings.Join(lines, "\n")) - } - if gap := inputLine - assistantLine - 1; gap < 2 { - t.Fatalf("gap between finished response and input body = %d, want at least 2:\n%s", gap, strings.Join(lines, "\n")) - } -} - -func TestChatViewHidesEmptyHintWhileTyping(t *testing.T) { - m := chatModel{ - input: []rune("next"), - width: 72, - height: 12, - } - - lines := strings.Split(stripANSI(m.View()), "\n") - hintLine := lineIndexContaining(lines, "Try:") - inputLine := lineIndexContaining(lines, "│ next") - if hintLine >= 0 { - t.Fatalf("view should not show empty hint while typing:\n%s", strings.Join(lines, "\n")) - } - if inputLine < 0 { - t.Fatalf("view missing input line:\n%s", strings.Join(lines, "\n")) - } -} - -func renderedInputLine(view string) int { - for i, line := range strings.Split(stripANSI(view), "\n") { - if strings.Contains(line, "│ next") { - return i - } - } - return -1 -} - -func lineIndexContaining(lines []string, needle string) int { - for i, line := range lines { - if strings.Contains(line, needle) { - return i - } - } - return -1 -} - -func TestChatScrollsTranscript(t *testing.T) { - m := chatModel{ - width: 80, - height: 10, - } - for range 12 { - m.entries = append(m.entries, chatEntry{role: "user", content: "line"}) - } - - if m.maxScroll() == 0 { - t.Fatal("test setup should produce scrollable transcript") - } - - updated, _ := m.Update(tea.KeyMsg{Type: tea.KeyPgUp}) - m = updated.(chatModel) - if m.scroll == 0 { - t.Fatal("page up should scroll transcript") - } - - updated, _ = m.Update(tea.KeyMsg{Type: tea.KeyPgDown}) - m = updated.(chatModel) - if m.scroll != 0 { - t.Fatalf("scroll = %d, want 0", m.scroll) - } - - updated, _ = m.Update(tea.KeyMsg{Type: tea.KeyCtrlHome}) - m = updated.(chatModel) - if m.scroll != m.maxScroll() { - t.Fatalf("scroll = %d, want max %d", m.scroll, m.maxScroll()) - } - - updated, _ = m.Update(tea.KeyMsg{Type: tea.KeyCtrlEnd}) - m = updated.(chatModel) - if m.scroll != 0 { - t.Fatalf("scroll = %d, want 0", m.scroll) - } -} - -func TestChatStreamingAssistantOutputHoldsLiveMarkdown(t *testing.T) { - m := chatModel{ - width: 80, - height: 12, - running: true, - events: make(chan tea.Msg), - } - - updated, _ := m.Update(chatAgentMsg{event: coreagent.Event{Type: coreagent.EventMessageDelta, Content: "generated line 00"}}) - m = updated.(chatModel) - if !strings.Contains(stripANSI(m.View()), "generated line 00") { - t.Fatalf("live output should be visible in flow view:\n%s", stripANSI(m.View())) - } - - updated, _ = m.Update(chatAgentMsg{event: coreagent.Event{Type: coreagent.EventMessageDelta, Content: "\n\ngenerated line 01"}}) - m = updated.(chatModel) - view := stripANSI(m.View()) - if !strings.Contains(view, "generated line 00") { - t.Fatalf("live assistant output should remain visible while streaming:\n%s", view) - } - if !strings.Contains(view, "generated line 01") { - t.Fatalf("latest generated line should remain visible:\n%s", view) - } -} - -func TestChatStreamingRendersBoldBareURLAfterCompletion(t *testing.T) { - const response = "Draft PR opened: **https://github.com/ollama/ollama/pull/17203**" - m := chatModel{width: 80, height: 12, running: true, events: make(chan tea.Msg)} - - updated, _ := m.Update(chatAgentMsg{event: coreagent.Event{Type: coreagent.EventMessageDelta, Content: "Draft PR opened: **https://github.com/ollama/"}}) - m = updated.(chatModel) - if got := stripANSI(m.renderTranscript(80)); !strings.Contains(got, "**https://github.com/ollama/") { - t.Fatalf("incomplete Markdown should remain visible while streaming: %q", got) - } - - updated, _ = m.Update(chatAgentMsg{event: coreagent.Event{Type: coreagent.EventMessageDelta, Content: "ollama/pull/17203**"}}) - m = updated.(chatModel) - if got := m.entries[0].content; got != response { - t.Fatalf("streamed content = %q, want %q", got, response) - } - rendered := m.renderTranscript(80) - plain := stripANSI(rendered) - if strings.Contains(plain, "**") { - t.Fatalf("rendered response should not contain Markdown delimiters: %q", plain) - } - if !strings.Contains(plain, "Draft PR opened: https://github.com/ollama/ollama/pull/17203") { - t.Fatalf("rendered response missing URL: %q", plain) - } - if !strings.Contains(rendered, chatStrongStyle.Render("https://github.com/ollama/ollama/pull/17203")) { - t.Fatalf("URL should use the bold terminal style: %q", rendered) - } -} - -func TestRenderMarkdownInlineWrapsStrongTextWithoutDelimiters(t *testing.T) { - rendered := renderMarkdownForView("**alpha beta gamma delta epsilon**", 20) - plain := stripANSI(rendered) - if strings.Contains(plain, "**") { - t.Fatalf("wrapped strong text should not contain Markdown delimiters: %q", plain) - } - for _, line := range strings.Split(rendered, "\n") { - if got := lipgloss.Width(line); got > 20 { - t.Fatalf("rendered line width = %d, want <= 20: %q", got, line) - } - } -} - -func TestRenderMarkdownPreservesBareURLUnderscores(t *testing.T) { - const url = "https://example.com/a__b__" - if got := stripANSI(renderMarkdownForView(url, 80)); got != url { - t.Fatalf("bare URL = %q, want %q", got, url) - } -} - -func TestRenderMarkdownStrongAfterPunctuation(t *testing.T) { - for _, test := range []struct { - name string - input string - want string - emphasis string - }{ - { - name: "colon", - input: "Status: **ready**", - want: "Status: ready", - emphasis: "ready", - }, - { - name: "dash", - input: "Note-**important**", - want: "Note-important", - emphasis: "important", - }, - { - name: "closing parenthesis", - input: "Result) **complete**", - want: "Result) complete", - emphasis: "complete", - }, - { - name: "identifier", - input: "value__with_delimiters__", - want: "value__with_delimiters__", - }, - { - name: "URL", - input: "https://example.com/a__b__", - want: "https://example.com/a__b__", - }, - { - name: "URL punctuation", - input: "https://example.com/a-**b**", - want: "https://example.com/a-**b**", - }, - } { - t.Run(test.name, func(t *testing.T) { - rendered := renderMarkdownForView(test.input, 80) - if got := stripANSI(rendered); got != test.want { - t.Fatalf("rendered = %q, want %q", got, test.want) - } - if test.emphasis != "" && !strings.Contains(rendered, chatStrongStyle.Render(test.emphasis)) { - t.Fatalf("rendered output should emphasize %q: %q", test.emphasis, rendered) - } - }) - } -} - -func TestChatMouseWheelScrollsTranscriptWhileRunning(t *testing.T) { - m := chatModel{ - width: 80, - height: 10, - running: true, - input: []rune("current draft"), - promptHistory: []string{"previous one", "previous two"}, - } - for range 12 { - m.entries = append(m.entries, chatEntry{role: "user", content: "line"}) - } - if m.maxScroll() == 0 { - t.Fatal("test setup should produce scrollable transcript") - } - - updated, _ := m.Update(tea.MouseMsg{Type: tea.MouseWheelUp}) - m = updated.(chatModel) - if m.scroll == 0 { - t.Fatal("mouse wheel up should scroll transcript while running") - } - if got := string(m.input); got != "current draft" { - t.Fatalf("mouse wheel should not navigate prompt history, input = %q", got) - } - - updated, _ = m.Update(tea.MouseMsg{Type: tea.MouseWheelDown}) - m = updated.(chatModel) - if m.scroll != 0 { - t.Fatalf("mouse wheel down should return to bottom, got scroll %d", m.scroll) - } - if got := string(m.input); got != "current draft" { - t.Fatalf("mouse wheel should leave draft alone, input = %q", got) - } -} - -func TestChatWindowsMouseWheelScrollsTranscript(t *testing.T) { - oldGOOS := chatRuntimeGOOS - chatRuntimeGOOS = "windows" - defer func() { - chatRuntimeGOOS = oldGOOS - }() - - m := chatModel{ - width: 80, - height: 8, - } - for i := range 20 { - m.entries = append(m.entries, chatEntry{role: "user", content: fmt.Sprintf("line-%02d", i)}) - } - - updated, _ := m.Update(tea.MouseMsg{Type: tea.MouseWheelUp}) - m = updated.(chatModel) - if m.scroll == 0 { - t.Fatal("windows mouse wheel up should scroll transcript") - } - - updated, _ = m.Update(tea.MouseMsg{Type: tea.MouseWheelDown}) - m = updated.(chatModel) - if m.scroll != 0 { - t.Fatalf("windows mouse wheel down should return to bottom, got scroll %d", m.scroll) - } -} - -func TestChatMouseDragSelectsTranscriptWithoutAutoCopy(t *testing.T) { - oldGOOS := chatRuntimeGOOS - chatRuntimeGOOS = "darwin" - defer func() { - chatRuntimeGOOS = oldGOOS - }() - - m := chatModel{ - width: 80, - height: 10, - entries: []chatEntry{ - {role: "user", content: "alpha beta"}, - }, - } - top, _ := m.transcriptLayout() - - updated, _ := m.Update(tea.MouseMsg{Type: tea.MouseLeft, Button: tea.MouseButtonLeft, Action: tea.MouseActionPress, X: 2, Y: top}) - m = updated.(chatModel) - updated, _ = m.Update(tea.MouseMsg{Type: tea.MouseLeft, Button: tea.MouseButtonLeft, Action: tea.MouseActionMotion, X: 7, Y: top}) - m = updated.(chatModel) - if got := m.selectedTranscriptText(80); got != "alpha" { - t.Fatalf("selected text = %q, want alpha", got) - } - if !m.selection.active { - t.Fatal("selection should stay active during drag") - } - if !m.selection.dragging { - t.Fatal("selection should track drag before release") - } - - updated, cmd := m.Update(tea.MouseMsg{Type: tea.MouseRelease, Action: tea.MouseActionRelease, X: 7, Y: top}) - m = updated.(chatModel) - if cmd != nil { - t.Fatal("mouse release should not auto-copy selected text") - } - if !m.selection.active { - t.Fatal("selection should stay visible on release") - } - if m.selection.dragging { - t.Fatal("selection should stop tracking drag on release") - } - if got := m.selectedTranscriptText(80); got != "alpha" { - t.Fatalf("selected text after release = %q, want alpha", got) - } -} - -func TestChatWindowsMouseDragCopiesAndClearsTranscriptSelection(t *testing.T) { - oldGOOS := chatRuntimeGOOS - chatRuntimeGOOS = "windows" - defer func() { - chatRuntimeGOOS = oldGOOS - }() - - var copied string - oldClipboard := writeClipboard - writeClipboard = func(_ context.Context, text string) error { - copied = text - return nil - } - defer func() { - writeClipboard = oldClipboard - }() - - m := chatModel{ - width: 80, - height: 10, - entries: []chatEntry{ - {role: "user", content: "alpha beta"}, - }, - } - top, _ := m.transcriptLayout() - - updated, _ := m.Update(tea.MouseMsg{Type: tea.MouseLeft, Button: tea.MouseButtonLeft, Action: tea.MouseActionPress, X: 2, Y: top}) - m = updated.(chatModel) - updated, _ = m.Update(tea.MouseMsg{Type: tea.MouseLeft, Button: tea.MouseButtonLeft, Action: tea.MouseActionMotion, X: 7, Y: top}) - m = updated.(chatModel) - if got := m.selectedTranscriptText(80); got != "alpha" { - t.Fatalf("selected text = %q, want alpha", got) - } - - updated, cmd := m.Update(tea.MouseMsg{Type: tea.MouseRelease, Action: tea.MouseActionRelease, X: 7, Y: top}) - m = updated.(chatModel) - if cmd == nil { - t.Fatal("mouse release should copy selected text on windows") - } - if msg := cmd(); msg != nil { - t.Fatalf("clipboard command message = %#v, want nil", msg) - } - if copied != "alpha" { - t.Fatalf("copied text = %q, want alpha", copied) - } - if m.selection.active { - t.Fatal("windows selection should clear after copy") - } - if m.status != "copied" { - t.Fatalf("status = %q, want copied", m.status) - } -} - -func TestChatMouseDragSelectionUsesDisplayColumns(t *testing.T) { - m := chatModel{ - width: 80, - height: 10, - entries: []chatEntry{ - {role: "user", content: "a界b"}, - }, - } - top, _ := m.transcriptLayout() - contentX := 2 - - updated, _ := m.Update(tea.MouseMsg{Type: tea.MouseLeft, Button: tea.MouseButtonLeft, Action: tea.MouseActionPress, X: contentX, Y: top}) - m = updated.(chatModel) - updated, _ = m.Update(tea.MouseMsg{Type: tea.MouseLeft, Button: tea.MouseButtonLeft, Action: tea.MouseActionMotion, X: contentX + 3, Y: top}) - m = updated.(chatModel) - - if got := m.selectedTranscriptText(80); got != "a界" { - t.Fatalf("selected text = %q, want a界", got) - } -} - -func TestChatBoundedViewDoesNotRenderScrollHeader(t *testing.T) { - m := chatModel{ - width: 80, - height: 6, - } - for i := range 12 { - m.entries = append(m.entries, chatEntry{role: "user", content: fmt.Sprintf("line-%02d", i)}) - } - m.scroll = m.maxScroll() - top, _ := m.transcriptLayout() - lines := strings.Split(stripANSI(m.View()), "\n") - if top != 0 { - t.Fatalf("transcript top = %d, want no status header", top) - } - if strings.Contains(lines[0], "more") { - t.Fatalf("view should not render scroll status header: %q", lines[0]) - } - if strings.TrimSpace(lines[top]) == "" { - t.Fatalf("transcript should start at layout top %d, line=%q view=%q", top, lines[top], strings.Join(lines, "\n")) - } -} - -func TestChatMouseDragSelectionUsesScrolledTranscriptCoordinates(t *testing.T) { - oldGOOS := chatRuntimeGOOS - chatRuntimeGOOS = "darwin" - defer func() { - chatRuntimeGOOS = oldGOOS - }() - - m := chatModel{ - width: 80, - height: 8, - } - for i := range 10 { - m.entries = append(m.entries, chatEntry{role: "user", content: fmt.Sprintf("line-%02d", i)}) - } - m.scroll = m.maxScroll() - top, _ := m.transcriptLayout() - - updated, _ := m.Update(tea.MouseMsg{Type: tea.MouseLeft, Button: tea.MouseButtonLeft, Action: tea.MouseActionPress, X: 2, Y: top}) - m = updated.(chatModel) - updated, _ = m.Update(tea.MouseMsg{Type: tea.MouseMotion, Button: tea.MouseButtonLeft, Action: tea.MouseActionMotion, X: 9, Y: top}) - m = updated.(chatModel) - if got := m.selectedTranscriptText(80); got != "line-00" { - t.Fatalf("selected text = %q, want line-00", got) - } - updated, cmd := m.Update(tea.MouseMsg{Type: tea.MouseRelease, Action: tea.MouseActionRelease, X: 9, Y: top}) - m = updated.(chatModel) - if cmd != nil { - t.Fatal("mouse release should not auto-copy selected text") - } - if !m.selection.active { - t.Fatal("selection should stay visible on release") - } - if m.selection.dragging { - t.Fatal("selection should stop tracking drag on release") - } -} - -func TestChatArrowKeysNavigatePromptHistoryWhenTranscriptScrollable(t *testing.T) { - m := chatModel{ - input: []rune("current draft"), - promptHistory: []string{"old prompt"}, - width: 80, - height: 10, - running: true, - } - for range 12 { - m.entries = append(m.entries, chatEntry{role: "user", content: "line"}) - } - if m.maxScroll() == 0 { - t.Fatal("test setup should produce scrollable transcript") - } - - updated, _ := m.Update(tea.KeyMsg{Type: tea.KeyUp}) - m = updated.(chatModel) - if m.scroll != 0 { - t.Fatalf("key up should not scroll transcript, got %d", m.scroll) - } - if got := string(m.input); got != "old prompt" { - t.Fatalf("key up should recall prompt history, got %q", got) - } - - updated, _ = m.Update(tea.KeyMsg{Type: tea.KeyDown}) - m = updated.(chatModel) - if m.scroll != 0 { - t.Fatalf("key down should not scroll transcript, got %d", m.scroll) - } - if got := string(m.input); got != "current draft" { - t.Fatalf("key down should restore draft, got %q", got) - } -} - -func TestEntriesFromMessagesSkipsToolCallOnlyAssistant(t *testing.T) { - args := api.NewToolCallFunctionArguments() - args.Set("command", "pwd") - messages := []api.Message{ - {Role: "user", Content: "pwd"}, - { - Role: "assistant", - ToolCalls: []api.ToolCall{{ - ID: "call-1", - Function: api.ToolCallFunction{ - Name: "bash", - Arguments: args, - }, - }}, - }, - {Role: "tool", ToolName: "bash", ToolCallID: "call-1", Content: "/tmp/project\n"}, - {Role: "assistant", Content: "The current directory is /tmp/project."}, - } - - entries := entriesFromMessages(messages) - if len(entries) != 3 { - t.Fatalf("entries = %d, want user/tool/assistant: %#v", len(entries), entries) - } - if entries[1].role != "tool" || entries[1].label != "Bash(\"pwd\")" { - t.Fatalf("tool entry = %#v", entries[1]) - } - if line := stripANSI(toolStatusLine(entries[1])); line != `Bash("pwd")` { - t.Fatalf("tool status line = %q, want command label", line) - } - - transcript := stripANSI((chatModel{entries: entries}).renderTranscript(120)) - if strings.Contains(transcript, "• \n") || strings.Contains(transcript, "•\n") { - t.Fatalf("transcript has blank assistant bullet: %q", transcript) - } -} - -func TestToolActionPhraseLoadsSkills(t *testing.T) { - if got := toolActionPhrase("skill", 1); got != "Loaded a skill" { - t.Fatalf("single skill action = %q", got) - } - if got := toolActionPhrase("skill", 2); got != "Loaded 2 skills" { - t.Fatalf("multiple skill action = %q", got) - } -} - -func TestEntriesFromMessagesRendersDeniedCommandAsDenied(t *testing.T) { - args := api.NewToolCallFunctionArguments() - args.Set("command", "pwd") - messages := []api.Message{ - { - Role: "assistant", - ToolCalls: []api.ToolCall{{ - ID: "call-1", - Function: api.ToolCallFunction{ - Name: "bash", - Arguments: args, - }, - }}, - }, - {Role: "tool", ToolName: "bash", ToolCallID: "call-1", Content: "Tool execution denied."}, - } - - entries := entriesFromMessages(messages) - if len(entries) != 1 { - t.Fatalf("entries = %d, want one tool entry: %#v", len(entries), entries) - } - if entries[0].status != "denied" { - t.Fatalf("tool status = %q, want denied: %#v", entries[0].status, entries[0]) - } - if line := stripANSI(toolStatusLine(entries[0])); line != `Bash("pwd") denied` { - t.Fatalf("tool status line = %q, want denied command label", line) - } -} - -func TestEntriesFromMessagesRendersDisabledToolAsSkipped(t *testing.T) { - args := api.NewToolCallFunctionArguments() - args.Set("command", "pwd") - messages := []api.Message{ - { - Role: "assistant", - ToolCalls: []api.ToolCall{{ - ID: "call-1", - Function: api.ToolCallFunction{ - Name: "bash", - Arguments: args, - }, - }}, - }, - {Role: "tool", ToolName: "bash", ToolCallID: "call-1", Content: "Tool execution disabled."}, - } - - entries := entriesFromMessages(messages) - if len(entries) != 1 { - t.Fatalf("entries = %d, want one tool entry: %#v", len(entries), entries) - } - if entries[0].status != "disabled" { - t.Fatalf("tool status = %q, want disabled: %#v", entries[0].status, entries[0]) - } - if action := toolActionForEntry(entries[0]); action != "disabled_command" { - t.Fatalf("tool action = %q, want disabled_command", action) - } - if phrase := toolActionPhrase(toolActionForEntry(entries[0]), 1); phrase != "Skipped a command" { - t.Fatalf("tool action phrase = %q, want skipped command", phrase) - } -} - -func TestEntriesFromMessagesRendersCompactionSummaryCollapsed(t *testing.T) { - entries := entriesFromMessages([]api.Message{ - {Role: "assistant", ToolCalls: []api.ToolCall{{ - ID: coreagent.CompactionToolCallID, - Function: api.ToolCallFunction{ - Name: coreagent.CompactionToolName, - }, - }}}, - {Role: "tool", ToolName: coreagent.CompactionToolName, ToolCallID: coreagent.CompactionToolCallID, Content: coreagent.CompactionSummaryMessagePrefix + "- old work\n- decisions"}, - {Role: "user", Content: "recent request"}, - }) - if len(entries) != 2 { - t.Fatalf("entries = %d, want summary plus user: %#v", len(entries), entries) - } - if entries[0].role != "compaction_summary" || entries[0].content != "- old work\n- decisions" { - t.Fatalf("summary entry = %#v", entries[0]) - } - - m := chatModel{entries: entries} - transcript := stripANSI(m.renderTranscript(100)) - if !strings.Contains(transcript, "Compacted summary") { - t.Fatalf("collapsed summary row missing: %q", transcript) - } - if strings.Contains(transcript, "Compacted summary done") { - t.Fatalf("summary row should not include done word: %q", transcript) - } - if strings.Contains(transcript, "old work") { - t.Fatalf("summary body should be collapsed: %q", transcript) - } - - updated, _ := m.Update(tea.KeyMsg{Type: tea.KeyCtrlO}) - m = updated.(chatModel) - view := stripANSI(m.renderTranscript(100)) - if !strings.Contains(view, "old work") || !strings.Contains(view, "decisions") { - t.Fatalf("inline summary body missing: %q", view) - } -} - -func TestEntriesFromMessagesHidesAutomaticCompactionInstruction(t *testing.T) { - entries := entriesFromMessages([]api.Message{ - { - Role: "tool", - ToolName: coreagent.CompactionToolName, - ToolCallID: coreagent.CompactionToolCallID, - Content: coreagent.CompactionSummaryMessagePrefix + "old work summary\n\n" + coreagent.CompactionContinueInstruction, - }, - }) - if len(entries) != 1 || entries[0].role != "compaction_summary" || entries[0].content != "old work summary" { - t.Fatalf("entries = %#v", entries) - } -} - -func TestEntriesFromMessagesRecognizesLegacySystemCompactionSummary(t *testing.T) { - entries := entriesFromMessages([]api.Message{ - {Role: "system", Content: "Conversation summary:\nlegacy summary"}, - }) - if len(entries) != 1 || entries[0].role != "compaction_summary" || entries[0].content != "legacy summary" { - t.Fatalf("entries = %#v", entries) - } -} - -func TestEntriesFromMessagesGroupsMultiToolHistoryInMiddle(t *testing.T) { - readArgs := api.NewToolCallFunctionArguments() - readArgs.Set("path", "feedback") - listArgs := api.NewToolCallFunctionArguments() - listArgs.Set("path", ".") - messages := []api.Message{ - {Role: "user", Content: "read feedback file"}, - { - Role: "assistant", - ToolCalls: []api.ToolCall{ - { - ID: "call-read", - Function: api.ToolCallFunction{ - Name: "read", - Arguments: readArgs, - }, - }, - { - ID: "call-list", - Function: api.ToolCallFunction{ - Name: "list", - Arguments: listArgs, - }, - }, - }, - }, - {Role: "tool", ToolName: "read", ToolCallID: "call-read", Content: "Error: no such file"}, - {Role: "tool", ToolName: "list", ToolCallID: "call-list", Content: "feedback.md\n"}, - {Role: "assistant", Content: "There is a feedback.md file."}, - } - - entries := entriesFromMessages(messages) - if len(entries) != 3 { - t.Fatalf("entries = %d, want user/tool-group/assistant: %#v", len(entries), entries) - } - if entries[1].role != "tool_group" || len(entries[1].tools) != 2 { - t.Fatalf("middle entry should be grouped tool history: %#v", entries[1]) - } - if entries[1].tools[0].label != "Read(\"feedback\")" || entries[1].tools[1].label != "List(\".\")" { - t.Fatalf("group labels = %#v", entries[1].tools) - } -} - -func TestChatToolCallDetectedDoesNotRenderQueuedRows(t *testing.T) { - args := api.NewToolCallFunctionArguments() - args.Set("command", "pwd") - m := chatModel{} - - m.applyAgentEvent(coreagent.Event{ - Type: coreagent.EventToolCallDetected, - ToolCalls: []api.ToolCall{{ - ID: "call-1", - Function: api.ToolCallFunction{ - Name: "bash", - Arguments: args, - }, - }}, - }) - - if len(m.entries) != 0 { - t.Fatalf("queued tool call should not create history entries: %#v", m.entries) - } - - m.applyAgentEvent(coreagent.Event{ - Type: coreagent.EventToolStarted, - ToolCallID: "call-1", - ToolName: "bash", - Args: args.ToMap(), - }) - if len(m.entries) != 1 { - t.Fatalf("started tool should create one visible entry, got %d", len(m.entries)) - } - if got := m.activityLabel(); got != "" { - t.Fatalf("activityLabel = %q, want no transient tool label", got) - } - if line := strings.TrimSpace(stripANSI(m.activityLine())); line != "" { - t.Fatalf("activityLine = %q, want no transient tool action line", line) - } -} - -func TestChatToolOutputIsHiddenUntilExpanded(t *testing.T) { - fullOutput := strings.Repeat("x", maxCtrlOToolOutputRunes+25) + "tail-marker" - m := chatModel{} - m.applyAgentEvent(coreagent.Event{ - Type: coreagent.EventToolFinished, - ToolName: "bash", - Content: fullOutput, - }) - - if len(m.entries) != 1 { - t.Fatalf("entries = %d, want 1", len(m.entries)) - } - if m.entries[0].content != fullOutput { - t.Fatal("tool entry should keep full content before rendering") - } - - transcript := m.renderTranscript(100) - body := stripANSI(transcript) - if strings.Contains(body, "tail-marker") || strings.Contains(body, strings.Repeat("x", 20)) { - t.Fatalf("collapsed tool output should be hidden: %q", body) - } - - updated, _ := m.Update(tea.KeyMsg{Type: tea.KeyCtrlO}) - m = updated.(chatModel) - - if !m.toolOutputMode || !m.toolOutputOpen || !m.entries[0].expanded { - t.Fatalf("ctrl+o should expand tool output inline: %#v", m.entries[0]) - } - body = stripANSI(m.renderTranscript(100)) - if strings.Contains(body, "tail-marker") { - t.Fatalf("expanded transcript should cap tool output: %q", body) - } - if !strings.Contains(body, "...") { - t.Fatalf("expanded transcript should show capped output ellipsis: %q", body) - } - if got := strings.Count(body, "x"); got != maxCtrlOToolOutputRunes-3 { - t.Fatalf("expanded transcript x count = %d, want %d:\n%s", got, maxCtrlOToolOutputRunes-3, body) - } - - updated, _ = m.Update(tea.KeyMsg{Type: tea.KeyCtrlO}) - m = updated.(chatModel) - body = stripANSI(m.renderTranscript(100)) - if m.toolOutputOpen || m.entries[0].expanded || strings.Contains(body, strings.Repeat("x", 20)) { - t.Fatalf("second ctrl+o should collapse tool output: %q", body) - } -} - -func TestChatCompletedToolsGroupWhenNextStepStarts(t *testing.T) { - firstArgs := map[string]any{"command": "pwd"} - secondArgs := map[string]any{"command": "ls"} - thirdArgs := map[string]any{"command": "date"} - m := chatModel{} - - m.applyAgentEvent(coreagent.Event{Type: coreagent.EventToolStarted, ToolCallID: "call-1", ToolName: "bash", Args: firstArgs}) - m.applyAgentEvent(coreagent.Event{Type: coreagent.EventToolFinished, ToolCallID: "call-1", ToolName: "bash", Args: firstArgs, Content: "one"}) - m.applyAgentEvent(coreagent.Event{Type: coreagent.EventToolStarted, ToolCallID: "call-2", ToolName: "bash", Args: secondArgs}) - m.applyAgentEvent(coreagent.Event{Type: coreagent.EventToolFinished, ToolCallID: "call-2", ToolName: "bash", Args: secondArgs, Content: "two"}) - m.applyAgentEvent(coreagent.Event{Type: coreagent.EventToolStarted, ToolCallID: "call-3", ToolName: "bash", Args: thirdArgs}) - - if len(m.entries) != 2 { - t.Fatalf("entries = %d, want grouped history plus active tool: %#v", len(m.entries), m.entries) - } - if m.entries[0].role != "tool_group" || len(m.entries[0].tools) != 2 { - t.Fatalf("first entry should be grouped tool history: %#v", m.entries[0]) - } - if line := stripANSI(toolGroupStatusLine(m.entries[0])); line != "Ran 2 commands" { - t.Fatalf("grouped command line = %q", line) - } - if m.entries[1].status != "running" || m.entries[1].label != "Bash(\"date\")" { - t.Fatalf("second entry should be active tool: %#v", m.entries[1]) - } - if line := stripANSI(toolStatusLine(m.entries[1])); line != `Bash("date")` { - t.Fatalf("running command line = %q", line) - } -} - -func TestChatCompletedToolsGroupAcrossEmptyAssistantEntries(t *testing.T) { - entries := groupCompletedToolEntries([]chatEntry{ - newChatEntry(chatEntry{role: "tool", detail: "bash", label: `Bash("ls")`, status: "done", content: "listed"}), - newChatEntry(chatEntry{role: "assistant"}), - newChatEntry(chatEntry{role: "tool", detail: "read", label: `Read("agent")`, status: "error", err: "is directory", content: "is directory"}), - newChatEntry(chatEntry{role: "assistant", content: " "}), - newChatEntry(chatEntry{role: "tool", detail: "read", label: `Read("AGENTS.md")`, status: "done", content: "instructions"}), - newChatEntry(chatEntry{role: "assistant", content: "The files were listed."}), - }) - - if len(entries) != 2 { - t.Fatalf("entries = %d, want grouped tools plus assistant: %#v", len(entries), entries) - } - if entries[0].role != "tool_group" || len(entries[0].tools) != 3 { - t.Fatalf("first entry should group tools across empty assistant entries: %#v", entries[0]) - } - if line := stripANSI(toolGroupStatusLine(entries[0])); line != "Ran 1 command and read 2 files" { - t.Fatalf("grouped tool line = %q", line) - } - if entries[1].role != "assistant" || entries[1].content != "The files were listed." { - t.Fatalf("second entry should keep real assistant content: %#v", entries[1]) - } -} - -func TestChatCompletedToolsPreserveCollapsedThoughts(t *testing.T) { - entries := groupCompletedToolEntries([]chatEntry{ - newChatEntry(chatEntry{role: "tool", detail: "bash", label: `Bash("pwd")`, status: "done", content: "one"}), - newChatEntry(chatEntry{role: "thinking", label: "Thinking", content: "choose next tool", status: "done"}), - newChatEntry(chatEntry{role: "tool", detail: "bash", label: `Bash("ls")`, status: "done", content: "two"}), - newChatEntry(chatEntry{role: "assistant"}), - newChatEntry(chatEntry{role: "tool", detail: "read", label: `Read("AGENTS.md")`, status: "done", content: "instructions"}), - }) - - if len(entries) != 3 { - t.Fatalf("entries = %d, want tool, thought, and grouped tools: %#v", len(entries), entries) - } - if entries[0].role != "tool" || entries[1].role != "thinking" || entries[2].role != "tool_group" || len(entries[2].tools) != 2 { - t.Fatalf("collapsed thought should separate tool groups: %#v", entries) - } - if line := stripANSI(toolGroupStatusLine(entries[2])); line != "Ran 1 command and read a file" { - t.Fatalf("grouped tool line = %q", line) - } -} - -func TestChatCtrlOTogglesInlineToolOutput(t *testing.T) { - m := chatModel{ - width: 100, - height: 20, - entries: []chatEntry{ - {role: "tool", detail: "bash", label: "Bash(\"pwd\")", status: "done", content: "one"}, - {role: "assistant", content: "between"}, - {role: "tool", detail: "read", label: "Read(\"file\")", status: "error", err: "nope", content: "two"}, - }, - } - - updated, _ := m.Update(tea.KeyMsg{Type: tea.KeyCtrlO}) - m = updated.(chatModel) - for _, index := range []int{0, 2} { - if !m.entries[index].expanded { - t.Fatalf("tool entry %d should be expanded inline", index) - } - } - view := stripANSI(m.renderTranscript(100)) - if strings.Contains(view, "Tool details") { - t.Fatalf("ctrl+o should keep tool output inline: %q", view) - } - if !strings.Contains(view, "one") || !strings.Contains(view, "two") { - t.Fatalf("view missing inline expanded tool output: %q", view) - } - if !strings.Contains(view, "between") { - t.Fatalf("inline tool output should keep surrounding chat visible: %q", view) - } - - updated, _ = m.Update(tea.KeyMsg{Type: tea.KeyCtrlO}) - m = updated.(chatModel) - for _, index := range []int{0, 2} { - if m.entries[index].expanded { - t.Fatalf("tool entry %d should remain collapsed", index) - } - } -} - -func TestChatCtrlOTogglesInlineOutput(t *testing.T) { - m := chatModel{ - width: 100, - height: 24, - entries: []chatEntry{ - {role: "user", content: "who is parth sareen"}, - {role: "tool", detail: "web_search", label: "Web Search(\"who is Parth Sareen\")", status: "done", content: "Search results"}, - {role: "assistant", content: "Based on public search results, Parth Sareen works on AI tooling."}, - }, - } - - updated, _ := m.Update(tea.KeyMsg{Type: tea.KeyCtrlO}) - m = updated.(chatModel) - if !m.toolOutputOpen || !m.entries[1].expanded { - t.Fatalf("tool output should be expanded inline: %#v", m.entries[1]) - } - - updated, _ = m.Update(tea.KeyMsg{Type: tea.KeyCtrlO}) - m = updated.(chatModel) - if m.toolOutputOpen || m.entries[1].expanded { - t.Fatalf("tool output should be collapsed inline: %#v", m.entries[1]) - } -} - -func TestChatCtrlOTogglesCompletedThinkingDetails(t *testing.T) { - m := chatModel{ - entries: []chatEntry{ - newChatEntry(chatEntry{role: "thinking", label: "Thinking", status: "done", content: "private reasoning", tokenCount: 12}), - }, - } - - if view := stripANSI(m.renderTranscript(100)); !strings.Contains(view, "Thought") || strings.Contains(view, "12 tokens") { - t.Fatalf("collapsed thinking should hide its token count:\n%s", view) - } - - updated, _ := m.Update(tea.KeyMsg{Type: tea.KeyCtrlO}) - m = updated.(chatModel) - if !m.thinkingDetailsOpen || !m.entries[0].expanded { - t.Fatalf("ctrl+o should expand completed thinking entries: %#v", m.entries[0]) - } - if view := stripANSI(m.renderTranscript(100)); !strings.Contains(view, "Thought (12 tokens)") || !strings.Contains(view, "private reasoning") { - t.Fatalf("ctrl+o should render completed thinking content:\n%s", view) - } - - m.liveMessages = []api.Message{{Role: "assistant", Thinking: "live private reasoning"}} - m.syncThinkingEntry("live private reasoning") - if !m.entries[1].expanded { - t.Fatalf("live thinking should always be expanded: %#v", m.entries[1]) - } - if view := stripANSI(m.renderTranscript(100)); !strings.Contains(view, "live private reasoning") { - t.Fatalf("live thinking should render while streaming:\n%s", view) - } - m.thinkingTokens = 9 - m.finishThinkingEntry() - if m.entries[1].status != "done" || !m.entries[1].expanded { - t.Fatalf("completed thinking should honor the open details mode: %#v", m.entries[1]) - } - - updated, _ = m.Update(tea.KeyMsg{Type: tea.KeyCtrlO}) - m = updated.(chatModel) - if m.thinkingDetailsOpen || m.entries[0].expanded || m.entries[1].expanded { - t.Fatalf("ctrl+o should collapse all completed thinking details: %#v", m.entries) - } -} - -func TestChatThinkingBodyUsesSecondaryGrey(t *testing.T) { - for _, tt := range []struct { - entry chatEntry - header string - }{ - {entry: chatEntry{role: "thinking", status: "running", content: "Let me inspect the files.", expanded: true}, header: "Thinking"}, - {entry: chatEntry{role: "thinking", status: "done", content: "Let me inspect the files.", tokenCount: 15, expanded: true}, header: "Thought (15 tokens)"}, - } { - lines := renderThinkingLines(tt.entry, 80) - if len(lines) < 3 { - t.Fatalf("thinking entry did not render its body: %#v", lines) - } - if lines[0] != tt.header { - t.Fatalf("thinking header = %q, want unmuted %q", lines[0], tt.header) - } - if got, want := lines[2], chatToolOutputStyle.Render("Let me inspect the files."); got != want { - t.Fatalf("thinking body = %q, want secondary grey %q", got, want) - } - } -} - -func TestChatThinkingBodyAlignsWithStatusText(t *testing.T) { - m := chatModel{entries: []chatEntry{ - {role: "thinking", status: "done", content: "Let me inspect the files.", tokenCount: 15, expanded: true}, - }} - - lines := strings.Split(stripANSI(m.renderTranscript(80)), "\n") - if len(lines) < 3 { - t.Fatalf("thinking entry did not render its body: %#v", lines) - } - if got, want := lines[0], "• Thought (15 tokens)"; got != want { - t.Fatalf("thinking header = %q, want %q", got, want) - } - if got, want := lines[2], " Let me inspect the files."; got != want { - t.Fatalf("thinking body = %q, want aligned with status text %q", got, want) - } -} - -func TestChatLiveThinkingUsesBoundedSanitizedTail(t *testing.T) { - longThinking := "first-visible-marker\n" + strings.Repeat("x", maxLiveThinkingRunes+100) + "\nlast-visible-marker\x1b[2J" - m := chatModel{entries: []chatEntry{{role: "thinking", label: "Thinking ↓ 10 tokens", status: "running", content: longThinking, expanded: true}}} - - view := stripANSI(m.renderTranscript(80)) - if strings.Contains(view, "first-visible-marker") || !strings.Contains(view, "earlier thinking omitted while streaming") || !strings.Contains(view, "last-visible-marker") { - t.Fatalf("live thinking should render a bounded tail:\n%s", view) - } - if strings.Contains(view, "\x1b") || strings.Contains(view, "[2J") { - t.Fatalf("live thinking should remove terminal control sequences: %q", view) - } -} - -func TestChatThinkingSanitizesTerminalControlsLiveAndReopened(t *testing.T) { - content := "safe\x1b]52;c;clipboard-payload\a visible\x1bPqdevice-control\x1b\\ done\b!\u009b31m red\a" - tests := []struct { - name string - entry chatEntry - }{ - {name: "live", entry: chatEntry{role: "thinking", status: "running", content: content, expanded: true}}, - {name: "reopened", entry: chatEntry{role: "thinking", status: "done", content: content, expanded: true}}, - } - - for _, tt := range tests { - t.Run(tt.name, func(t *testing.T) { - view := chatModel{entries: []chatEntry{tt.entry}}.renderTranscript(80) - plain := stripANSI(view) - if !strings.Contains(plain, "safe") || !strings.Contains(plain, "visible") || !strings.Contains(plain, "done!") || !strings.Contains(plain, "red") { - t.Fatalf("sanitized thinking lost visible text: %q", plain) - } - for _, unsafe := range []string{"clipboard-payload", "device-control", "\x1b]", "\x1bP", "\a", "\b", "\u009b"} { - if strings.Contains(view, unsafe) { - t.Fatalf("sanitized thinking contains unsafe value %q: %q", unsafe, view) - } - } - }) - } -} - -func TestLiveThinkingTailPreservesUTF8Boundary(t *testing.T) { - content := strings.Repeat("🙂", maxLiveThinkingRunes+100) - tail, omitted := liveThinkingTail(content, maxLiveThinkingRunes) - if !omitted { - t.Fatal("long thinking trace should report an omitted prefix") - } - if !utf8.ValidString(tail) || utf8.RuneCountInString(tail) != maxLiveThinkingRunes { - t.Fatalf("tail should contain %d complete runes, got %d", maxLiveThinkingRunes, utf8.RuneCountInString(tail)) - } -} - -func BenchmarkRenderLiveThinkingLines(b *testing.B) { - content := strings.Repeat("old thinking that is outside the visible tail\n", 100_000) + "visible tail" - entry := chatEntry{role: "thinking", status: "running", content: content, expanded: true} - b.ReportAllocs() - for b.Loop() { - renderThinkingLines(entry, 80) - } -} - -func TestChatThinkingDetailsRespectNarrowAndResizedViews(t *testing.T) { - m := chatModel{ - width: 80, - height: 10, - entries: []chatEntry{ - {role: "thinking", label: "Thinking ↓ 42 tokens", status: "running", content: "first delta\nsecond delta", expanded: true}, - }, - } - - updated, _ := m.Update(tea.WindowSizeMsg{Width: 24, Height: 8}) - m = updated.(chatModel) - transcript := stripANSI(m.renderTranscript(24)) - if !strings.Contains(transcript, "Thinking") || !strings.Contains(transcript, "second delta") { - t.Fatalf("narrow transcript should retain the live thinking tail:\n%s", transcript) - } - if len(m.transcriptLines(24)) == 0 { - t.Fatal("resized transcript should remain selectable and scrollable") - } -} - -func TestChatCtrlOCollapseKeepsInputPromptVisible(t *testing.T) { - m := chatModel{ - width: 80, - height: 8, - opts: Options{Model: "gemma4"}, - entries: []chatEntry{ - {role: "user", content: "inspect this"}, - {role: "tool", detail: "bash", label: "Bash(\"pwd\")", status: "done", content: strings.Repeat("output\n", 20)}, - {role: "assistant", content: "Done."}, - }, - } - - updated, _ := m.Update(tea.KeyMsg{Type: tea.KeyCtrlO}) - m = updated.(chatModel) - updated, _ = m.Update(tea.KeyMsg{Type: tea.KeyCtrlO}) - m = updated.(chatModel) - - view := stripANSI(m.View()) - if !strings.Contains(view, "│ █") { - t.Fatalf("input prompt disappeared after collapsing tool output:\n%s", view) - } -} - -func TestChatCtrlOShowsRunningToolOutputInline(t *testing.T) { - args := map[string]any{"command": "pwd"} - m := chatModel{ - width: 100, - height: 20, - running: true, - opts: Options{Model: "test-model"}, - } - m.applyAgentEvent(coreagent.Event{ - Type: coreagent.EventToolStarted, - ToolCallID: "call-1", - ToolName: "bash", - Args: args, - }) - - updated, _ := m.Update(tea.KeyMsg{Type: tea.KeyCtrlO}) - m = updated.(chatModel) - if !m.entries[0].expanded { - t.Fatalf("ctrl+o should expand the running tool inline") - } - view := stripANSI(m.View()) - if !strings.Contains(view, `Bash("pwd")`) || !strings.Contains(view, "$ pwd") { - t.Fatalf("expanded transcript missing running tool details:\n%s", view) - } - - m.applyAgentEvent(coreagent.Event{ - Type: coreagent.EventToolFinished, - ToolCallID: "call-1", - ToolName: "bash", - Args: args, - Content: "/tmp/project\n", - }) - - view = stripANSI(m.View()) - if !strings.Contains(view, "/tmp/project") { - t.Fatalf("finished tool output should be visible inline: %q", view) - } - if !strings.Contains(view, "│ █") { - t.Fatalf("input prompt disappeared while tool output is expanded:\n%s", view) - } - if !strings.Contains(view, "test-model") { - t.Fatalf("footer/model line disappeared while tool output is expanded:\n%s", view) - } - - updated, _ = m.Update(tea.KeyMsg{Type: tea.KeyCtrlO}) - m = updated.(chatModel) - view = stripANSI(m.View()) - if !strings.Contains(view, "│ █") { - t.Fatalf("input prompt disappeared after collapsing tool output:\n%s", view) - } - if !strings.Contains(view, "test-model") { - t.Fatalf("footer/model line disappeared after collapsing tool output:\n%s", view) - } -} - -func TestChatLongCommandLabelTruncatesAndCtrlOShowsFullCommand(t *testing.T) { - command := strings.Repeat("x", 120) - m := chatModel{ - width: 180, - height: 20, - } - m.applyAgentEvent(coreagent.Event{ - Type: coreagent.EventToolStarted, - ToolCallID: "call-1", - ToolName: "bash", - Args: map[string]any{"command": command}, - }) - - wantLabel := `Bash("` + strings.Repeat("x", 100) + `...")` - if m.entries[0].label != wantLabel { - t.Fatalf("label = %q, want %q", m.entries[0].label, wantLabel) - } - if strings.Contains(m.entries[0].label, strings.Repeat("x", 101)) { - t.Fatalf("collapsed label should truncate command: %q", m.entries[0].label) - } - - updated, _ := m.Update(tea.KeyMsg{Type: tea.KeyCtrlO}) - m = updated.(chatModel) - view := stripANSI(m.renderTranscript(180)) - if !strings.Contains(view, "$ "+command) { - t.Fatalf("expanded ctrl+o view should show full command:\n%s", view) - } -} - -func TestChatCtrlOInlineOutputSurvivesToolGrouping(t *testing.T) { - firstArgs := map[string]any{"command": "pwd"} - secondArgs := map[string]any{"command": "ls"} - m := chatModel{ - width: 100, - height: 24, - entries: []chatEntry{ - newChatEntry(chatEntry{role: "tool", detail: "bash", label: "Bash(\"pwd\")", status: "done", content: "one", args: firstArgs}), - newChatEntry(chatEntry{role: "tool", detail: "bash", label: "Bash(\"ls\")", status: "done", content: "two", args: secondArgs}), - }, - } - - updated, _ := m.Update(tea.KeyMsg{Type: tea.KeyCtrlO}) - m = updated.(chatModel) - m.applyAgentEvent(coreagent.Event{ - Type: coreagent.EventToolStarted, - ToolCallID: "call-3", - ToolName: "bash", - Args: map[string]any{"command": "date"}, - }) - - if len(m.entries) != 2 { - t.Fatalf("entries = %d, want tool group plus active tool: %#v", len(m.entries), m.entries) - } - if m.entries[0].role != "tool_group" || !m.entries[0].expanded { - t.Fatalf("grouped tool history should stay expanded inline: %#v", m.entries[0]) - } - if line := stripANSI(toolGroupStatusLine(m.entries[0])); line != "Ran 2 commands" { - t.Fatalf("grouped command line = %q", line) - } - - view := stripANSI(m.renderTranscript(100)) - for _, want := range []string{`Bash("pwd")`, `Bash("ls")`, `Bash("date")`, "one", "two"} { - if !strings.Contains(view, want) { - t.Fatalf("grouped tool output missing %q:\n%s", want, view) - } - } - if strings.Contains(view, " Ran 1 command") { - t.Fatalf("expanded grouped children should show concrete invocations, not generic summaries:\n%s", view) - } -} - -func TestChatExpandedToolGroupRendersIndentedToolBlocks(t *testing.T) { - startedAt := time.Date(2026, 6, 11, 12, 0, 0, 0, time.UTC) - m := chatModel{ - entries: []chatEntry{ - newChatEntry(chatEntry{ - role: "tool_group", - label: "Tool calls (2)", - status: "done", - expanded: true, - tools: []chatEntry{ - newChatEntry(chatEntry{ - role: "tool", - detail: "web_search", - label: "Web Search(\"Parth Sareen\")", - status: "done", - content: "Search results for: Parth Sareen", - startedAt: startedAt, - finishedAt: startedAt.Add(823 * time.Millisecond), - }), - newChatEntry(chatEntry{ - role: "tool", - detail: "read", - label: "Read(\"feedback.md\")", - status: "done", - content: "Looks good.", - startedAt: startedAt.Add(time.Second), - finishedAt: startedAt.Add(3 * time.Second), - }), - }, - }), - }, - } - - transcript := stripANSI(m.renderTranscript(120)) - expected := strings.Join([]string{ - " Web Search(\"Parth Sareen\")", - " Search results for: Parth Sareen", - " ", - " Read(\"feedback.md\")", - " Looks good.", - }, "\n") - if !strings.Contains(transcript, expected) { - t.Fatalf("expanded tool group did not render expected block spacing:\n%s", transcript) - } -} - -func TestChatExpandedToolGroupCapsChildOutput(t *testing.T) { - longOutput := strings.Repeat("y", maxCtrlOToolOutputRunes+25) + "group-tail-marker" - m := chatModel{ - entries: []chatEntry{ - newChatEntry(chatEntry{ - role: "tool_group", - status: "done", - expanded: true, - tools: []chatEntry{ - newChatEntry(chatEntry{ - role: "tool", - detail: "read", - label: "Read(\"big.log\")", - status: "done", - content: longOutput, - }), - }, - }), - }, - } - - transcript := stripANSI(m.renderTranscript(120)) - if strings.Contains(transcript, "group-tail-marker") { - t.Fatalf("expanded grouped tool output should be capped:\n%s", transcript) - } - if !strings.Contains(transcript, "...") { - t.Fatalf("expanded grouped tool output should show capped output ellipsis:\n%s", transcript) - } - if got := strings.Count(transcript, "y"); got != maxCtrlOToolOutputRunes-3 { - t.Fatalf("expanded grouped tool output y count = %d, want %d:\n%s", got, maxCtrlOToolOutputRunes-3, transcript) - } - if got := m.entries[0].tools[0].content; got != longOutput { - t.Fatal("ctrl+o rendering should not mutate grouped tool content") - } -} - -func TestChatToolGroupStatusOmitsResultCounts(t *testing.T) { - startedAt := time.Date(2026, 6, 15, 17, 0, 0, 0, time.UTC) - entry := newChatEntry(chatEntry{ - role: "tool_group", - label: "Tool calls (3)", - status: "error", - startedAt: startedAt, - finishedAt: startedAt.Add(3 * time.Second), - tools: []chatEntry{ - newChatEntry(chatEntry{role: "tool", status: "done"}), - newChatEntry(chatEntry{role: "tool", status: "done"}), - newChatEntry(chatEntry{role: "tool", status: "error", err: "failed"}), - }, - }) - - line := stripANSI(toolGroupStatusLine(entry)) - for _, word := range []string{"succeeded", "failed", "done", "in 3s", "3s"} { - if strings.Contains(line, word) { - t.Fatalf("group status = %q, should not show status word or elapsed %q", line, word) - } - } - if line != "Used 3 tools" { - t.Fatalf("group status = %q, want action summary", line) - } -} - -func TestChatToolGroupStatusShowsAllSuccessCount(t *testing.T) { - entry := newChatEntry(chatEntry{ - role: "tool_group", - label: "Tool calls (2)", - status: "done", - tools: []chatEntry{ - newChatEntry(chatEntry{role: "tool", status: "done"}), - newChatEntry(chatEntry{role: "tool", status: "done"}), - }, - }) - - line := stripANSI(toolGroupStatusLine(entry)) - if strings.Contains(line, "succeeded") || strings.Contains(line, "done") { - t.Fatalf("group status = %q, should not show success count or done", line) - } - if line != "Used 2 tools" { - t.Fatalf("group status = %q, want action summary", line) - } -} - -func TestChatToolGroupStatusRecomputesCachedLabel(t *testing.T) { - entries := groupCompletedToolEntries([]chatEntry{ - newChatEntry(chatEntry{ - role: "tool_group", - label: "Ran 1 command", - status: "done", - tools: []chatEntry{ - newChatEntry(chatEntry{role: "tool", detail: "bash", label: `Bash("pwd")`, status: "done"}), - }, - }), - newChatEntry(chatEntry{ - role: "tool", - detail: "bash", - label: `Bash("ls")`, - status: "done", - }), - }) - - if len(entries) != 1 || entries[0].role != "tool_group" { - t.Fatalf("entries = %#v, want regrouped tool group", entries) - } - if line := stripANSI(toolGroupStatusLine(entries[0])); line != "Ran 2 commands" { - t.Fatalf("group status = %q, want recomputed action summary", line) - } -} - -func TestChatToolGroupStatusKeepsHiddenDetectedCount(t *testing.T) { - args := map[string]any{"command": "ls"} - entry := newChatEntry(chatEntry{ - role: "tool_group", - label: "Ran 2 commands", - status: "done", - tools: []chatEntry{ - newChatEntry(chatEntry{role: "tool", detail: "bash", label: `Bash("pwd")`, status: "done"}), - }, - }) - entry.tools[0].args = map[string]any{"command": "pwd"} - detected := []chatEntry{ - newChatEntry(chatEntry{role: "tool", detail: "bash", label: `Bash("ls")`, status: "queued", args: args}), - } - - grouped := groupCompletedToolEntries([]chatEntry{entry}, detected...) - if len(grouped) != 1 || grouped[0].role != "tool_group" || len(grouped[0].tools) != 1 { - t.Fatalf("grouped entries = %#v, want visible finished child only", grouped) - } - if line := stripANSI(toolGroupStatusLine(grouped[0])); line != "Ran 2 commands" { - t.Fatalf("group status = %q, want hidden detected command counted", line) - } -} - -func TestChatToolOutputRendersUnifiedDiff(t *testing.T) { - diff := strings.Join([]string{ - "diff --git a/file.go b/file.go", - "index 1111111..2222222 100644", - "--- a/file.go", - "+++ b/file.go", - "@@ -1,3 +1,3 @@", - " package main", - "-var old = true", - "+var newer = true", - }, "\n") - m := chatModel{ - entries: []chatEntry{{ - role: "tool", - detail: "bash", - label: "Bash(\"git diff\")", - status: "done", - content: diff, - }}, - } - - updated, _ := m.Update(tea.KeyMsg{Type: tea.KeyCtrlO}) - m = updated.(chatModel) - rendered := m.renderTranscript(100) - body := stripANSI(rendered) - if !strings.Contains(body, "diff --git a/file.go b/file.go") || - !strings.Contains(body, "-var old = true") || - !strings.Contains(body, "+var newer = true") { - t.Fatalf("rendered diff missing expected lines: %q", body) - } - if !looksLikeUnifiedDiff(diff) { - t.Fatal("diff output should be detected as a unified diff") - } -} - -func TestChatToolCallRendersPrettyInvocationAndResult(t *testing.T) { - args := api.NewToolCallFunctionArguments() - args.Set("query", "Parth Sareen Ollama software engineer") - - m := chatModel{width: 100, height: 30} - m.applyAgentEvent(coreagent.Event{ - Type: coreagent.EventToolCallDetected, - ToolCalls: []api.ToolCall{{ - ID: "call-1", - Function: api.ToolCallFunction{ - Name: "web_search", - Arguments: args, - }, - }}, - }) - m.applyAgentEvent(coreagent.Event{ - Type: coreagent.EventToolStarted, - ToolCallID: "call-1", - ToolName: "web_search", - Args: args.ToMap(), - }) - m.applyAgentEvent(coreagent.Event{ - Type: coreagent.EventToolFinished, - ToolCallID: "call-1", - ToolName: "web_search", - Args: args.ToMap(), - Content: "**Search results for:** Parth Sareen\n\n1. Parth Sareen\n URL: https://parthsareen.com\n", - }) - - transcript := stripANSI(m.renderTranscript(100)) - if !strings.Contains(transcript, "Web Search(\"Parth Sareen Ollama software engineer\")") { - t.Fatalf("transcript missing invocation: %q", transcript) - } - for _, word := range []string{"done", "in 6s", "6s"} { - if strings.Contains(transcript, word) { - t.Fatalf("transcript should not include status word/elapsed %q: %q", word, transcript) - } - } - if strings.Contains(transcript, "https://parthsareen.com") || strings.Contains(transcript, "Search results for:") { - t.Fatalf("tool output should be collapsed by default: %q", transcript) - } - if strings.Contains(transcript, "web_search") { - t.Fatalf("transcript should use display name instead of raw tool name: %q", transcript) - } - if m.entries[0].status != "done" { - t.Fatalf("tool status = %q, want done", m.entries[0].status) - } - - updated, _ := m.Update(tea.KeyMsg{Type: tea.KeyCtrlO}) - m = updated.(chatModel) - view := stripANSI(m.renderTranscript(100)) - if strings.Contains(view, "**") { - t.Fatalf("inline web output should render Markdown, not show delimiters: %q", view) - } - if !strings.Contains(view, "Search results for:") || !strings.Contains(view, "https://parthsareen.com") { - t.Fatalf("inline web output missing content: %q", view) - } -} - -func TestChatToolOutputHidesInternalTruncationMarker(t *testing.T) { - entry := chatEntry{ - role: "tool", - detail: "bash", - status: "done", - expanded: true, - content: "head\n\n[tool output truncated: showing first ~10 tokens and last ~10 tokens; omitted ~25 tokens. Use a narrower command, line range, or search query if more detail is needed.]\n\ntail", - } - - rendered := stripANSI(strings.Join(renderToolResultLines(entry, 80), "\n")) - if strings.Contains(rendered, "tool output truncated") || strings.Contains(rendered, "omitted ~25 tokens") { - t.Fatalf("internal truncation marker should be hidden from tool output: %q", rendered) - } - if !strings.Contains(rendered, "head") || !strings.Contains(rendered, "tail") { - t.Fatalf("tool output should keep visible content: %q", rendered) - } -} - -func TestChatRenderTranscriptCachesEntryLinesUntilDirty(t *testing.T) { - m := chatModel{ - entries: []chatEntry{ - newChatEntry(chatEntry{role: "assistant", content: "**hello**"}), - }, - } - - _ = m.renderTranscript(80) - if len(m.entries[0].renderLines) == 0 { - t.Fatal("rendered entry lines should be cached") - } - - m.entries[0].renderLines = []string{"cached line"} - if got := stripANSI(m.renderTranscript(80)); !strings.Contains(got, "cached line") { - t.Fatalf("render should reuse cached lines for unchanged entry: %q", got) - } - - m.entries[0].content = "**changed**" - m.markEntryDirty(0) - if got := stripANSI(m.renderTranscript(80)); strings.Contains(got, "cached line") || !strings.Contains(got, "changed") { - t.Fatalf("dirty entry should re-render instead of using stale cache: %q", got) - } -} - -func TestWrapChatTextSplitsLongLines(t *testing.T) { - lines := wrapChatText("alpha beta gamma delta", 12) - if len(lines) < 2 { - t.Fatalf("lines = %#v, want split text", lines) - } - if strings.Contains(lines[0], "delta") { - t.Fatalf("first line was not wrapped: %#v", lines) - } -} - -func TestWrapChatTextUsesDisplayWidth(t *testing.T) { - lines := wrapChatText(strings.Repeat("界", 20), 20) - if len(lines) < 2 { - t.Fatalf("lines = %#v, want full-width text split", lines) - } - for _, line := range lines { - if got := lipgloss.Width(line); got > 20 { - t.Fatalf("line %q width = %d, want <= 20", line, got) - } - } -} - -func TestRenderMarkdownTableWrapsLongCells(t *testing.T) { - markdown := strings.Join([]string{ - "| # | Item | Why it matters |", - "|---|---|---|", - "| **C** | **Bash filesystem/network confinement** | Biggest asymmetry: file tools are sandboxed via `os.Root`, bash is not and this text should survive until tail-token. |", - }, "\n") - - rendered := renderMarkdownForView(markdown, 72) - plain := stripANSI(rendered) - if !strings.Contains(plain, "tail-token") { - t.Fatalf("long table cell was truncated:\n%s", plain) - } - if strings.Contains(plain, "tail-toke...") { - t.Fatalf("long table cell should wrap, not ellipsize:\n%s", plain) - } - for _, line := range strings.Split(rendered, "\n") { - if got := lipgloss.Width(line); got > 72 { - t.Fatalf("rendered table line width = %d, want <= 72: %q\n%s", got, stripANSI(line), plain) - } - } -} - -func TestRenderMarkdownProseWithPipeExamplesIsNotTable(t *testing.T) { - markdown := strings.Join([]string{ - "- **Regression confirmed** vs. `fdbe8d33`: the bare `< | open | >` remains prose.", - "| | |", - "- **Severity**: `< | close | >` is another inline example.", - }, "\n") - - rendered := renderMarkdownForView(markdown, 120) - plain := stripANSI(rendered) - if !strings.Contains(plain, "| | |") || !strings.Contains(plain, "the bare < | open | > remains") || !strings.Contains(plain, "< | close | > is another") { - t.Fatalf("pipe-delimited prose rendered as a table:\n%s", plain) - } - if !strings.Contains(rendered, chatStrongStyle.Render("Regression confirmed")) { - t.Fatalf("bold prose was not emphasized: %q", rendered) - } - if !strings.Contains(rendered, chatInlineCodeStyle.Render("fdbe8d33")) { - t.Fatalf("inline code was not styled: %q", rendered) - } -} - -func TestRenderMarkdownTablePreservesValidSeparator(t *testing.T) { - markdown := strings.Join([]string{ - "| Name | State |", - "| --- | :---: |", - "| Ollama | Ready |", - }, "\n") - - plain := stripANSI(renderMarkdownForView(markdown, 80)) - if strings.Contains(plain, "---") { - t.Fatalf("table separator should not render as prose:\n%s", plain) - } - if !strings.Contains(plain, "Name") || !strings.Contains(plain, "Ollama") { - t.Fatalf("valid Markdown table was not rendered:\n%s", plain) - } -} diff --git a/cmd/tui/chat/test_helpers_test.go b/cmd/tui/chat/test_helpers_test.go deleted file mode 100644 index 96af086e9a0..00000000000 --- a/cmd/tui/chat/test_helpers_test.go +++ /dev/null @@ -1,113 +0,0 @@ -package chat - -import ( - "context" - "fmt" - "regexp" - "testing" - "time" - - tea "github.com/charmbracelet/bubbletea" - - coreagent "github.com/ollama/ollama/agent" - "github.com/ollama/ollama/api" -) - -type chatTestTool struct{} - -type chatTestClient struct{} - -type chatCaptureClient struct { - requests []*api.ChatRequest -} - -type chatToolLoopClient struct { - calls int - toolRounds int -} - -func (chatTestClient) Chat(ctx context.Context, req *api.ChatRequest, fn api.ChatResponseFunc) error { - if err := ctx.Err(); err != nil { - return err - } - return fn(api.ChatResponse{ - Message: api.Message{Role: "assistant", Content: "ok"}, - Done: true, - }) -} - -func (c *chatCaptureClient) Chat(ctx context.Context, req *api.ChatRequest, fn api.ChatResponseFunc) error { - if err := ctx.Err(); err != nil { - return err - } - c.requests = append(c.requests, req) - return fn(api.ChatResponse{ - Message: api.Message{Role: "assistant", Content: "ok"}, - Done: true, - }) -} - -func (c *chatToolLoopClient) Chat(ctx context.Context, req *api.ChatRequest, fn api.ChatResponseFunc) error { - if err := ctx.Err(); err != nil { - return err - } - c.calls++ - if c.calls > c.toolRounds { - return fn(api.ChatResponse{Message: api.Message{Role: "assistant", Content: "done"}, Done: true}) - } - - args := api.NewToolCallFunctionArguments() - args.Set("value", "keep going") - return fn(api.ChatResponse{Message: api.Message{Role: "assistant", ToolCalls: []api.ToolCall{{ - ID: fmt.Sprintf("call-%d", c.calls), - Function: api.ToolCallFunction{ - Name: "fake_tool", - Arguments: args, - }, - }}}}) -} - -func (chatTestTool) Name() string { - return "fake_tool" -} - -func (chatTestTool) Description() string { - return "does test work" -} - -func (chatTestTool) Schema() api.ToolFunction { - return api.ToolFunction{ - Name: "fake_tool", - Description: "does test work", - Parameters: api.ToolFunctionParameters{ - Type: "object", - }, - } -} - -func (chatTestTool) Execute(context.Context, coreagent.ToolContext, map[string]any) (coreagent.ToolResult, error) { - return coreagent.ToolResult{Content: "ok"}, nil -} - -func waitForRunDone(t *testing.T, events <-chan tea.Msg) chatRunDoneMsg { - t.Helper() - timeout := time.After(2 * time.Second) - for { - select { - case msg, ok := <-events: - if !ok { - t.Fatal("events closed before run done") - } - if done, ok := msg.(chatRunDoneMsg); ok { - return done - } - case <-timeout: - t.Fatal("timed out waiting for run done") - } - } -} - -func stripANSI(s string) string { - re := regexp.MustCompile(`\x1b\[[0-9;:]*[A-Za-z]`) - return re.ReplaceAllString(s, "") -} diff --git a/cmd/tui/chat/theme.go b/cmd/tui/chat/theme.go deleted file mode 100644 index 375f063d1ed..00000000000 --- a/cmd/tui/chat/theme.go +++ /dev/null @@ -1,129 +0,0 @@ -package chat - -import "github.com/charmbracelet/lipgloss" - -const ( - chatAnsiRed = "1" - chatAnsiGreen = "2" - chatAnsiYellow = "3" - chatAnsiBlue = "4" - chatAnsiCyan = "6" - chatAnsiBrightBlack = "8" -) - -var ( - chatHeaderStyle = lipgloss.NewStyle(). - Bold(true) - - chatMetaStyle = lipgloss.NewStyle(). - Faint(true) - - chatFooterStyle = lipgloss.NewStyle(). - Faint(true) - - chatInputBorderStyle = lipgloss.NewStyle(). - Faint(true) - - chatInputPlaceholderStyle = lipgloss.NewStyle(). - Foreground(lipgloss.Color("8")) - - chatCursorStyle = lipgloss.NewStyle(). - Reverse(true) - - chatBlankCursorStyle = lipgloss.NewStyle(). - Faint(true) - - chatNotificationStyle = chatMetaStyle - - chatUserStyle = lipgloss.NewStyle() - - chatUserBlockStyle = lipgloss.NewStyle(). - Foreground(lipgloss.AdaptiveColor{Light: "#777777", Dark: "#8a8a8a"}) - - chatToolStyle = lipgloss.NewStyle() - - chatInlineCodeStyle = lipgloss.NewStyle(). - Bold(true) - - chatStrongStyle = lipgloss.NewStyle(). - Bold(true) - - chatCodeBlockStyle = lipgloss.NewStyle() - - chatTableBorderStyle = lipgloss.NewStyle(). - Faint(true) - - chatToolRunningStyle = lipgloss.NewStyle(). - Foreground(lipgloss.Color(chatAnsiYellow)) - - chatToolDoneStyle = lipgloss.NewStyle(). - Foreground(lipgloss.Color(chatAnsiGreen)) - - // chatToolMixedStyle marks a tool group with both succeeded and failed - // calls (partial success). Amber/orange is distinct from green (success), - // red (failure), and yellow (running). - chatToolMixedStyle = lipgloss.NewStyle(). - Foreground(lipgloss.Color("208")) - - chatToolOutputStyle = lipgloss.NewStyle(). - Foreground(lipgloss.AdaptiveColor{Light: "#666666", Dark: "#a0a0a0"}) - - chatDiffMetaStyle = lipgloss.NewStyle(). - Faint(true) - - chatDiffFileStyle = lipgloss.NewStyle(). - Bold(true). - Foreground(lipgloss.Color(chatAnsiCyan)) - - chatDiffHunkStyle = lipgloss.NewStyle(). - Foreground(lipgloss.Color(chatAnsiBlue)) - - chatDiffAddStyle = lipgloss.NewStyle(). - Foreground(lipgloss.Color(chatAnsiGreen)) - - chatDiffDeleteStyle = lipgloss.NewStyle(). - Foreground(lipgloss.Color(chatAnsiRed)) - - chatErrorStyle = lipgloss.NewStyle(). - Foreground(lipgloss.Color(chatAnsiRed)) - - chatFullAccessStyle = lipgloss.NewStyle(). - Foreground(lipgloss.AdaptiveColor{Light: "#9f5f5f", Dark: "#b87373"}) - - chatCommandNameStyle = lipgloss.NewStyle() - - chatPickerTextStyle = lipgloss.NewStyle() - - chatPickerTitleStyle = lipgloss.NewStyle(). - Bold(true) - - chatPickerSelectedStyle = lipgloss.NewStyle(). - Bold(true) - - chatPickerMetaStyle = lipgloss.NewStyle(). - Faint(true) - - chatHistoryTitleStyle = lipgloss.NewStyle(). - Bold(true) - - chatHistorySystemRoleStyle = lipgloss.NewStyle(). - Bold(true). - Faint(true) - - chatHistoryUserRoleStyle = lipgloss.NewStyle(). - Bold(true). - Foreground(lipgloss.Color(chatAnsiBlue)) - - chatHistoryAssistantRoleStyle = lipgloss.NewStyle(). - Bold(true). - Foreground(lipgloss.Color(chatAnsiYellow)) - - chatHistoryToolRoleStyle = lipgloss.NewStyle(). - Bold(true). - Foreground(lipgloss.Color(chatAnsiGreen)) - - chatHistoryLabelStyle = lipgloss.NewStyle(). - Faint(true) - - chatHistoryTextStyle = lipgloss.NewStyle() -) diff --git a/cmd/tui/chat/think.go b/cmd/tui/chat/think.go deleted file mode 100644 index 074bfa68384..00000000000 --- a/cmd/tui/chat/think.go +++ /dev/null @@ -1,166 +0,0 @@ -package chat - -import ( - "fmt" - "strings" - - tea "github.com/charmbracelet/bubbletea" - - "github.com/ollama/ollama/api" -) - -type chatThinkOption struct { - value string - label string - description string -} - -type chatThinkPicker struct { - options []chatThinkOption - cursor int -} - -var chatThinkOptions = []chatThinkOption{ - {value: "auto", label: "auto", description: "use the model default"}, - {value: "on", label: "on", description: "enable thinking"}, - {value: "off", label: "off", description: "disable thinking"}, - {value: "low", label: "low", description: "use low thinking effort"}, - {value: "medium", label: "medium", description: "use medium thinking effort"}, - {value: "high", label: "high", description: "use high thinking effort"}, - {value: "max", label: "max", description: "use maximum thinking effort"}, -} - -func (m *chatModel) openThinkPicker() (tea.Model, tea.Cmd) { - m.thinkPicker = newChatThinkPicker(m.opts.Think) - m.status = "think" - return *m, nil -} - -func newChatThinkPicker(current *api.ThinkValue) *chatThinkPicker { - picker := &chatThinkPicker{options: append([]chatThinkOption(nil), chatThinkOptions...)} - currentValue := thinkValueLabel(current) - for i, option := range picker.options { - if option.value == currentValue { - picker.cursor = i - break - } - } - return picker -} - -func (m chatModel) updateThinkPicker(msg tea.KeyMsg) (tea.Model, tea.Cmd) { - switch msg.Type { - case tea.KeyCtrlC, tea.KeyEsc: - m.thinkPicker = nil - m.status = "ready" - case tea.KeyEnter: - return m.selectThinkOption() - case tea.KeyUp: - m.thinkPicker.move(-1) - case tea.KeyDown: - m.thinkPicker.move(1) - } - return m, nil -} - -func (p *chatThinkPicker) move(delta int) { - if p == nil || len(p.options) == 0 || delta == 0 { - return - } - p.cursor = clamp(p.cursor+delta, 0, len(p.options)-1) -} - -func (p *chatThinkPicker) selected() (chatThinkOption, bool) { - if p == nil || len(p.options) == 0 { - return chatThinkOption{}, false - } - return p.options[clamp(p.cursor, 0, len(p.options)-1)], true -} - -func (m chatModel) selectThinkOption() (tea.Model, tea.Cmd) { - option, ok := m.thinkPicker.selected() - if !ok { - return m, nil - } - m.thinkPicker = nil - return m.applyThinkValue(option.value) -} - -func (m *chatModel) handleThinkCommand(value string) (tea.Model, tea.Cmd) { - return m.applyThinkValue(value) -} - -func (m *chatModel) applyThinkValue(value string) (tea.Model, tea.Cmd) { - think, label, err := parseThinkValue(value) - if err != nil { - m.entries = append(m.entries, newChatEntry(chatEntry{role: "error", content: err.Error(), err: err.Error()})) - m.status = "error" - return *m, nil - } - m.opts.Think = think - m.status = "think " + label - return *m, nil -} - -func parseThinkValue(value string) (*api.ThinkValue, string, error) { - switch strings.ToLower(strings.TrimSpace(value)) { - case "", "auto", "default", "unset": - return nil, "auto", nil - case "on", "true", "think", "thinking": - return &api.ThinkValue{Value: true}, "on", nil - case "off", "false", "nothink", "no-think": - return &api.ThinkValue{Value: false}, "off", nil - case "low", "medium", "high", "max": - value = strings.ToLower(strings.TrimSpace(value)) - return &api.ThinkValue{Value: value}, value, nil - default: - return nil, "", fmt.Errorf("Usage: /think [auto|on|off|low|medium|high|max]") - } -} - -func thinkValueLabel(value *api.ThinkValue) string { - if value == nil || value.Value == nil { - return "auto" - } - switch v := value.Value.(type) { - case bool: - if v { - return "on" - } - return "off" - case string: - return strings.ToLower(v) - default: - return "auto" - } -} - -func (m chatModel) renderThinkPicker(width int) string { - picker := m.thinkPicker - if picker == nil { - return "" - } - - var b strings.Builder - b.WriteString(chatPickerTitleStyle.Render("Thinking mode")) - b.WriteString("\n\n") - for i, option := range picker.options { - selected := i == picker.cursor - if selected { - b.WriteString(chatPickerSelectedStyle.Render("› " + option.label)) - } else { - b.WriteString(" ") - b.WriteString(chatPickerTextStyle.Render(option.label)) - } - b.WriteByte('\n') - b.WriteString(chatPickerMetaStyle.Render(" " + option.description)) - b.WriteByte('\n') - if i < len(picker.options)-1 { - b.WriteByte('\n') - } - } - - b.WriteString("\n") - b.WriteString(chatPickerMetaStyle.Render("↑/↓ navigate • enter select • esc cancel")) - return b.String() -} diff --git a/cmd/tui/tui.go b/cmd/tui/tui.go index 998a91e4c66..0e16a0a3bfb 100644 --- a/cmd/tui/tui.go +++ b/cmd/tui/tui.go @@ -46,8 +46,8 @@ type menuItem struct { } var runModelMenuItem = menuItem{ - title: "Chat, Code, & Work", - description: "Chat with models, code, search the web, and delegate real work", + title: "Chat with a model", + description: "Start an interactive chat with a model", isRunModel: true, } diff --git a/cmd/tui/tui_test.go b/cmd/tui/tui_test.go index 728bbda1f26..7063297f5ae 100644 --- a/cmd/tui/tui_test.go +++ b/cmd/tui/tui_test.go @@ -114,8 +114,8 @@ func TestMenuRendersRootLaunchChoices(t *testing.T) { view := menu.View() for _, want := range []string{ - "Chat, Code, & Work", - "Chat with models, code, search the web, and delegate real work", + "Chat with a model", + "Start an interactive chat with a model", "Launch Claude Code", "Launch OpenCode", "Launch Hermes Agent", diff --git a/go.mod b/go.mod index b26cd0ed991..c60eac37024 100644 --- a/go.mod +++ b/go.mod @@ -23,7 +23,6 @@ require ( github.com/agnivade/levenshtein v1.1.1 github.com/charmbracelet/bubbletea v1.3.10 github.com/charmbracelet/lipgloss v1.1.0 - github.com/charmbracelet/x/ansi v0.10.1 github.com/d4l3k/go-bfloat16 v0.0.0-20211005043715-690c3bdd05f1 github.com/dlclark/regexp2 v1.11.5 github.com/emirpasic/gods/v2 v2.0.0-alpha @@ -49,6 +48,7 @@ require ( github.com/buger/jsonparser v1.1.1 // indirect github.com/bytedance/sonic/loader v0.1.1 // indirect github.com/charmbracelet/colorprofile v0.2.3-0.20250311203215-f60798e515dc // indirect + github.com/charmbracelet/x/ansi v0.10.1 // indirect github.com/charmbracelet/x/cellbuf v0.0.13-0.20250311204145-2c3ea96c31dd // indirect github.com/charmbracelet/x/term v0.2.1 // indirect github.com/chewxy/hm v1.0.0 // indirect From 53fed26112817f7c55f664efb9e3f65f06cab7db Mon Sep 17 00:00:00 2001 From: Daniel Hiltgen Date: Fri, 11 Sep 2026 13:25:08 -0700 Subject: [PATCH 15/24] llm: keep gemma3n projector off the CPU (#18376) Gemma3n's MobileNetV5 projector silently produces corrupted image embeddings on the CPU backend - no error, the model just describes the wrong image (reproduced on llama.cpp b10760; gemma4's encoder is fine on CPU). Without this guard the existing partial-offload, limited-VRAM, and OOM-retry fallbacks would pick the CPU projector on exactly the small GPUs where gemma3n lands. --- llm/llama_server.go | 13 ++++++++++ llm/llama_server_test.go | 51 ++++++++++++++++++++++++++++++++++++++++ 2 files changed, 64 insertions(+) diff --git a/llm/llama_server.go b/llm/llama_server.go index 0f40197995e..c3e9172656a 100644 --- a/llm/llama_server.go +++ b/llm/llama_server.go @@ -661,12 +661,22 @@ func appendMMProjArgs(params []string, launch llamaServerLaunchConfig) []string } func (launch llamaServerLaunchConfig) mmprojOffloadDisabled() (bool, string) { + if launch.requiresMMProjGPUOffload() { + return false, "" + } if launch.forceNoMMProjOffload { return true, "startup-oom-retry" } return shouldDisableMMProjOffload(launch.opts, launch.gpus, launch.modelLayers, launch.mmprojMemory) } +func (launch llamaServerLaunchConfig) requiresMMProjGPUOffload() bool { + // Gemma3n's MobileNetV5 projector silently produces corrupted image + // embeddings when it runs on the CPU backend, so keep it on the GPU + // whenever one is in play. + return launch.modelArch == "gemma3n" && launch.opts.NumGPU != 0 && len(launch.gpus) > 0 +} + func shouldDisableMMProjOffload(opts api.Options, gpus []ml.DeviceInfo, modelLayers, mmprojMemory uint64) (bool, string) { if opts.NumGPU == 0 { return true, "cpu-only" @@ -1107,6 +1117,9 @@ func (s *llamaServerRunner) shouldRetryMMProjCPUOffload(err error) bool { if err == nil || s.mmprojOffloadOOMRetried || !IsOutOfMemory(err) || len(s.launch.projectors) == 0 { return false } + if s.launch.requiresMMProjGPUOffload() { + return false + } // llama-server --fit can select a text-layer placement that fits before // mtmd/CLIP allocates the multimodal projector. Retry once with the // projector on CPU so the scheduler can keep the text model placement. diff --git a/llm/llama_server_test.go b/llm/llama_server_test.go index 4d4300b5ee5..256b57dc35b 100644 --- a/llm/llama_server_test.go +++ b/llm/llama_server_test.go @@ -2195,6 +2195,7 @@ func TestAppendMMProjArgs(t *testing.T) { tests := []struct { name string + modelArch string projectors []string opts api.Options gpus []ml.DeviceInfo @@ -2299,12 +2300,23 @@ func TestAppendMMProjArgs(t *testing.T) { retry: true, want: []string{"base", "--mmproj", "model.gguf", "--no-mmproj-offload"}, }, + { + name: "gemma3n keeps projector offload under partial text offload", + modelArch: "gemma3n", + projectors: []string{"model.gguf"}, + opts: partialOpts, + gpus: []ml.DeviceInfo{{DeviceID: ml.DeviceID{Library: "CUDA"}, FreeMemory: 24 << 30}}, + mmprojMemory: 933 << 20, + modelLayers: 81, + want: []string{"base", "--mmproj", "model.gguf"}, + }, } for _, tt := range tests { t.Run(tt.name, func(t *testing.T) { got := appendMMProjArgs([]string{"base"}, llamaServerLaunchConfig{ modelPath: "model.gguf", + modelArch: tt.modelArch, projectors: tt.projectors, mmprojMemory: tt.mmprojMemory, opts: tt.opts, @@ -2319,6 +2331,45 @@ func TestAppendMMProjArgs(t *testing.T) { } } +func TestShouldRetryMMProjCPUOffload(t *testing.T) { + launch := func(modelArch string) llamaServerLaunchConfig { + return llamaServerLaunchConfig{ + modelArch: modelArch, + projectors: []string{"mmproj.gguf"}, + opts: api.DefaultOptions(), + gpus: []ml.DeviceInfo{{DeviceID: ml.DeviceID{Library: "CUDA"}, FreeMemory: 24 << 30}}, + mmprojMemory: 933 << 20, + modelLayers: 81, + } + } + + tests := []struct { + name string + launch llamaServerLaunchConfig + want bool + }{ + { + name: "oom retries with projector on cpu", + launch: launch("qwen3vl"), + want: true, + }, + { + name: "gemma3n oom does not retry with projector on cpu", + launch: launch("gemma3n"), + want: false, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + s := &llamaServerRunner{launch: tt.launch} + if got := s.shouldRetryMMProjCPUOffload(errors.New("out of memory")); got != tt.want { + t.Fatalf("shouldRetryMMProjCPUOffload = %v, want %v", got, tt.want) + } + }) + } +} + func TestMMProjFitTargetExtraEnvs(t *testing.T) { t.Setenv(llamaArgFitTargetEnv, "") _ = os.Unsetenv(llamaArgFitTargetEnv) From b17427b3caaafd4eadffd097fd788c65e60432d4 Mon Sep 17 00:00:00 2001 From: Eva H <63033505+hoyyeva@users.noreply.github.com> Date: Mon, 14 Sep 2026 11:54:46 -0400 Subject: [PATCH 16/24] app: refresh Apps layout and command copy feedback (#18372) --- app/cmd/app/webview.go | 7 + .../src/components/CodexDesktopRow.test.tsx | 265 +++-- app/ui/app/src/components/CodexDesktopRow.tsx | 65 +- .../components/IntegrationConnectButton.tsx | 39 + app/ui/app/src/components/Onboarding.test.tsx | 987 ++++++++++++------ app/ui/app/src/components/Onboarding.tsx | 398 ++++--- .../src/components/OnboardingRoute.test.tsx | 182 ++++ app/ui/app/src/index.css | 43 +- app/ui/app/src/lib/onboarding.ts | 16 +- 9 files changed, 1305 insertions(+), 697 deletions(-) create mode 100644 app/ui/app/src/components/IntegrationConnectButton.tsx create mode 100644 app/ui/app/src/components/OnboardingRoute.test.tsx diff --git a/app/cmd/app/webview.go b/app/cmd/app/webview.go index 0427601a74e..89da905075d 100644 --- a/app/cmd/app/webview.go +++ b/app/cmd/app/webview.go @@ -238,6 +238,13 @@ func (w *Webview) Run(path string) unsafe.Pointer { return } + if runtime.GOOS == "darwin" { + // Keep the current frame through the handoff. SetSize also + // recenters the macOS window and would jump before Apps paints. + setOnboardingWindowStyle(wv.Window(), false) + return + } + width, height := defaultWindowWidth, defaultWindowHeight if w.Store != nil { storedWidth, storedHeight, err := w.Store.WindowSize() diff --git a/app/ui/app/src/components/CodexDesktopRow.test.tsx b/app/ui/app/src/components/CodexDesktopRow.test.tsx index daf32b4e227..81748aed2e3 100644 --- a/app/ui/app/src/components/CodexDesktopRow.test.tsx +++ b/app/ui/app/src/components/CodexDesktopRow.test.tsx @@ -6,7 +6,7 @@ import { notifyManager, } from "@tanstack/react-query"; import { renderToStaticMarkup } from "react-dom/server"; -import { act, create } from "react-test-renderer"; +import { act, create, type ReactTestRenderer } from "react-test-renderer"; import { afterEach, beforeEach, describe, expect, it, vi } from "vitest"; import { CodexConnectedIntro } from "./CodexConnectedIntro"; import { @@ -64,18 +64,14 @@ function status( }; } -describe("CodexDesktopRow", () => { - it("renders a disconnected ChatGPT toggle", () => { - const html = renderToStaticMarkup( - , - ); - - expect(html).toContain(">ChatGPT (Desktop)

"); - expect(html).toContain("Use Ollama models in ChatGPT"); - expect(html).toContain('aria-label="Add Ollama models to ChatGPT"'); - expect(html).toContain('aria-checked="false"'); - }); +function connectionButton(renderer: ReactTestRenderer) { + return renderer.root.find( + (node) => + node.type === "button" && typeof node.props["aria-pressed"] === "boolean", + ); +} +describe("CodexDesktopRow", () => { it("matches Claude's connected copy before the first request", () => { const html = renderToStaticMarkup( { expect(html).not.toContain("Codex + Ollama"); expect(html).not.toContain("3 Ollama models"); expect(html).toContain('aria-label="Remove Ollama models from ChatGPT"'); - expect(html).toContain('aria-checked="true"'); + expect(html).toContain('aria-pressed="true"'); }); it.each([ @@ -119,21 +115,7 @@ describe("CodexDesktopRow", () => { }, ); - it("offers installation when ChatGPT is not installed", () => { - const html = renderToStaticMarkup( - , - ); - - expect(html).toContain("Use Ollama models in ChatGPT"); - expect(html).not.toContain('disabled=""'); - expect(html).toContain('title="Install ChatGPT and add Ollama models"'); - expect(html).toContain("Download & connect"); - }); - - it("matches Claude's download and install progress states", async () => { + it("keeps the connection busy until ChatGPT installation is detected", async () => { const notInstalled = status({ installed: false }); let finishInstall!: (result: "opened") => void; const install = new Promise<"opened">((resolve) => { @@ -162,26 +144,19 @@ describe("CodexDesktopRow", () => { />, ); }); - const toggle = renderer!.root.findByProps({ role: "switch" }); + const toggle = connectionButton(renderer!); + expect(toggle.props["aria-pressed"]).toBe(false); + expect(toggle.props.disabled).toBe(false); await act(async () => { toggle.props.onClick(); await Promise.resolve(); }); - expect(toggle.props["aria-checked"]).toBe(true); + expect(window.installCodexDesktop).toHaveBeenCalledOnce(); + expect(toggle.props["aria-pressed"]).toBe(true); expect(toggle.props["aria-busy"]).toBe(true); expect(toggle.props.disabled).toBe(true); - expect(toggle.props.className).toContain("disabled:cursor-wait"); - expect(renderer!.root.findByProps({ role: "status" }).children).toContain( - "Downloading…", - ); - expect( - renderer!.root.findAll((node) => - node.children.includes( - "Ollama is downloading the ChatGPT installer…", - ), - ), - ).toHaveLength(1); + expect(toggle.findByProps({ role: "status" })).toBeTruthy(); await act(async () => { finishInstall("opened"); @@ -189,25 +164,16 @@ describe("CodexDesktopRow", () => { await Promise.resolve(); }); - expect(toggle.props["aria-checked"]).toBe(true); + expect(toggle.props["aria-pressed"]).toBe(true); expect(toggle.props["aria-busy"]).toBe(true); expect(toggle.props.disabled).toBe(true); - expect(renderer!.root.findByProps({ role: "status" }).children).toContain( - "Finish installing…", - ); - expect( - renderer!.root.findAll((node) => - node.children.includes( - "Finish installing ChatGPT. Ollama will connect it automatically.", - ), - ), - ).toHaveLength(1); + expect(toggle.findByProps({ role: "status" })).toBeTruthy(); } finally { await act(async () => renderer?.unmount()); } }); - it("matches Claude's connecting state", async () => { + it("disables the connection button while connecting ChatGPT", async () => { let finishConnect!: (result: { status: CodexDesktopStatus }) => void; const connect = new Promise<{ status: CodexDesktopStatus }>((resolve) => { finishConnect = resolve; @@ -229,23 +195,16 @@ describe("CodexDesktopRow", () => { />, ); }); - const toggle = renderer!.root.findByProps({ role: "switch" }); + const toggle = connectionButton(renderer!); await act(async () => { toggle.props.onClick(); await Promise.resolve(); }); - expect(toggle.props["aria-checked"]).toBe(true); + expect(toggle.props["aria-pressed"]).toBe(true); expect(toggle.props["aria-busy"]).toBe(true); expect(toggle.props.disabled).toBe(true); - expect(renderer!.root.findByProps({ role: "status" }).children).toContain( - "Connecting…", - ); - expect( - renderer!.root.findAll((node) => - node.children.includes("Connecting ChatGPT to Ollama…"), - ), - ).toHaveLength(1); + expect(toggle.findByProps({ role: "status" })).toBeTruthy(); await act(async () => { finishConnect({ status: status({ connected: true }) }); @@ -304,7 +263,7 @@ describe("CodexDesktopRow", () => { expect( renderer!.root.findByProps({ "aria-label": "Remove Ollama models from ChatGPT", - }).props["aria-checked"], + }).props["aria-pressed"], ).toBe(true); expect(renderer!.root.findByProps({ role: "status" }).children).toContain( "Ollama models added alongside Codex models", @@ -357,7 +316,7 @@ describe("CodexDesktopRow", () => { />, ); }); - const toggle = renderer!.root.findByProps({ role: "switch" }); + const toggle = connectionButton(renderer!); await act(async () => toggle.props.onClick()); await act(async () => vi.advanceTimersByTimeAsync(CODEX_DESKTOP_INSTALL_TIMEOUT_MS - 2000), @@ -369,7 +328,7 @@ describe("CodexDesktopRow", () => { if (!retry) settleOldCheck(); }); expect(toggle.props.disabled).toBe(false); - expect(toggle.props["aria-checked"]).toBe(false); + expect(toggle.props["aria-pressed"]).toBe(false); expect( renderer!.root.findByProps({ role: "alert" }).children, ).toContain("ChatGPT installation wasn’t detected. Try again."); @@ -391,7 +350,7 @@ describe("CodexDesktopRow", () => { await act(async () => vi.advanceTimersByTimeAsync(1000)); expect(connect).toHaveBeenCalledOnce(); expect(toggle.props.disabled).toBe(false); - expect(toggle.props["aria-checked"]).toBe(true); + expect(toggle.props["aria-pressed"]).toBe(true); } } finally { await act(async () => renderer?.unmount()); @@ -426,7 +385,7 @@ describe("CodexDesktopRow", () => { />, ); }); - const toggle = renderer!.root.findByProps({ role: "switch" }); + const toggle = connectionButton(renderer!); await act(async () => toggle.props.onClick()); if (step === "connection") { await act(async () => check.resolve(status())); @@ -491,15 +450,11 @@ describe("CodexDesktopRow", () => { expect(renderer!.root.findAllByType(CodexConnectedIntro)).toHaveLength( 0, ); - expect( - renderer!.root.findByProps({ role: "alert" }).children, - ).toContain( - "ChatGPT is installed. Turn on the switch to restart it with Ollama models.", - ); + expect(renderer!.root.findByProps({ role: "alert" })).toBeTruthy(); expect( renderer!.root.findByProps({ "aria-label": "Add Ollama models to ChatGPT", - }).props["aria-checked"], + }).props["aria-pressed"], ).toBe(false); } finally { await act(async () => renderer?.unmount()); @@ -536,7 +491,7 @@ describe("CodexDesktopRow", () => { }); expect(getStatus).not.toHaveBeenCalled(); - expect(toggle.props["aria-checked"]).toBe(false); + expect(toggle.props["aria-pressed"]).toBe(false); expect(toggle.props.disabled).toBe(false); expect(renderer!.root.findAllByProps({ role: "alert" })).toHaveLength(0); } finally { @@ -635,7 +590,7 @@ describe("CodexDesktopRow", () => { expect( renderer!.root.findByProps({ "aria-label": "Remove Ollama models from ChatGPT", - }).props["aria-checked"], + }).props["aria-pressed"], ).toBe(true); } finally { await act(async () => renderer?.unmount()); @@ -692,12 +647,62 @@ describe("CodexDesktopRow", () => { expect(renderer!.root.findByProps({ role: "alert" }).children).toContain( "quit ChatGPT: timed out waiting for ChatGPT to exit", ); - expect(toggle.props["aria-checked"]).toBe(false); + expect(toggle.props["aria-pressed"]).toBe(false); } finally { await act(async () => renderer?.unmount()); } }); + it.each([false, true])( + "honors the native ChatGPT disconnect confirmation: %s", + async (confirmed) => { + const connected = status({ connected: true, running: true }); + const confirm = vi.fn(() => confirmed); + const disconnect = vi + .fn() + .mockResolvedValueOnce({ + status: connected, + restartConfirmationRequired: true, + }) + .mockResolvedValue({ status: status({ running: true }) }); + vi.stubGlobal("IS_REACT_ACT_ENVIRONMENT", true); + vi.stubGlobal("window", { + addEventListener: vi.fn(), + removeEventListener: vi.fn(), + setCodexDesktopConnected: disconnect, + confirm, + }); + let renderer; + try { + await act(async () => { + renderer = create( + , + ); + }); + const toggle = connectionButton(renderer!); + await act(async () => { + toggle.props.onClick(); + toggle.props.onClick(); + }); + expect(confirm).toHaveBeenCalledOnce(); + expect(disconnect).toHaveBeenNthCalledWith(1, false, false); + if (confirmed) { + expect(disconnect).toHaveBeenCalledTimes(2); + expect(disconnect).toHaveBeenLastCalledWith(false, true); + } else { + expect(disconnect).toHaveBeenCalledOnce(); + expect(toggle.props.disabled).toBe(false); + expect(toggle.props["aria-pressed"]).toBe(true); + } + } finally { + await act(async () => renderer?.unmount()); + } + }, + ); + it("allows the normal profile to be restored if ChatGPT is removed", async () => { const html = renderToStaticMarkup( { ); }); await act(async () => { - renderer!.root.findByProps({ role: "switch" }).props.onClick(); + connectionButton(renderer!).props.onClick(); }); expect(connect).not.toHaveBeenCalled(); expect(save).not.toHaveBeenCalled(); expect(confirm).toHaveBeenCalledTimes(running ? 1 : 0); - const toggle = renderer!.root.findByProps({ role: "switch" }); - expect(toggle.props["aria-checked"]).toBe(true); + const toggle = connectionButton(renderer!); + expect(toggle.props["aria-pressed"]).toBe(true); expect(toggle.props.disabled).toBe(true); const intro = renderer!.root.findByType(CodexConnectedIntro); await act(async () => { @@ -791,7 +796,7 @@ describe("ChatGPT first connection intro", () => { expect(confirm).toHaveBeenCalledTimes(running ? 1 : 0); expect(save).toHaveBeenCalledOnce(); expect(toggle.props.disabled).toBe(false); - expect(toggle.props["aria-checked"]).toBe(true); + expect(toggle.props["aria-pressed"]).toBe(true); expect(renderer!.root.findAllByType(CodexConnectedIntro)).toHaveLength( 0, ); @@ -829,7 +834,7 @@ describe("ChatGPT first connection intro", () => { ); }); await act(async () => { - renderer!.root.findByProps({ role: "switch" }).props.onClick(); + connectionButton(renderer!).props.onClick(); }); expect(connect).not.toHaveBeenCalled(); expect(renderer!.root.findAllByType(CodexConnectedIntro)).toHaveLength(1); @@ -862,10 +867,10 @@ it("leaves the Apps page usable when the initial restart is cancelled", async () />, ); }); - const toggle = renderer!.root.findByProps({ role: "switch" }); + const toggle = connectionButton(renderer!); await act(async () => toggle.props.onClick()); expect(renderer!.root.findAllByType(CodexConnectedIntro)).toHaveLength(0); - expect(toggle.props["aria-checked"]).toBe(false); + expect(toggle.props["aria-pressed"]).toBe(false); expect(toggle.props.disabled).toBe(false); expect(connect).not.toHaveBeenCalled(); expect(save).not.toHaveBeenCalled(); @@ -903,18 +908,18 @@ it.each(["failed", "rejected", "save failed"])( />, ); }); - const toggle = renderer!.root.findByProps({ role: "switch" }); + const toggle = connectionButton(renderer!); await act(async () => { toggle.props.onClick(); }); - expect(toggle.props["aria-checked"]).toBe(true); + expect(toggle.props["aria-pressed"]).toBe(true); expect(toggle.props.disabled).toBe(true); const intro = renderer!.root.findByType(CodexConnectedIntro); await act(async () => { intro.props.onDone(); }); expect(renderer!.root.findAllByType(CodexConnectedIntro)).toHaveLength(0); - expect(toggle.props["aria-checked"]).toBe(outcome === "save failed"); + expect(toggle.props["aria-pressed"]).toBe(outcome === "save failed"); expect(toggle.props.disabled).toBe(false); expect(renderer!.root.findByProps({ role: "alert" }).children).toContain( outcome === "failed" @@ -958,7 +963,7 @@ it("dismisses before launch and prevents duplicate Continue requests", async () />, ); }); - const toggle = renderer!.root.findByProps({ role: "switch" }); + const toggle = connectionButton(renderer!); await act(async () => toggle.props.onClick()); const intro = renderer!.root.findByType(CodexConnectedIntro); await act(async () => { @@ -1010,18 +1015,16 @@ it.each(["cancelled", "status failed"])( />, ); }); - await act(async () => - renderer!.root.findByProps({ role: "switch" }).props.onClick(), - ); + await act(async () => connectionButton(renderer!).props.onClick()); await act(async () => renderer!.root.findByType(CodexConnectedIntro).props.onDone(), ); expect(connect).not.toHaveBeenCalled(); expect(save).not.toHaveBeenCalled(); expect(renderer!.root.findAllByType(CodexConnectedIntro)).toHaveLength(0); - const toggle = renderer!.root.findByProps({ role: "switch" }); + const toggle = connectionButton(renderer!); expect(toggle.props.disabled).toBe(false); - expect(toggle.props["aria-checked"]).toBe(false); + expect(toggle.props["aria-pressed"]).toBe(false); } finally { await act(async () => renderer?.unmount()); } @@ -1059,9 +1062,7 @@ it("does not open a restart prompt after leaving the Apps page", async () => { , ); }); - await act(async () => - renderer!.root.findByProps({ role: "switch" }).props.onClick(), - ); + await act(async () => connectionButton(renderer!).props.onClick()); await act(async () => renderer!.unmount()); await act(async () => action.resolve({ @@ -1105,9 +1106,7 @@ it("keeps the latest status when focus refreshes complete out of order", async ( }); await act(async () => newer.resolve(status({ connected: true }))); await act(async () => older.resolve(status())); - expect( - renderer!.root.findByProps({ role: "switch" }).props["aria-checked"], - ).toBe(true); + expect(connectionButton(renderer!).props["aria-pressed"]).toBe(true); } finally { await act(async () => renderer?.unmount()); } @@ -1134,7 +1133,7 @@ it("does not start overlapping installers before the switch rerenders", async () />, ); }); - const toggle = renderer!.root.findByProps({ role: "switch" }); + const toggle = connectionButton(renderer!); await act(async () => { toggle.props.onClick(); toggle.props.onClick(); @@ -1142,7 +1141,7 @@ it("does not start overlapping installers before the switch rerenders", async () expect(install).toHaveBeenCalledOnce(); await act(async () => result.resolve("cancelled")); expect(toggle.props.disabled).toBe(false); - expect(toggle.props["aria-checked"]).toBe(false); + expect(toggle.props["aria-pressed"]).toBe(false); } finally { await act(async () => renderer?.unmount()); } @@ -1179,7 +1178,7 @@ it("does not retry acknowledgment during a disconnect, or reconnect to save afte />, ); }); - const toggle = renderer!.root.findByProps({ role: "switch" }); + const toggle = connectionButton(renderer!); await act(async () => toggle.props.onClick()); await act(async () => renderer!.root.findByType(CodexConnectedIntro).props.onDone(), @@ -1198,7 +1197,7 @@ it("does not retry acknowledgment during a disconnect, or reconnect to save afte await act(async () => retry.props.onClick()); expect(connect).toHaveBeenCalledTimes(2); expect(save).toHaveBeenCalledTimes(2); - expect(toggle.props["aria-checked"]).toBe(false); + expect(toggle.props["aria-pressed"]).toBe(false); expect( renderer!.root.findAllByProps({ "aria-label": "Retry saving progress" }), ).toHaveLength(0); @@ -1239,9 +1238,7 @@ it.each(["returned", "rejected"])( />, ); }); - await act(async () => - renderer!.root.findByProps({ role: "switch" }).props.onClick(), - ); + await act(async () => connectionButton(renderer!).props.onClick()); await act(async () => renderer!.root.findByType(CodexConnectedIntro).props.onDone(), ); @@ -1257,7 +1254,7 @@ it.each(["returned", "rejected"])( const retryButton = renderer!.root.findByProps({ "aria-label": "Retry saving progress", }); - const toggle = renderer!.root.findByProps({ role: "switch" }); + const toggle = connectionButton(renderer!); await act(async () => { retryButton.props.onClick(); retryButton.props.onClick(); @@ -1267,7 +1264,7 @@ it.each(["returned", "rejected"])( expect(save).toHaveBeenCalledTimes(2); expect(retryButton.props.disabled).toBe(true); expect(toggle.props.disabled).toBe(true); - expect(toggle.props["aria-checked"]).toBe(true); + expect(toggle.props["aria-pressed"]).toBe(true); await act(async () => retry.resolve("")); expect( renderer!.root.findAllByProps({ @@ -1276,7 +1273,7 @@ it.each(["returned", "rejected"])( ).toHaveLength(0); expect(renderer!.root.findAllByProps({ role: "alert" })).toHaveLength(0); expect(toggle.props.disabled).toBe(false); - expect(toggle.props["aria-checked"]).toBe(true); + expect(toggle.props["aria-pressed"]).toBe(true); expect(connect).toHaveBeenCalledOnce(); expect(save).toHaveBeenCalledTimes(2); } finally { @@ -1323,9 +1320,7 @@ it.each([ await act(async () => { renderer = create(); }); - await act(async () => - renderer!.root.findByProps({ role: "switch" }).props.onClick(), - ); + await act(async () => connectionButton(renderer!).props.onClick()); await act(async () => renderer!.root.findByType(CodexConnectedIntro).props.onDone(), ); @@ -1337,7 +1332,7 @@ it.each([ await act(async () => { renderer = create(); }); - const toggle = renderer!.root.findByProps({ role: "switch" }); + const toggle = connectionButton(renderer!); if (timing === "after returning") { expect(toggle.props.disabled).toBe(true); await act(async () => toggle.props.onClick()); @@ -1361,9 +1356,7 @@ it.each([ await act(async () => { renderer = create(); }); - expect( - renderer!.root.findByProps({ role: "switch" }).props.disabled, - ).toBe(true); + expect(connectionButton(renderer!).props.disabled).toBe(true); await act(async () => renderer!.root .findByProps({ "aria-label": "Retry saving progress" }) @@ -1377,12 +1370,8 @@ it.each([ }), ).toHaveLength(0); expect(renderer!.root.findAllByProps({ role: "alert" })).toHaveLength(0); - expect( - renderer!.root.findByProps({ role: "switch" }).props.disabled, - ).toBe(false); - expect( - renderer!.root.findByProps({ role: "switch" }).props["aria-checked"], - ).toBe(true); + expect(connectionButton(renderer!).props.disabled).toBe(false); + expect(connectionButton(renderer!).props["aria-pressed"]).toBe(true); expect(connect).toHaveBeenCalledOnce(); expect(save).toHaveBeenCalledTimes(2); } finally { @@ -1415,9 +1404,7 @@ it("refreshes connection status when Retry overtakes the returning page's status await act(async () => { renderer = create(); }); - await act(async () => - renderer!.root.findByProps({ role: "switch" }).props.onClick(), - ); + await act(async () => connectionButton(renderer!).props.onClick()); await act(async () => renderer!.root.findByType(CodexConnectedIntro).props.onDone(), ); @@ -1435,12 +1422,8 @@ it("refreshes connection status when Retry overtakes the returning page's status ); await act(async () => stale.resolve(connectedStatus)); await act(async () => retrySave.resolve("")); - expect( - renderer!.root.findByProps({ role: "switch" }).props["aria-checked"], - ).toBe(true); - expect(renderer!.root.findByProps({ role: "switch" }).props.disabled).toBe( - false, - ); + expect(connectionButton(renderer!).props["aria-pressed"]).toBe(true); + expect(connectionButton(renderer!).props.disabled).toBe(false); expect(connect).toHaveBeenCalledOnce(); expect(save).toHaveBeenCalledTimes(2); } finally { @@ -1468,9 +1451,7 @@ it("observes a successful pending save after returning without saving again", as await act(async () => { renderer = create(); }); - await act(async () => - renderer!.root.findByProps({ role: "switch" }).props.onClick(), - ); + await act(async () => connectionButton(renderer!).props.onClick()); await act(async () => renderer!.root.findByType(CodexConnectedIntro).props.onDone(), ); @@ -1479,17 +1460,13 @@ it("observes a successful pending save after returning without saving again", as await act(async () => { renderer = create(); }); - expect(renderer!.root.findByProps({ role: "switch" }).props.disabled).toBe( - true, - ); + expect(connectionButton(renderer!).props.disabled).toBe(true); await act(async () => pendingSave.resolve("")); expect( renderer!.root.findAllByProps({ "aria-label": "Retry saving progress" }), ).toHaveLength(0); expect(renderer!.root.findAllByType(CodexConnectedIntro)).toHaveLength(0); - expect(renderer!.root.findByProps({ role: "switch" }).props.disabled).toBe( - false, - ); + expect(connectionButton(renderer!).props.disabled).toBe(false); expect(connect).toHaveBeenCalledOnce(); expect(save).toHaveBeenCalledOnce(); } finally { @@ -1562,9 +1539,7 @@ it.each(["resolved", "rejected"])( ); }); await act(async () => onFocus?.()); - await act(async () => - renderer!.root.findByProps({ role: "switch" }).props.onClick(), - ); + await act(async () => connectionButton(renderer!).props.onClick()); await act(async () => renderer!.root.findByType(CodexConnectedIntro).props.onDone(), ); @@ -1572,9 +1547,7 @@ it.each(["resolved", "rejected"])( if (outcome === "resolved") stale.resolve(firstUseStatus); else stale.reject(new Error("stale status failure")); }); - expect( - renderer!.root.findByProps({ role: "switch" }).props["aria-checked"], - ).toBe(true); + expect(connectionButton(renderer!).props["aria-pressed"]).toBe(true); expect(renderer!.root.findByProps({ role: "alert" }).children).toContain( "Ollama couldn’t save your progress. Please try again.", ); diff --git a/app/ui/app/src/components/CodexDesktopRow.tsx b/app/ui/app/src/components/CodexDesktopRow.tsx index 12e8071ad3c..034ab0c0872 100644 --- a/app/ui/app/src/components/CodexDesktopRow.tsx +++ b/app/ui/app/src/components/CodexDesktopRow.tsx @@ -5,7 +5,8 @@ import type { CodexDesktopActionResult, CodexDesktopStatus, } from "@/types/webview"; -import { ArrowPathIcon, CommandLineIcon } from "@heroicons/react/24/outline"; +import { CommandLineIcon } from "@heroicons/react/24/outline"; +import { IntegrationConnectButton } from "@/components/IntegrationConnectButton"; import { useMutation, useMutationState, @@ -254,7 +255,7 @@ export function CodexDesktopRow({ if (next.running) { setError( - "ChatGPT is installed. Turn on the switch to restart it with Ollama models.", + "ChatGPT is installed. Click Connect to restart it with Ollama models.", ); return; } @@ -270,7 +271,7 @@ export function CodexDesktopRow({ setStatus(result.status); if (result.restartConfirmationRequired) { setError( - "ChatGPT is installed. Turn on the switch to restart it with Ollama models.", + "ChatGPT is installed. Click Connect to restart it with Ollama models.", ); } else if (result.error || !result.status.connected) { setError( @@ -316,8 +317,7 @@ export function CodexDesktopRow({ phase === "waiting-for-install" || phase === "connecting"; const progress = connectionProgress[savingAcknowledgment ? "saving" : phase]; - const statusLabel = - progress?.label ?? (!connected && !installed ? "Download & connect" : null); + const statusLabel = progress?.label ?? null; const actionError = error ?? (acknowledgmentFailed @@ -327,7 +327,12 @@ export function CodexDesktopRow({ actionError ?? notice ?? progress?.description ?? - codexDesktopDescription(status, integration.description); + codexDesktopDescription( + status, + installed + ? "Use Ollama models in Codex mode in ChatGPT." + : "We’ll download ChatGPT and connect it to Ollama.", + ); const saveAcknowledgment = async (): Promise => { if (queryClient.isMutating({ mutationKey: acknowledgmentKey })) @@ -413,9 +418,9 @@ export function CodexDesktopRow({ let result: CodexDesktopActionResult = await window.setCodexDesktopConnected(enabled, restartConfirmed); + if (!mounted.current) return; setStatus(result.status); if (result.restartConfirmationRequired) { - if (!mounted.current) return; // Keep focus-driven status refreshes from discarding this operation // while the native confirmation dialog temporarily owns focus. if ( @@ -423,11 +428,13 @@ export function CodexDesktopRow({ enabled ? "Restart ChatGPT to add Ollama models? Any running task will stop." : "Restart ChatGPT to remove Ollama models? Any running task will stop.", - ) + ) || + !mounted.current ) { return; } result = await window.setCodexDesktopConnected(enabled, true); + if (!mounted.current) return; setStatus(result.status); } @@ -467,16 +474,19 @@ export function CodexDesktopRow({ }; return ( -
-
+
+
-

- ChatGPT (Desktop) +

+ ChatGPT

{description}

@@ -494,22 +504,11 @@ export function CodexDesktopRow({ Retry )} - {statusLabel && ( - - {pending && } - {statusLabel} - - )} - + />
{showIntro && ( void toggleConnection(true)} /> diff --git a/app/ui/app/src/components/IntegrationConnectButton.tsx b/app/ui/app/src/components/IntegrationConnectButton.tsx new file mode 100644 index 00000000000..a83b5dba8d0 --- /dev/null +++ b/app/ui/app/src/components/IntegrationConnectButton.tsx @@ -0,0 +1,39 @@ +import { ArrowPathIcon } from "@heroicons/react/24/outline"; +import type { ComponentProps } from "react"; + +export function IntegrationConnectButton({ + connected, + busy, + progress, + label, + ...props +}: Pick, "onClick" | "disabled" | "title"> & { + connected: boolean; + label: string; + busy: boolean; + progress?: string | null; +}) { + return ( + + ); +} diff --git a/app/ui/app/src/components/Onboarding.test.tsx b/app/ui/app/src/components/Onboarding.test.tsx index 17f4146114d..2f35aa93d9c 100644 --- a/app/ui/app/src/components/Onboarding.test.tsx +++ b/app/ui/app/src/components/Onboarding.test.tsx @@ -1,3 +1,4 @@ +import { StrictMode } from "react"; import { renderToStaticMarkup } from "react-dom/server"; import { QueryClient } from "@tanstack/react-query"; import { act, create, type ReactTestRenderer } from "react-test-renderer"; @@ -18,13 +19,27 @@ import { isClaudeConnectionComplete, scheduleClaudeInstallTimeout, } from "@/lib/claudeDesktop"; -import { isWindowsPlatform } from "@/lib/platform"; import { authenticationTimeoutAction, - nextOnboardingStep, onboardingConnectUrl, } from "@/lib/onboarding"; import type { IntegrationStatuses } from "@/api"; +import * as clipboard from "@/utils/clipboard"; + +// Transition timing is checked in the browser; these tests cover copy-notice +// lifetimes with the real component logic. +vi.mock("@headlessui/react", async (importOriginal) => { + const original = await importOriginal(); + return Object.assign({}, original, { + Transition: ({ + show, + children, + }: { + show: boolean; + children: React.ReactNode; + }) => (show ?
{children}
: null), + }); +}); let queryClient: QueryClient; beforeEach(() => { @@ -44,6 +59,16 @@ vi.mock("@tanstack/react-query", async (importOriginal) => { }); }); +function claudeConnectionButton(renderer: ReactTestRenderer) { + return renderer.root + .findByProps({ id: "integration-claude-desktop" }) + .find( + (node) => + node.type === "button" && + typeof node.props["aria-pressed"] === "boolean", + ); +} + describe("Onboarding", () => { it("explains what Ollama is before asking the user to choose a path", () => { const html = renderToStaticMarkup(); @@ -81,59 +106,6 @@ describe("Onboarding", () => { } }); - it("hides the Claude and ChatGPT desktop integrations on Windows", () => { - vi.stubGlobal("window", { - OLLAMA_PLATFORM: "windows", - innerHeight: 660, - }); - vi.stubGlobal("navigator", { platform: "MacIntel" }); - try { - expect(isWindowsPlatform()).toBe(true); - const html = renderToStaticMarkup( - , - ); - - expect(html).not.toContain('id="desktop-heading"'); - expect(html).not.toContain("Use Ollama models in Claude Desktop"); - expect(html).not.toContain("Use Ollama models in ChatGPT"); - expect(html).not.toContain("ollama launch chatgpt"); - expect(html).toContain('id="terminal-heading"'); - expect(html).toContain("ollama launch claude"); - } finally { - vi.unstubAllGlobals(); - } - }); - - it("shows the account choice only to signed-out users", () => { - expect(nextOnboardingStep("intro", "continue", false)).toBe("welcome"); - expect(nextOnboardingStep("intro", "continue", true)).toBe("apps"); - expect(nextOnboardingStep("welcome", "authenticated", true)).toBe("apps"); - expect(nextOnboardingStep("apps", "continue", true)).toBe("apps"); - expect(nextOnboardingStep("welcome", "local", false)).toBe("run"); - }); - it("lets an in-flight authentication check finish before timing out", () => { expect(authenticationTimeoutAction(false, true)).toBe("defer"); expect(authenticationTimeoutAction(false, false)).toBe("fail"); @@ -187,7 +159,7 @@ describe("Onboarding", () => { } }); - it("keeps the Claude switch on and busy through installer detection", async () => { + it("keeps the Claude connection busy through installer detection", async () => { vi.stubGlobal("IS_REACT_ACT_ENVIRONMENT", true); const disconnectedStatus = { @@ -242,10 +214,9 @@ describe("Onboarding", () => { await Promise.resolve(); }); - const claudeSwitch = () => - renderer!.root - .findAllByProps({ role: "switch" }) - .find((node) => String(node.props["aria-label"]).endsWith("Claude"))!; + const claudeSwitch = () => claudeConnectionButton(renderer!); + expect(claudeSwitch().props["aria-pressed"]).toBe(false); + expect(claudeSwitch().props.disabled).toBe(false); let clickResult!: Promise; await act(async () => { clickResult = claudeSwitch().props.onClick(); @@ -253,22 +224,11 @@ describe("Onboarding", () => { await Promise.resolve(); }); - expect(claudeSwitch().props["aria-checked"]).toBe(true); + expect(window.installClaudeDesktop).toHaveBeenCalledOnce(); + expect(claudeSwitch().props["aria-pressed"]).toBe(true); expect(claudeSwitch().props["aria-busy"]).toBe(true); expect(claudeSwitch().props.disabled).toBe(true); - expect(claudeSwitch().props.className).toContain("disabled:opacity-50"); - expect( - renderer.root - .findAllByProps({ role: "status" }) - .some((node) => node.children.includes("Downloading…")), - ).toBe(true); - expect( - renderer.root.findAll( - (node) => - typeof node.props.className === "string" && - node.props.className.includes("animate-spin"), - ), - ).not.toHaveLength(0); + expect(claudeSwitch().findByProps({ role: "status" })).toBeTruthy(); await act(async () => { finishInstall("opened"); @@ -276,22 +236,10 @@ describe("Onboarding", () => { await Promise.resolve(); }); - expect(claudeSwitch().props["aria-checked"]).toBe(true); + expect(claudeSwitch().props["aria-pressed"]).toBe(true); expect(claudeSwitch().props["aria-busy"]).toBe(true); expect(claudeSwitch().props.disabled).toBe(true); - expect(claudeSwitch().props.className).toContain("disabled:opacity-50"); - expect( - renderer.root - .findAllByProps({ role: "status" }) - .some((node) => node.children.includes("Finish installing…")), - ).toBe(true); - expect( - renderer.root.findAll( - (node) => - typeof node.props.className === "string" && - node.props.className.includes("animate-spin"), - ), - ).not.toHaveLength(0); + expect(claudeSwitch().findByProps({ role: "status" })).toBeTruthy(); } finally { if (renderer) { act(() => renderer?.unmount()); @@ -368,11 +316,8 @@ describe("Onboarding", () => { await Promise.resolve(); }); - const claudeSwitch = () => - renderer!.root - .findAllByProps({ role: "switch" }) - .find((node) => String(node.props["aria-label"]).endsWith("Claude"))!; - expect(claudeSwitch().props["aria-checked"]).toBe(false); + const claudeSwitch = () => claudeConnectionButton(renderer!); + expect(claudeSwitch().props["aria-pressed"]).toBe(false); expect(claudeSwitch().props["aria-busy"]).toBeUndefined(); expect(claudeSwitch().props.disabled).toBe(false); @@ -384,7 +329,7 @@ describe("Onboarding", () => { }); expect(setClaudeConnected).toHaveBeenCalledWith(true, false); - expect(claudeSwitch().props["aria-checked"]).toBe(true); + expect(claudeSwitch().props["aria-pressed"]).toBe(true); expect(claudeSwitch().props["aria-busy"]).toBe(true); expect(claudeSwitch().props.disabled).toBe(true); @@ -393,7 +338,7 @@ describe("Onboarding", () => { await clickResult; }); - expect(claudeSwitch().props["aria-checked"]).toBe(false); + expect(claudeSwitch().props["aria-pressed"]).toBe(false); expect(claudeSwitch().props["aria-busy"]).toBeUndefined(); expect(claudeSwitch().props.disabled).toBe(false); expect( @@ -410,7 +355,7 @@ describe("Onboarding", () => { }); expect(getClaudeStatus).toHaveBeenCalledOnce(); - expect(claudeSwitch().props["aria-checked"]).toBe(true); + expect(claudeSwitch().props["aria-pressed"]).toBe(true); expect(claudeSwitch().props["aria-busy"]).toBeUndefined(); expect(claudeSwitch().props.disabled).toBe(false); expect( @@ -504,224 +449,11 @@ describe("Onboarding", () => { expect(html).not.toContain("Sign up"); }); - it("groups disconnected Claude with a scrollable terminal list", () => { - const integrations: IntegrationStatuses = [ - { - id: "claude-desktop", - name: "Claude Code (Desktop)", - description: "Use Ollama models in Claude Desktop", - installed: true, - action: "connect", - }, - { - id: "claude", - name: "Claude Code", - description: "Anthropic's coding tool with subagents", - installed: true, - action: "copy", - command: "ollama launch claude", - }, - { - id: "codex", - name: "Codex CLI", - description: "OpenAI's open-source coding agent", - installed: true, - action: "copy", - command: "ollama launch codex", - }, - { - id: "openclaw", - name: "OpenClaw", - description: "Personal AI with 100+ skills", - installed: true, - action: "copy", - command: "ollama launch openclaw", - }, - { - id: "opencode", - name: "OpenCode", - description: "Anomaly's open-source coding agent", - installed: false, - action: "copy", - command: "ollama launch opencode", - }, - { - id: "droid", - name: "Droid", - description: "AI software engineering agent", - installed: false, - action: "copy", - command: "ollama launch droid", - }, - { - id: "dsh", - name: "DeepSeek Harness", - description: "DeepSeek's open-source agent harness", - installed: false, - action: "copy", - command: "ollama launch dsh", - }, - { - id: "cline", - name: "Cline", - description: "Autonomous coding agent", - installed: false, - action: "copy", - command: "ollama launch cline", - }, - { - id: "terminal", - name: "Terminal", - description: "Run local models from your terminal", - action: "copy", - command: "ollama", - }, - ]; - const html = renderToStaticMarkup( - , - ); - - expect(html).not.toContain( - "Connect Claude, or copy a command to run in your terminal.", - ); - expect(html).toContain("Claude Code (Desktop)"); - expect(html).toContain("Use Ollama models in Claude Desktop"); - expect(html).toContain("Claude Code"); - expect(html).toContain("Codex CLI"); - expect(html).not.toContain("Search apps"); - expect(html).not.toContain('type="search"'); - expect(html).toContain("Desktop"); - expect(html).toContain('id="desktop-heading"'); - expect(html).toContain('id="terminal-heading"'); - expect(html).not.toContain("Ready to launch"); - expect(html).not.toContain('id="claude-apps-heading"'); - expect(html.indexOf("Desktop")).toBeLessThan( - html.indexOf("Use Ollama models in Claude Desktop"), - ); - expect(html).not.toContain(">Command"); - expect(html).toContain("ollama launch claude"); - expect(html).not.toContain("Installed"); - expect(html).toContain("Use Ollama models in ChatGPT"); - expect(html).toContain('aria-label="Connect Claude"'); - expect(html).toContain('role="switch"'); - expect(html).toContain('aria-checked="false"'); - expect(html).not.toContain("Inactive"); - expect(html).toContain("Download & connect"); - expect(html).not.toContain("Active"); - expect(html).toContain("bg-transparent"); - expect(html).toContain('aria-label="Copy OpenCode command"'); - expect(html).toContain('aria-label="Copy Terminal command"'); - expect(html).not.toContain(">Copy command"); - expect(html).toContain("ChatGPT (Desktop)"); - expect(html).toContain("OpenCode"); - expect(html).toContain("Terminal"); - expect(html).toContain("overflow-y-auto"); - expect(html).not.toContain('aria-label="Show more apps"'); - expect(html).not.toContain("aria-expanded"); - expect(html).not.toContain("grid-rows-[0fr]"); - expect(html).not.toContain("inert"); - expect(html).toContain("/launch-icons/claude.svg"); - expect(html).toContain("/launch-icons/claude-code.svg"); - expect(html).toContain("/launch-icons/codex-color.svg"); - expect(html).toMatch( - /src="\/launch-icons\/cline\.svg"[^>]*class="[^"]*dark:invert/, - ); - expect(html).toContain("/launch-icons/deepseek-harness.svg"); - expect(html).not.toMatch( - /src="\/launch-icons\/deepseek-harness\.svg"[^>]*class="[^"]*dark:invert/, - ); - expect(html).not.toContain(" { - const html = renderToStaticMarkup( - , - ); - - expect(html.indexOf("Use Ollama models in Claude Desktop")).toBeLessThan( - html.indexOf(">ChatGPT (Desktop)

"), - ); - expect(html).toContain('aria-label="Add Ollama models to ChatGPT"'); - expect(html).not.toContain('aria-label="Copy ChatGPT command"'); - expect(html).toContain('aria-label="Copy Codex CLI command"'); - }); - - it("keeps connected Claude in Desktop without an idle status", () => { + it("offers ChatGPT when the catalog has no desktop metadata", () => { const html = renderToStaticMarkup( - , + , ); - - expect(html).toContain('id="desktop-heading"'); - expect(html).not.toContain('id="claude-apps-heading"'); - expect(html).not.toContain("Ready to launch"); - expect(html).not.toContain("Active"); - expect(html).not.toContain("Inactive"); - expect(html).toContain('aria-checked="true"'); - expect(html).toContain('aria-label="Disconnect Claude"'); - expect(html).toContain("Connected to Ollama · 12 requests this session"); + expect(html).toContain('id="integration-chatgpt"'); }); it("shows initial Claude recovery guidance without error styling", () => { @@ -757,7 +489,7 @@ describe("Onboarding", () => { ); expect(html).toContain('role="alert"'); expect(html).not.toContain("text-red"); - expect(html).toContain('aria-checked="true"'); + expect(html).toContain('aria-pressed="true"'); expect(html).toContain('aria-label="Disconnect Claude"'); }); @@ -810,34 +542,6 @@ describe("Onboarding", () => { expect(html).not.toContain("Built-in defaults"); }); - it("keeps Claude available without a separate not-installed group", () => { - const html = renderToStaticMarkup( - , - ); - - expect(html).toContain("Use Ollama models in Claude Desktop"); - expect(html).toContain('aria-label="Connect Claude"'); - expect(html).toContain("Download & connect"); - expect(html).not.toContain("Inactive"); - const claudeButton = html.match( - /]*aria-label="Connect Claude"[^>]*>/, - )?.[0]; - expect(claudeButton).toBeDefined(); - expect(claudeButton).not.toContain('disabled=""'); - }); - it("uses branded icons for the remaining launcher integrations", () => { const html = renderToStaticMarkup( { expect(html).toContain("Try again"); }); }); + +function appsIntegrations(claudeInstalled: boolean): IntegrationStatuses { + const launcher = (id: string, name: string) => ({ + id, + name, + description: `${name} description`, + installed: false, + command: `ollama launch ${id}`, + }); + return [ + { + id: "claude-desktop", + name: "Claude", + description: "Use Ollama models in Claude Desktop", + installed: claudeInstalled, + }, + launcher("claude", "Claude Code"), + launcher("codex", "Codex CLI"), + launcher("opencode", "OpenCode"), + launcher("pi", "Pi"), + launcher("hermes", "Hermes Agent"), + { + id: "terminal", + name: "Terminal", + description: "Run local models from your terminal", + command: "ollama", + }, + ]; +} + +function onboardingProps(onOpenApps: () => Promise) { + return { + isAuthenticated: true, + isSigningIn: false, + signInError: null, + completionError: null, + onOpenApps, + onSignIn: vi.fn(), + onSignUp: vi.fn(), + onRetryCompletion: vi.fn(), + onUseLocal: vi.fn(), + }; +} + +const DISCONNECTED_CLAUDE = { + supported: true, + used: false, + installed: true, + configured: false, + connected: false, + running: false, + startFailed: false, + portConflict: false, +}; + +async function settle(ticks = 8) { + for (let i = 0; i < ticks; i++) { + await Promise.resolve(); + } +} + +function stubOnboardingWindow(platform = "darwin") { + vi.stubGlobal("IS_REACT_ACT_ENVIRONMENT", true); + vi.stubGlobal("navigator", { platform: "MacIntel" }); + vi.stubGlobal("window", { + OLLAMA_PLATFORM: platform, + setOnboardingWindow: vi.fn(), + }); +} + +describe("Onboarding handoff", () => { + it("preserves local setup without opening Apps", async () => { + stubOnboardingWindow(); + const props = { ...onboardingProps(vi.fn()), isAuthenticated: false }; + let renderer: ReactTestRenderer | undefined; + try { + await act(async () => { + renderer = create(); + }); + await act(async () => { + renderer!.root.findByType(IntroScreen).props.onContinue(); + }); + expect(props.onOpenApps).not.toHaveBeenCalled(); + await act(async () => { + renderer!.root.findByType(WelcomeScreen).props.onLocal(); + }); + expect(renderer!.root.findByType(RunOllamaScreen)).toBeTruthy(); + expect(props.onUseLocal).toHaveBeenCalledOnce(); + expect(props.onOpenApps).not.toHaveBeenCalled(); + } finally { + if (renderer) act(() => renderer?.unmount()); + vi.unstubAllGlobals(); + } + }); +}); + +describe("ConnectAppsScreen interactions", () => { + function stubAppsWindow(overrides: Record = {}) { + vi.stubGlobal("IS_REACT_ACT_ENVIRONMENT", true); + vi.stubGlobal("navigator", { platform: "MacIntel" }); + vi.stubGlobal("window", { + OLLAMA_PLATFORM: "darwin", + innerHeight: 660, + addEventListener: vi.fn(), + removeEventListener: vi.fn(), + setTimeout: globalThis.setTimeout, + clearTimeout: globalThis.clearTimeout, + setInterval: globalThis.setInterval, + clearInterval: globalThis.clearInterval, + ...overrides, + }); + } + + it.each([false, true])( + "honors the native Claude disconnect confirmation: %s", + async (confirmed) => { + const connected = { + ...DISCONNECTED_CLAUDE, + used: true, + running: true, + configured: true, + connected: true, + }; + const confirm = vi.fn(() => confirmed); + const disconnect = vi.fn().mockResolvedValue({ + status: { ...connected, configured: false, connected: false }, + }); + stubAppsWindow({ + getClaudeDesktopConnectionSummary: vi.fn().mockResolvedValue(connected), + setClaudeDesktopConnected: disconnect, + confirm, + }); + let renderer: ReactTestRenderer | undefined; + try { + await act(async () => { + renderer = create( + , + ); + await settle(); + }); + const toggle = claudeConnectionButton(renderer!); + await act(async () => { + await toggle.props.onClick(); + }); + expect(confirm).toHaveBeenCalledOnce(); + if (confirmed) { + expect(disconnect).toHaveBeenCalledExactlyOnceWith(false, true); + } else { + expect(disconnect).not.toHaveBeenCalled(); + expect(toggle.props.disabled).toBe(false); + expect(toggle.props["aria-pressed"]).toBe(true); + } + } finally { + await act(async () => renderer?.unmount()); + vi.unstubAllGlobals(); + } + }, + ); + + it("lets the user retry Connect after cancelling the native confirmation", async () => { + const running = { ...DISCONNECTED_CLAUDE, running: true, used: true }; + const confirm = vi.fn().mockReturnValueOnce(false).mockReturnValue(true); + const connect = vi.fn().mockResolvedValue({ + status: { ...running, configured: true, connected: true }, + }); + stubAppsWindow({ + getClaudeDesktopConnectionSummary: vi.fn().mockResolvedValue(running), + setClaudeDesktopConnected: connect, + confirm, + }); + let renderer: ReactTestRenderer | undefined; + const clickConnect = () => + claudeConnectionButton(renderer!).props.onClick(); + try { + await act(async () => { + renderer = create( + + + , + ); + await settle(); + }); + await act(async () => { + await clickConnect(); + }); + expect(confirm).toHaveBeenCalledTimes(1); + expect(connect).not.toHaveBeenCalled(); + await act(async () => { + await clickConnect(); + }); + expect(confirm).toHaveBeenCalledTimes(2); + expect(connect).toHaveBeenCalledExactlyOnceWith(true, true); + } finally { + if (renderer) act(() => renderer?.unmount()); + vi.unstubAllGlobals(); + } + }); + + it.each(["darwin", "windows"])( + "copies every catalog command and limits desktop connections to macOS (%s)", + async (platform) => { + const copyCommand = vi + .spyOn(clipboard, "copyTextToClipboard") + .mockResolvedValue(true); + stubAppsWindow({ OLLAMA_PLATFORM: platform }); + const integrations = [ + ...appsIntegrations(true), + ...Array.from({ length: 20 }, (_, index) => ({ + id: `extra-${index}`, + name: `Extra app ${index}`, + description: "Another supported integration", + command: `ollama launch extra-${index}`, + })), + ]; + const launchers = integrations.filter((item) => item.command); + let renderer: ReactTestRenderer | undefined; + try { + await act(async () => { + renderer = create( + , + ); + }); + const cards = renderer!.root.findAll( + (node) => + node.type === "button" && node.props.id?.startsWith("integration-"), + ); + expect(new Set(cards.map((card) => card.props.id))).toEqual( + new Set(launchers.map((item) => `integration-${item.id}`)), + ); + expect( + renderer!.root.findAll( + (node) => + node.type === "button" && + typeof node.props["aria-pressed"] === "boolean", + ), + ).toHaveLength(platform === "darwin" ? 2 : 0); + for (const item of launchers) { + await act(async () => { + await renderer!.root + .findByProps({ id: `integration-${item.id}` }) + .props.onClick(); + }); + expect(copyCommand).toHaveBeenLastCalledWith(item.command); + } + expect(copyCommand).toHaveBeenCalledTimes(launchers.length); + } finally { + if (renderer) act(() => renderer?.unmount()); + copyCommand.mockRestore(); + vi.unstubAllGlobals(); + } + }, + ); + + it("renews the copy notification on each click and dismisses it after inactivity", async () => { + vi.useFakeTimers(); + const copyCommand = vi + .spyOn(clipboard, "copyTextToClipboard") + .mockResolvedValue(true); + stubAppsWindow(); + let renderer: ReactTestRenderer | undefined; + try { + await act(async () => { + renderer = create( + , + ); + await settle(); + }); + const card = () => + renderer!.root.findByProps({ id: "integration-codex" }); + await act(async () => { + await card().props.onClick(); + }); + expect(copyCommand).toHaveBeenCalledExactlyOnceWith( + "ollama launch codex", + ); + const notice = () => renderer!.root.findByProps({ role: "status" }); + expect(notice()).toBeTruthy(); + act(() => vi.advanceTimersByTime(5000)); + await act(async () => { + await card().props.onClick(); + }); + act(() => vi.advanceTimersByTime(1001)); + expect(notice()).toBeTruthy(); + act(() => vi.advanceTimersByTime(5000)); + expect(renderer!.root.findAllByProps({ role: "status" })).toHaveLength(0); + expect(copyCommand).toHaveBeenCalledTimes(2); + } finally { + if (renderer) act(() => renderer?.unmount()); + copyCommand.mockRestore(); + vi.unstubAllGlobals(); + vi.useRealTimers(); + } + }); + + it.each(["denied", "throws"])( + "offers manual copying instead of success when clipboard access %s", + async (outcome) => { + vi.useFakeTimers(); + const copyCommand = vi + .spyOn(clipboard, "copyTextToClipboard") + .mockImplementationOnce(() => + outcome === "throws" + ? Promise.reject(new Error("clipboard unavailable")) + : Promise.resolve(false), + ) + .mockResolvedValue(true); + stubAppsWindow(); + let renderer: ReactTestRenderer | undefined; + try { + await act(async () => { + renderer = create( + , + ); + await settle(); + }); + const card = () => + renderer!.root.findByProps({ id: "integration-codex" }); + await act(async () => { + await card().props.onClick(); + }); + act(() => vi.advanceTimersByTime(20_000)); + expect( + renderer!.root.findByProps({ role: "alert" }).findByType("code") + .children, + ).toEqual(["ollama launch codex"]); + expect(renderer!.root.findAllByProps({ role: "status" })).toHaveLength( + 0, + ); + await act(async () => { + await card().props.onClick(); + }); + expect(renderer!.root.findAllByProps({ role: "alert" })).toHaveLength( + 0, + ); + expect(renderer!.root.findByProps({ role: "status" })).toBeTruthy(); + } finally { + if (renderer) act(() => renderer?.unmount()); + copyCommand.mockRestore(); + vi.unstubAllGlobals(); + vi.useRealTimers(); + } + }, + ); + + it.each(["Escape", "outside pointer"])( + "dismisses a copy error with %s while preserving manual copying", + async (dismissal) => { + const events = new EventTarget(); + const commandNode = {}; + const noticeNode = { + contains: (target: unknown) => target === commandNode, + }; + const copyCommand = vi + .spyOn(clipboard, "copyTextToClipboard") + .mockResolvedValueOnce(false) + .mockResolvedValue(true); + stubAppsWindow({ + addEventListener: events.addEventListener.bind(events), + removeEventListener: events.removeEventListener.bind(events), + }); + let renderer: ReactTestRenderer | undefined; + try { + await act(async () => { + renderer = create( + , + { + createNodeMock: (element) => + element.props.role === "alert" ? noticeNode : null, + }, + ); + await settle(); + }); + const card = () => + renderer!.root.findByProps({ id: "integration-codex" }); + await act(async () => { + await card().props.onClick(); + }); + + // Selecting the command and unrelated keys must keep it available. + const selection = new Event("pointerdown"); + Object.defineProperty(selection, "target", { value: commandNode }); + act(() => { + events.dispatchEvent(selection); + events.dispatchEvent( + Object.assign(new Event("keydown"), { key: "c" }), + ); + }); + expect( + renderer!.root.findByProps({ role: "alert" }).findByType("code") + .children, + ).toEqual(["ollama launch codex"]); + + act(() => { + events.dispatchEvent( + dismissal === "Escape" + ? Object.assign(new Event("keydown"), { key: "Escape" }) + : new Event("pointerdown"), + ); + }); + expect(renderer!.root.findAllByProps({ role: "alert" })).toHaveLength( + 0, + ); + + // A later successful copy keeps its usual notification lifetime. + await act(async () => { + await card().props.onClick(); + }); + act(() => { + events.dispatchEvent( + Object.assign(new Event("keydown"), { key: "Escape" }), + ); + events.dispatchEvent(new Event("pointerdown")); + }); + expect(renderer!.root.findByProps({ role: "status" })).toBeTruthy(); + } finally { + if (renderer) act(() => renderer?.unmount()); + copyCommand.mockRestore(); + vi.unstubAllGlobals(); + } + }, + ); + + it.each(["connected", "disconnected", "failed"])( + "lets users copy commands while Claude status is still loading (%s)", + async (outcome) => { + let resolveClaude!: (status: typeof DISCONNECTED_CLAUDE) => void; + let rejectClaude!: (error: Error) => void; + const claude = new Promise( + (resolve, reject) => { + resolveClaude = resolve; + rejectClaude = reject; + }, + ); + const copyCommand = vi + .spyOn(clipboard, "copyTextToClipboard") + .mockResolvedValue(true); + const connect = vi.fn(); + stubAppsWindow({ + getClaudeDesktopConnectionSummary: vi.fn().mockReturnValue(claude), + setClaudeDesktopConnected: connect, + }); + vi.stubGlobal( + "fetch", + vi + .fn() + .mockResolvedValue( + new Response(JSON.stringify(appsIntegrations(true))), + ), + ); + + let renderer: ReactTestRenderer | undefined; + try { + await act(async () => { + renderer = create(); + await settle(); + }); + + const claudeToggle = () => claudeConnectionButton(renderer!); + expect(claudeToggle().props.disabled).toBe(true); + expect(claudeToggle().props["aria-busy"]).toBe(true); + + await act(async () => { + await renderer!.root + .findByProps({ id: "integration-codex" }) + .props.onClick(); + }); + expect(copyCommand).toHaveBeenCalledWith("ollama launch codex"); + expect(connect).not.toHaveBeenCalled(); + + await act(async () => { + if (outcome === "failed") { + rejectClaude(new Error("Claude status unavailable")); + } else { + resolveClaude({ + ...DISCONNECTED_CLAUDE, + used: true, + configured: outcome === "connected", + connected: outcome === "connected", + }); + } + await settle(); + }); + expect(claudeToggle().props.disabled).toBe(false); + expect(claudeToggle().props["aria-pressed"]).toBe( + outcome === "connected", + ); + if (outcome === "failed") { + expect(renderer!.root.findByProps({ role: "alert" })).toBeTruthy(); + } + expect(connect).not.toHaveBeenCalled(); + } finally { + if (renderer) act(() => renderer?.unmount()); + copyCommand.mockRestore(); + vi.unstubAllGlobals(); + } + }, + ); + + it("shows an app-list error without waiting for Claude status", async () => { + stubAppsWindow({ + getClaudeDesktopConnectionSummary: vi + .fn() + .mockReturnValue(new Promise(() => {})), + }); + vi.stubGlobal( + "fetch", + vi.fn().mockResolvedValue(new Response(null, { status: 500 })), + ); + + let renderer: ReactTestRenderer | undefined; + try { + await act(async () => { + renderer = create(); + await settle(); + }); + expect(renderer!.root.findByProps({ role: "alert" })).toBeTruthy(); + } finally { + if (renderer) act(() => renderer?.unmount()); + vi.unstubAllGlobals(); + } + }); + + it("shows the first-use intro after the user connects Claude", async () => { + const connectedStatus = { + ...DISCONNECTED_CLAUDE, + configured: true, + connected: true, + }; + const setClaudeDesktopConnected = vi + .fn() + .mockResolvedValue({ status: connectedStatus }); + stubAppsWindow({ + getClaudeDesktopConnectionSummary: vi + .fn() + .mockResolvedValue(DISCONNECTED_CLAUDE), + setClaudeDesktopConnected, + activateOllama: vi.fn(), + }); + + let renderer: ReactTestRenderer | undefined; + try { + await act(async () => { + renderer = create( + + + , + ); + await settle(); + }); + expect(setClaudeDesktopConnected).not.toHaveBeenCalled(); + await act(async () => { + await claudeConnectionButton(renderer!).props.onClick(); + }); + + expect(setClaudeDesktopConnected).toHaveBeenCalledTimes(1); + expect(setClaudeDesktopConnected).toHaveBeenCalledWith(true, false); + expect(claudeConnectionButton(renderer!).props["aria-pressed"]).toBe( + true, + ); + expect(renderer!.root.findByType(ClaudeConnectedIntro)).toBeTruthy(); + } finally { + if (renderer) act(() => renderer?.unmount()); + vi.unstubAllGlobals(); + } + }); +}); diff --git a/app/ui/app/src/components/Onboarding.tsx b/app/ui/app/src/components/Onboarding.tsx index 976d0672a98..3f5af1b2088 100644 --- a/app/ui/app/src/components/Onboarding.tsx +++ b/app/ui/app/src/components/Onboarding.tsx @@ -1,7 +1,9 @@ import CopyButton from "@/components/CopyButton"; import { CodexDesktopRow } from "@/components/CodexDesktopRow"; import Logo from "@/components/Logo"; -import { nextOnboardingStep, type OnboardingStep } from "@/lib/onboarding"; +import type { OnboardingStep } from "@/lib/onboarding"; +import { IntegrationConnectButton } from "@/components/IntegrationConnectButton"; +import { Transition } from "@headlessui/react"; import { getIntegrationStatuses, type IntegrationStatus, @@ -27,13 +29,10 @@ import type { } from "@/types/webview"; import { copyTextToClipboard } from "@/utils/clipboard"; import { - ArrowPathIcon, ArrowsRightLeftIcon, CommandLineIcon, ShieldCheckIcon, - Square2StackIcon, } from "@heroicons/react/24/outline"; -import { CheckIcon } from "@heroicons/react/20/solid"; import { useCallback, useEffect, @@ -84,6 +83,7 @@ interface ScreenProps { interface WelcomeScreenProps extends ScreenProps { isAuthenticated: boolean; + isLeaving?: boolean; completionError?: string | null; onRetryCompletion?: () => void; onLocal: () => void; @@ -170,8 +170,10 @@ export function IntroScreen({ completionError = null, onContinue, onRetryCompletion, + isLeaving = false, }: { completionError?: string | null; + isLeaving?: boolean; onContinue: () => void; onRetryCompletion?: () => void; }) { @@ -220,6 +222,8 @@ export function IntroScreen({ type="button" className="mt-8 flex h-11 w-full max-w-[240px] cursor-pointer items-center justify-center rounded-full bg-neutral-900 px-5 font-sans text-sm font-normal text-white transition-colors hover:bg-neutral-800 focus-visible:outline-2 focus-visible:outline-offset-2 focus-visible:outline-neutral-500" onClick={onContinue} + disabled={isLeaving} + aria-busy={isLeaving || undefined} > Continue @@ -258,6 +262,7 @@ function InlineError({ export function WelcomeScreen({ isAuthenticated, isSigningIn, + isLeaving = false, signInError, completionError = null, onSignIn, @@ -285,15 +290,20 @@ export function WelcomeScreen({ type="button" className="flex h-11 w-full cursor-pointer items-center justify-center rounded-full bg-neutral-900 px-5 font-sans text-sm font-normal text-white transition-colors hover:bg-neutral-800 focus-visible:outline-2 focus-visible:outline-offset-2 focus-visible:outline-neutral-500 disabled:cursor-wait disabled:opacity-70" onClick={onSignUp} - disabled={isSigningIn} - aria-busy={isSigningIn} + disabled={isSigningIn || isLeaving} + aria-busy={isSigningIn || isLeaving} > - {isSigningIn ? "Finish in your browser…" : "Sign up"} + {isLeaving + ? "Opening apps…" + : isSigningIn + ? "Finish in your browser…" + : "Sign up"} @@ -357,8 +367,8 @@ export function RunOllamaScreen({ ); } -function LaunchCommandIcon({ item }: { item: IntegrationStatus }) { - const icon = INTEGRATION_ICONS[item.id]; +function LaunchCommandIcon({ id }: { id: string }) { + const icon = INTEGRATION_ICONS[id]; return (
@@ -438,7 +448,20 @@ export function ConnectAppsScreen({ initialCodexStatus, }: ConnectAppsScreenProps) { const isWindows = isWindowsPlatform(); - const [copiedCommand, setCopiedCommand] = useState(null); + const [copyNotice, setCopyNotice] = useState<{ + sequence: number; + id: string; + name: string; + command: string; + copied: boolean; + visible: boolean; + } | null>(null); + const copyInFlight = useRef(false); + const copySequence = useRef(0); + const copyNoticeRef = useRef(null); + const [initialClaudeStatusSettled, setInitialClaudeStatusSettled] = useState( + Boolean(initialClaudeStatus), + ); const [claudeError, setClaudeError] = useState(null); const [claudeStatus, setClaudeStatus] = useState( initialClaudeStatus ?? null, @@ -461,14 +484,34 @@ export function ConnectAppsScreen({ }, []); useEffect(() => { - if (!copiedCommand) return; - + if (!copyNotice?.copied) return; const timeout = window.setTimeout(() => { - setCopiedCommand(null); - }, 5000); - + setCopyNotice((current) => current && { ...current, visible: false }); + }, 6000); return () => window.clearTimeout(timeout); - }, [copiedCommand]); + }, [copyNotice?.sequence, copyNotice?.copied]); + + useEffect(() => { + if (!copyNotice?.visible || copyNotice.copied) return; + + const dismiss = () => { + setCopyNotice((current) => + current && !current.copied ? { ...current, visible: false } : current, + ); + }; + const onPointerDown = (event: PointerEvent) => { + if (!copyNoticeRef.current?.contains(event.target as Node)) dismiss(); + }; + const onKeyDown = (event: KeyboardEvent) => { + if (event.key === "Escape" && !event.defaultPrevented) dismiss(); + }; + window.addEventListener("pointerdown", onPointerDown); + window.addEventListener("keydown", onKeyDown); + return () => { + window.removeEventListener("pointerdown", onPointerDown); + window.removeEventListener("keydown", onKeyDown); + }; + }, [copyNotice?.copied, copyNotice?.visible]); const refreshClaudeStatus = useCallback(async () => { if (isWindows) return null; @@ -497,30 +540,46 @@ export function ConnectAppsScreen({ const integrations = initialIntegrations ? Promise.resolve(initialIntegrations) : getIntegrationStatuses(); + + void integrations.then( + (statuses) => { + if (!active) return; + setIntegrationStatuses(statuses); + }, + () => { + if (!active) return; + setStatusError(true); + }, + ); + + return () => { + active = false; + }; + }, [initialIntegrations]); + + useEffect(() => { + let active = true; const claude = initialClaudeStatus ? Promise.resolve(initialClaudeStatus) : getClaudeConnectionSummary(); - void Promise.allSettled([integrations, claude]).then( - ([integrationResult, claudeResult]) => { + void claude.then( + (status) => { if (!active) return; - if (integrationResult.status === "fulfilled") { - setIntegrationStatuses(integrationResult.value); - } else { - setStatusError(true); - } - if (claudeResult.status === "fulfilled") { - setClaudeStatus(claudeResult.value); - } else { - setClaudeError("Ollama could not read the Claude connection status."); - } + setClaudeStatus(status); + setInitialClaudeStatusSettled(true); + }, + () => { + if (!active) return; + setClaudeError("Ollama could not read the Claude connection status."); + setInitialClaudeStatusSettled(true); }, ); return () => { active = false; }; - }, [initialClaudeStatus, initialIntegrations]); + }, [initialClaudeStatus]); const openConnectedClaude = useCallback( async (status: ClaudeDesktopStatus) => { @@ -593,12 +652,13 @@ export function ConnectAppsScreen({ ); const dismissClaudeConnectedIntro = async () => { - if (!window.setClaudeDesktopConnected) return; + if (!window.setClaudeDesktopConnected || claudePhase !== "idle") return; setClaudePhase("launching"); try { const liveStatus = await withClaudeConnectionTimeout( getClaudeConnectionSummary(), ); + if (!screenMounted.current) return; if (!liveStatus) { throw new Error("Claude Desktop connection status is unavailable"); } @@ -609,6 +669,7 @@ export function ConnectAppsScreen({ restartConfirmed = window.confirm( "Restart Claude Desktop to use Ollama? Any running task will stop.", ); + if (!screenMounted.current) return; if (!restartConfirmed) { setClaudePhase("idle"); return; @@ -755,10 +816,25 @@ export function ConnectAppsScreen({ }, [claudePhase, finishClaudeConnection, reconcileLateClaudeAction]); const copyLaunchCommand = async (item: IntegrationStatus) => { - if (item.command && (await copyTextToClipboard(item.command))) { - setClaudeError(null); - setCopiedCommand(item.command); + if (!item.command || copyInFlight.current) return; + copyInFlight.current = true; + let copied = false; + try { + copied = await copyTextToClipboard(item.command); + } catch { + // Keep the command available for manual copying when clipboard access fails. + } finally { + copyInFlight.current = false; } + if (!screenMounted.current) return; + setCopyNotice({ + sequence: ++copySequence.current, + id: item.id, + name: item.name, + command: item.command, + copied, + visible: true, + }); }; const connectClaude = async () => { @@ -772,7 +848,7 @@ export function ConnectAppsScreen({ return; } - setCopiedCommand(null); + setCopyNotice((current) => current && { ...current, visible: false }); setClaudeError(null); const enabling = claudeStatus ? !isClaudeConfigured(claudeStatus) : true; setClaudePhase(enabling ? "connecting" : "disconnecting"); @@ -789,6 +865,7 @@ export function ConnectAppsScreen({ ); return; } + if (!screenMounted.current) return; if (!status) { setClaudePhase("idle"); setClaudeError("Ollama could not read the Claude connection status."); @@ -842,6 +919,7 @@ export function ConnectAppsScreen({ ? "Restart Claude Desktop to use Ollama? Any running task will stop." : "Restart Claude Desktop to remove Ollama? Any running task will stop.", ); + if (!screenMounted.current) return; if (!restartConfirmed) { setClaudePhase("idle"); return; @@ -936,67 +1014,54 @@ export function ConnectAppsScreen({ ? "Opening…" : claudePhase === "disconnecting" ? "Disconnecting…" - : !claudeConfigured && !claudeInstalled - ? "Download & connect" - : null; + : null; const claudeGuidance = claudeDesktopRecoveryMessage( claudeStatus?.error, claudeError, ); - const launchIntegrationRow = (item: IntegrationStatus) => { - const copied = copiedCommand === item.command; + const launchIntegrationCard = (item: IntegrationStatus) => { + const copied = + copyNotice?.id === item.id && copyNotice.copied && copyNotice.visible; return ( -
copyLaunchCommand(item)} + aria-label={ + copied ? `${item.name} command copied` : `Copy ${item.name} command` + } + title={item.description} + className="relative isolate flex min-w-0 items-center gap-3 rounded-2xl border border-neutral-200 bg-white px-4 py-3 text-left transition-colors duration-700 hover:bg-neutral-50 focus-visible:outline-2 focus-visible:outline-offset-2 focus-visible:outline-neutral-500 motion-reduce:transition-none dark:border-neutral-700 dark:bg-neutral-900 dark:hover:bg-neutral-800" > -
- -
-

- {item.name} -

-

- {item.description} -

-
-
-
- - {item.command} - - -
-
+ {copied && ( +