app: match model picker to server models

Rather than adding models to the model picker to guide users on first use, take the ollama tags response as a source of truth.

app/ui/app/src/hooks/useModels.ts

@@ -1,13 +1,12 @@
 import { useQuery } from "@tanstack/react-query";
 import { Model } from "@/gotypes";
 import { getModels } from "@/api";
-import { mergeModels } from "@/utils/mergeModels";
-import { useSettings } from "./useSettings";
 import { useMemo } from "react";
 
+const DEFAULT_MODEL = "gemma3:4b";
+
 export function useModels(searchQuery = "") {
-  const { settings } = useSettings();
-  const localQuery = useQuery<Model[], Error>({
+  const query = useQuery<Model[], Error>({
     queryKey: ["models", searchQuery],
     queryFn: () => getModels(searchQuery),
     gcTime: 10 * 60 * 1000, // Keep in cache for 10 minutes
@@ -19,33 +18,18 @@ export function useModels(searchQuery = "") {
     refetchIntervalInBackground: true,
   });
 
-  const allModels = useMemo(() => {
-    const models = mergeModels(localQuery.data || [], settings.airplaneMode);
-
-    if (searchQuery && searchQuery.trim()) {
-      const query = searchQuery.toLowerCase().trim();
-      const filteredModels = models.filter((model) =>
-        model.model.toLowerCase().includes(query),
-      );
-
-      const seen = new Set<string>();
-      return filteredModels.filter((model) => {
-        const currentModel = model.model.toLowerCase();
-        if (seen.has(currentModel)) {
-          return false;
-        }
-        seen.add(currentModel);
-        return true;
-      });
+  const models = useMemo(() => {
+    const data = query.data || [];
+    if (data.length === 0) {
+      return [new Model({ model: DEFAULT_MODEL })];
     }
-
-    return models;
-  }, [localQuery.data, searchQuery, settings.airplaneMode]);
+    return data;
+  }, [query.data]);
 
   return {
-    ...localQuery,
-    data: allModels,
-    isLoading: localQuery.isLoading,
+    ...query,
+    data: models,
+    isLoading: query.isLoading,
   };
 }
 

app/ui/app/src/hooks/useSelectedModel.ts

@@ -4,7 +4,6 @@ import { useModels } from "./useModels";
 import { useChat } from "./useChats";
 import { useSettings } from "./useSettings.ts";
 import { Model } from "@/gotypes";
-import { FEATURED_MODELS } from "@/utils/mergeModels";
 import { getTotalVRAM } from "@/utils/vram.ts";
 import { getInferenceCompute } from "@/api";
 
@@ -46,77 +45,13 @@ export function useSelectedModel(currentChatId?: string, searchQuery?: string) {
   const restoredChatRef = useRef<string | null>(null);
 
   const selectedModel: Model | null = useMemo(() => {
-    // if airplane mode is on and selected model ends with cloud,
-    // switch to recommended default model
-    if (settings.airplaneMode && settings.selectedModel?.endsWith("cloud")) {
-      return (
-        models.find((m) => m.model === recommendedModel) ||
-        models.find((m) => m.isCloud) ||
-        models.find((m) => m.digest === undefined || m.digest === "") ||
-        models[0] ||
-        null
-      );
-    }
-
-    // Migration logic: if turboEnabled is true and selectedModel is a base model,
-    // migrate to the cloud version and disable turboEnabled permanently
-    // TODO: remove this logic in a future release
-    const baseModelsToMigrate = [
-      "gpt-oss:20b",
-      "gpt-oss:120b",
-      "deepseek-v3.1:671b",
-      "qwen3-coder:480b",
-    ];
-    const shouldMigrate =
-      !settings.airplaneMode &&
-      settings.turboEnabled &&
-      baseModelsToMigrate.includes(settings.selectedModel);
-
-    if (shouldMigrate) {
-      const cloudModel = `${settings.selectedModel}cloud`;
-      return (
-        models.find((m) => m.model === cloudModel) ||
-        new Model({
-          model: cloudModel,
-          cloud: true,
-          ollama_host: false,
-        })
-      );
-    }
-
     return (
       models.find((m) => m.model === settings.selectedModel) ||
       (settings.selectedModel &&
-        new Model({
-          model: settings.selectedModel,
-          cloud: FEATURED_MODELS.some(
-            (f) => f.endsWith("cloud") && f === settings.selectedModel,
-          ),
-          ollama_host: false,
-        })) ||
+        new Model({ model: settings.selectedModel })) ||
       null
     );
-  }, [models, settings.selectedModel, settings.airplaneMode, recommendedModel]);
-
-  useEffect(() => {
-    if (!selectedModel) return;
-
-    if (
-      settings.airplaneMode &&
-      settings.selectedModel?.endsWith("cloud") &&
-      selectedModel.model !== settings.selectedModel
-    ) {
-      setSettings({ SelectedModel: selectedModel.model });
-    }
-
-    if (
-      !settings.airplaneMode &&
-      settings.turboEnabled &&
-      selectedModel.model !== settings.selectedModel
-    ) {
-      setSettings({ SelectedModel: selectedModel.model, TurboEnabled: false });
-    }
-  }, [selectedModel, settings.airplaneMode, settings.selectedModel]);
+  }, [models, settings.selectedModel]);
 
   // Set model from chat history when chat data loads
   useEffect(() => {
@@ -169,8 +104,6 @@ export function useSelectedModel(currentChatId?: string, searchQuery?: string) {
 
     const defaultModel =
       models.find((m) => m.model === recommendedModel) ||
-      models.find((m) => m.isCloud()) ||
-      models.find((m) => m.digest === undefined || m.digest === "") ||
       models[0];
 
     if (defaultModel) {

app/ui/app/src/utils/mergeModels.test.ts

@@ -1,128 +0,0 @@
-import { describe, it, expect } from "vitest";
-import { Model } from "@/gotypes";
-import { mergeModels, FEATURED_MODELS } from "@/utils/mergeModels";
-import "@/api";
-
-describe("Model merging logic", () => {
-  it("should handle cloud models with -cloud suffix", () => {
-    const localModels: Model[] = [
-      new Model({ model: "gpt-oss:120b-cloud" }),
-      new Model({ model: "llama3:latest" }),
-      new Model({ model: "mistral:latest" }),
-    ];
-
-    const merged = mergeModels(localModels);
-
-    // First verify cloud models are first and in FEATURED_MODELS order
-    const cloudModels = FEATURED_MODELS.filter((m: string) =>
-      m.endsWith("cloud"),
-    );
-    for (let i = 0; i < cloudModels.length; i++) {
-      expect(merged[i].model).toBe(cloudModels[i]);
-      expect(merged[i].isCloud()).toBe(true);
-    }
-
-    // Then verify non-cloud featured models are next and in FEATURED_MODELS order
-    const nonCloudFeatured = FEATURED_MODELS.filter(
-      (m: string) => !m.endsWith("cloud"),
-    );
-    for (let i = 0; i < nonCloudFeatured.length; i++) {
-      const model = merged[i + cloudModels.length];
-      expect(model.model).toBe(nonCloudFeatured[i]);
-      expect(model.isCloud()).toBe(false);
-    }
-
-    // Verify local models are preserved and come after featured models
-    const featuredCount = FEATURED_MODELS.length;
-    expect(merged[featuredCount].model).toBe("llama3:latest");
-    expect(merged[featuredCount + 1].model).toBe("mistral:latest");
-
-    // Length should be exactly featured models plus our local models
-    expect(merged.length).toBe(FEATURED_MODELS.length + 2);
-  });
-
-  it("should hide cloud models in airplane mode", () => {
-    const localModels: Model[] = [
-      new Model({ model: "gpt-oss:120b-cloud" }),
-      new Model({ model: "llama3:latest" }),
-      new Model({ model: "mistral:latest" }),
-    ];
-
-    const merged = mergeModels(localModels, true); // airplane mode = true
-
-    // No cloud models should be present
-    const cloudModels = merged.filter((m) => m.isCloud());
-    expect(cloudModels.length).toBe(0);
-
-    // Should have non-cloud featured models
-    const nonCloudFeatured = FEATURED_MODELS.filter(
-      (m) => !m.endsWith("cloud"),
-    );
-    for (let i = 0; i < nonCloudFeatured.length; i++) {
-      const model = merged[i];
-      expect(model.model).toBe(nonCloudFeatured[i]);
-      expect(model.isCloud()).toBe(false);
-    }
-
-    // Local models should be preserved
-    const featuredCount = nonCloudFeatured.length;
-    expect(merged[featuredCount].model).toBe("llama3:latest");
-    expect(merged[featuredCount + 1].model).toBe("mistral:latest");
-  });
-
-  it("should handle empty input", () => {
-    const merged = mergeModels([]);
-
-    // First verify cloud models are first and in FEATURED_MODELS order
-    const cloudModels = FEATURED_MODELS.filter((m) => m.endsWith("cloud"));
-    for (let i = 0; i < cloudModels.length; i++) {
-      expect(merged[i].model).toBe(cloudModels[i]);
-      expect(merged[i].isCloud()).toBe(true);
-    }
-
-    // Then verify non-cloud featured models are next and in FEATURED_MODELS order
-    const nonCloudFeatured = FEATURED_MODELS.filter(
-      (m) => !m.endsWith("cloud"),
-    );
-    for (let i = 0; i < nonCloudFeatured.length; i++) {
-      const model = merged[i + cloudModels.length];
-      expect(model.model).toBe(nonCloudFeatured[i]);
-      expect(model.isCloud()).toBe(false);
-    }
-
-    // Length should be exactly FEATURED_MODELS length
-    expect(merged.length).toBe(FEATURED_MODELS.length);
-  });
-
-  it("should sort models correctly", () => {
-    const localModels: Model[] = [
-      new Model({ model: "zephyr:latest" }),
-      new Model({ model: "alpha:latest" }),
-      new Model({ model: "gpt-oss:120b-cloud" }),
-    ];
-
-    const merged = mergeModels(localModels);
-
-    // First verify cloud models are first and in FEATURED_MODELS order
-    const cloudModels = FEATURED_MODELS.filter((m) => m.endsWith("cloud"));
-    for (let i = 0; i < cloudModels.length; i++) {
-      expect(merged[i].model).toBe(cloudModels[i]);
-      expect(merged[i].isCloud()).toBe(true);
-    }
-
-    // Then verify non-cloud featured models are next and in FEATURED_MODELS order
-    const nonCloudFeatured = FEATURED_MODELS.filter(
-      (m) => !m.endsWith("cloud"),
-    );
-    for (let i = 0; i < nonCloudFeatured.length; i++) {
-      const model = merged[i + cloudModels.length];
-      expect(model.model).toBe(nonCloudFeatured[i]);
-      expect(model.isCloud()).toBe(false);
-    }
-
-    // Non-featured local models should be at the end in alphabetical order
-    const featuredCount = FEATURED_MODELS.length;
-    expect(merged[featuredCount].model).toBe("alpha:latest");
-    expect(merged[featuredCount + 1].model).toBe("zephyr:latest");
-  });
-});

app/ui/app/src/utils/mergeModels.ts

@@ -1,101 +0,0 @@
-import { Model } from "@/gotypes";
-
-// Featured models list (in priority order)
-export const FEATURED_MODELS = [
-  "gpt-oss:120b-cloud",
-  "gpt-oss:20b-cloud",
-  "deepseek-v3.1:671b-cloud",
-  "qwen3-coder:480b-cloud",
-  "qwen3-vl:235b-cloud",
-  "minimax-m2:cloud",
-  "glm-4.6:cloud",
-  "gpt-oss:120b",
-  "gpt-oss:20b",
-  "gemma3:27b",
-  "gemma3:12b",
-  "gemma3:4b",
-  "gemma3:1b",
-  "deepseek-r1:8b",
-  "qwen3-coder:30b",
-  "qwen3-vl:30b",
-  "qwen3-vl:8b",
-  "qwen3-vl:4b",
-  "qwen3:30b",
-  "qwen3:8b",
-  "qwen3:4b",
-];
-
-function alphabeticalSort(a: Model, b: Model): number {
-  return a.model.toLowerCase().localeCompare(b.model.toLowerCase());
-}
-
-//Merges models, sorting cloud models first, then other models
-export function mergeModels(
-  localModels: Model[],
-  airplaneMode: boolean = false,
-): Model[] {
-  const allModels = (localModels || []).map((model) => model);
-
-  // 1. Get cloud models from local models and featured list
-  const cloudModels = [...allModels.filter((m) => m.isCloud())];
-
-  // Add any cloud models from FEATURED_MODELS that aren't in local models
-  FEATURED_MODELS.filter((f) => f.endsWith("cloud")).forEach((cloudModel) => {
-    if (!cloudModels.some((m) => m.model === cloudModel)) {
-      cloudModels.push(new Model({ model: cloudModel }));
-    }
-  });
-
-  // 2. Get other featured models (non-cloud)
-  const featuredModels = FEATURED_MODELS.filter(
-    (f) => !f.endsWith("cloud"),
-  ).map((model) => {
-    // Check if this model exists in local models
-    const localMatch = allModels.find(
-      (m) => m.model.toLowerCase() === model.toLowerCase(),
-    );
-
-    if (localMatch) return localMatch;
-
-    return new Model({
-      model,
-    });
-  });
-
-  // 3. Get remaining local models that aren't featured and aren't cloud models
-  const remainingModels = allModels.filter(
-    (model) =>
-      !model.isCloud() &&
-      !FEATURED_MODELS.some(
-        (f) => f.toLowerCase() === model.model.toLowerCase(),
-      ),
-  );
-
-  cloudModels.sort((a, b) => {
-    const aIndex = FEATURED_MODELS.indexOf(a.model);
-    const bIndex = FEATURED_MODELS.indexOf(b.model);
-
-    // If both are featured, sort by their position in FEATURED_MODELS
-    if (aIndex !== -1 && bIndex !== -1) {
-      return aIndex - bIndex;
-    }
-
-    // If only one is featured, featured model comes first
-    if (aIndex !== -1 && bIndex === -1) return -1;
-    if (aIndex === -1 && bIndex !== -1) return 1;
-
-    // If neither is featured, sort alphabetically
-    return a.model.toLowerCase().localeCompare(b.model.toLowerCase());
-  });
-
-  featuredModels.sort(
-    (a, b) =>
-      FEATURED_MODELS.indexOf(a.model) - FEATURED_MODELS.indexOf(b.model),
-  );
-
-  remainingModels.sort(alphabeticalSort);
-
-  return airplaneMode
-    ? [...featuredModels, ...remainingModels]
-    : [...cloudModels, ...featuredModels, ...remainingModels];
-}

cmd/cmd.go

@@ -367,14 +367,25 @@ func loadOrUnloadModel(cmd *cobra.Command, opts *runOptions) error {
 		return err
 	} else if info.RemoteHost != "" {
 		// Cloud model, no need to load/unload
+
+		isCloud := strings.HasPrefix(info.RemoteHost, "https://ollama.com")
+
+		// Check if user is signed in for ollama.com cloud models
+		if isCloud {
+			if _, err := client.Whoami(cmd.Context()); err != nil {
+				return err
+			}
+		}
+
 		if opts.ShowConnect {
 			p.StopAndClear()
-			if strings.HasPrefix(info.RemoteHost, "https://ollama.com") {
+			if isCloud {
 				fmt.Fprintf(os.Stderr, "Connecting to '%s' on 'ollama.com' ⚡\n", info.RemoteModel)
 			} else {
 				fmt.Fprintf(os.Stderr, "Connecting to '%s' on '%s'\n", info.RemoteModel, info.RemoteHost)
 			}
 		}
+
 		return nil
 	}
 

cmd/cmd_test.go

@@ -3,6 +3,7 @@ package cmd
 import (
 	"bytes"
 	"encoding/json"
+	"errors"
 	"fmt"
 	"io"
 	"net/http"
@@ -1659,3 +1660,103 @@ func TestRunOptions_Copy_Independence(t *testing.T) {
 		t.Error("Copy Think should not be affected by original modification")
 	}
 }
+
+func TestLoadOrUnloadModel_CloudModelAuth(t *testing.T) {
+	tests := []struct {
+		name          string
+		remoteHost    string
+		whoamiStatus  int
+		whoamiResp    any
+		expectedError string
+	}{
+		{
+			name:         "ollama.com cloud model - user signed in",
+			remoteHost:   "https://ollama.com",
+			whoamiStatus: http.StatusOK,
+			whoamiResp:   api.UserResponse{Name: "testuser"},
+		},
+		{
+			name:         "ollama.com cloud model - user not signed in",
+			remoteHost:   "https://ollama.com",
+			whoamiStatus: http.StatusUnauthorized,
+			whoamiResp: map[string]string{
+				"error":      "unauthorized",
+				"signin_url": "https://ollama.com/signin",
+			},
+			expectedError: "unauthorized",
+		},
+		{
+			name:         "non-ollama.com remote - no auth check",
+			remoteHost:   "https://other-remote.com",
+			whoamiStatus: http.StatusUnauthorized, // should not be called
+			whoamiResp:   nil,
+		},
+	}
+
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			whoamiCalled := false
+			mockServer := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+				switch r.URL.Path {
+				case "/api/show":
+					w.Header().Set("Content-Type", "application/json")
+					if err := json.NewEncoder(w).Encode(api.ShowResponse{
+						RemoteHost:  tt.remoteHost,
+						RemoteModel: "test-model",
+					}); err != nil {
+						http.Error(w, err.Error(), http.StatusInternalServerError)
+					}
+				case "/api/me":
+					whoamiCalled = true
+					w.Header().Set("Content-Type", "application/json")
+					w.WriteHeader(tt.whoamiStatus)
+					if tt.whoamiResp != nil {
+						if err := json.NewEncoder(w).Encode(tt.whoamiResp); err != nil {
+							http.Error(w, err.Error(), http.StatusInternalServerError)
+						}
+					}
+				default:
+					http.NotFound(w, r)
+				}
+			}))
+			defer mockServer.Close()
+
+			t.Setenv("OLLAMA_HOST", mockServer.URL)
+
+			cmd := &cobra.Command{}
+			cmd.SetContext(t.Context())
+
+			opts := &runOptions{
+				Model:       "test-cloud-model",
+				ShowConnect: false,
+			}
+
+			err := loadOrUnloadModel(cmd, opts)
+
+			if strings.HasPrefix(tt.remoteHost, "https://ollama.com") {
+				if !whoamiCalled {
+					t.Error("expected whoami to be called for ollama.com cloud model")
+				}
+			} else {
+				if whoamiCalled {
+					t.Error("whoami should not be called for non-ollama.com remote")
+				}
+			}
+
+			if tt.expectedError != "" {
+				if err == nil {
+					t.Errorf("expected error containing %q, got nil", tt.expectedError)
+				} else {
+					var authErr api.AuthorizationError
+					if !errors.As(err, &authErr) {
+						t.Errorf("expected AuthorizationError, got %T: %v", err, err)
+					}
+				}
+			} else {
+				if err != nil {
+					t.Errorf("expected no error, got %v", err)
+				}
+			}
+		})
+	}
+}

cmd/config/integrations.go

@@ -48,7 +48,6 @@ var integrations = map[string]Runner{
 	"droid":    &Droid{},
 	"opencode": &OpenCode{},
 	"openclaw": &Openclaw{},
-	"pi":       &Pi{},
 }
 
 // recommendedModels are shown when the user has no models or as suggestions.
@@ -64,7 +63,6 @@ var recommendedModels = []selectItem{
 var integrationAliases = map[string]bool{
 	"clawdbot": true,
 	"moltbot":  true,
-	"pi":       true,
 }
 
 func selectIntegration() (string, error) {

cmd/config/pi.go

@@ -1,196 +0,0 @@
-package config
-
-import (
-	"encoding/json"
-	"fmt"
-	"os"
-	"os/exec"
-	"path/filepath"
-	"slices"
-
-	"github.com/ollama/ollama/envconfig"
-)
-
-// Pi implements Runner and Editor for Pi (Pi Coding Agent) integration
-type Pi struct{}
-
-func (p *Pi) String() string { return "Pi" }
-
-func (p *Pi) Run(model string, args []string) error {
-	if _, err := exec.LookPath("pi"); err != nil {
-		return fmt.Errorf("pi is not installed, install with: npm install -g @mariozechner/pi-coding-agent")
-	}
-
-	// Call Edit() to ensure config is up-to-date before launch
-	models := []string{model}
-	if config, err := loadIntegration("pi"); err == nil && len(config.Models) > 0 {
-		models = config.Models
-	}
-	if err := p.Edit(models); err != nil {
-		return fmt.Errorf("setup failed: %w", err)
-	}
-
-	cmd := exec.Command("pi", args...)
-	cmd.Stdin = os.Stdin
-	cmd.Stdout = os.Stdout
-	cmd.Stderr = os.Stderr
-	return cmd.Run()
-}
-
-func (p *Pi) Paths() []string {
-	home, err := os.UserHomeDir()
-	if err != nil {
-		return nil
-	}
-
-	var paths []string
-	modelsPath := filepath.Join(home, ".pi", "agent", "models.json")
-	if _, err := os.Stat(modelsPath); err == nil {
-		paths = append(paths, modelsPath)
-	}
-	settingsPath := filepath.Join(home, ".pi", "agent", "settings.json")
-	if _, err := os.Stat(settingsPath); err == nil {
-		paths = append(paths, settingsPath)
-	}
-	return paths
-}
-
-func (p *Pi) Edit(models []string) error {
-	if len(models) == 0 {
-		return nil
-	}
-
-	home, err := os.UserHomeDir()
-	if err != nil {
-		return err
-	}
-
-	configPath := filepath.Join(home, ".pi", "agent", "models.json")
-	if err := os.MkdirAll(filepath.Dir(configPath), 0o755); err != nil {
-		return err
-	}
-
-	config := make(map[string]any)
-	if data, err := os.ReadFile(configPath); err == nil {
-		_ = json.Unmarshal(data, &config)
-	}
-
-	providers, ok := config["providers"].(map[string]any)
-	if !ok {
-		providers = make(map[string]any)
-	}
-
-	ollama, ok := providers["ollama"].(map[string]any)
-	if !ok {
-		ollama = map[string]any{
-			"baseUrl": envconfig.Host().String() + "/v1",
-			"api":     "openai-completions",
-			"apiKey":  "ollama",
-		}
-	}
-
-	existingModels, ok := ollama["models"].([]any)
-	if !ok {
-		existingModels = make([]any, 0)
-	}
-
-	// Build set of selected models to track which need to be added
-	selectedSet := make(map[string]bool, len(models))
-	for _, m := range models {
-		selectedSet[m] = true
-	}
-
-	// Build new models list:
-	// 1. Keep user-managed models (no _launch marker) - untouched
-	// 2. Keep ollama-managed models (_launch marker) that are still selected
-	// 3. Add new ollama-managed models
-	var newModels []any
-	for _, m := range existingModels {
-		if modelObj, ok := m.(map[string]any); ok {
-			if id, ok := modelObj["id"].(string); ok {
-				// User-managed model (no _launch marker) - always preserve
-				if !isPiOllamaModel(modelObj) {
-					newModels = append(newModels, m)
-				} else if selectedSet[id] {
-					// Ollama-managed and still selected - keep it
-					newModels = append(newModels, m)
-					selectedSet[id] = false
-				}
-			}
-		}
-	}
-
-	// Add newly selected models that weren't already in the list
-	for _, model := range models {
-		if selectedSet[model] {
-			newModels = append(newModels, map[string]any{
-				"id":      model,
-				"_launch": true,
-			})
-		}
-	}
-
-	ollama["models"] = newModels
-	providers["ollama"] = ollama
-	config["providers"] = providers
-
-	configData, err := json.MarshalIndent(config, "", "  ")
-	if err != nil {
-		return err
-	}
-	if err := writeWithBackup(configPath, configData); err != nil {
-		return err
-	}
-
-	// Update settings.json with default provider and model
-	settingsPath := filepath.Join(home, ".pi", "agent", "settings.json")
-	settings := make(map[string]any)
-	if data, err := os.ReadFile(settingsPath); err == nil {
-		_ = json.Unmarshal(data, &settings)
-	}
-
-	settings["defaultProvider"] = "ollama"
-	settings["defaultModel"] = models[0]
-
-	settingsData, err := json.MarshalIndent(settings, "", "  ")
-	if err != nil {
-		return err
-	}
-	return writeWithBackup(settingsPath, settingsData)
-}
-
-func (p *Pi) Models() []string {
-	home, err := os.UserHomeDir()
-	if err != nil {
-		return nil
-	}
-
-	configPath := filepath.Join(home, ".pi", "agent", "models.json")
-	config, err := readJSONFile(configPath)
-	if err != nil {
-		return nil
-	}
-
-	providers, _ := config["providers"].(map[string]any)
-	ollama, _ := providers["ollama"].(map[string]any)
-	models, _ := ollama["models"].([]any)
-
-	var result []string
-	for _, m := range models {
-		if modelObj, ok := m.(map[string]any); ok {
-			if id, ok := modelObj["id"].(string); ok {
-				result = append(result, id)
-			}
-		}
-	}
-	slices.Sort(result)
-	return result
-}
-
-// isPiOllamaModel reports whether a model config entry is managed by ollama launch
-func isPiOllamaModel(cfg map[string]any) bool {
-	if v, ok := cfg["_launch"].(bool); ok && v {
-		return true
-	}
-	return false
-}

cmd/config/pi_test.go

@@ -1,609 +0,0 @@
-package config
-
-import (
-	"encoding/json"
-	"os"
-	"path/filepath"
-	"testing"
-)
-
-func TestPiIntegration(t *testing.T) {
-	pi := &Pi{}
-
-	t.Run("String", func(t *testing.T) {
-		if got := pi.String(); got != "Pi" {
-			t.Errorf("String() = %q, want %q", got, "Pi")
-		}
-	})
-
-	t.Run("implements Runner", func(t *testing.T) {
-		var _ Runner = pi
-	})
-
-	t.Run("implements Editor", func(t *testing.T) {
-		var _ Editor = pi
-	})
-}
-
-func TestPiPaths(t *testing.T) {
-	pi := &Pi{}
-
-	t.Run("returns empty when no config exists", func(t *testing.T) {
-		tmpDir := t.TempDir()
-		setTestHome(t, tmpDir)
-
-		paths := pi.Paths()
-		if len(paths) != 0 {
-			t.Errorf("Paths() = %v, want empty", paths)
-		}
-	})
-
-	t.Run("returns path when config exists", func(t *testing.T) {
-		tmpDir := t.TempDir()
-		setTestHome(t, tmpDir)
-
-		configDir := filepath.Join(tmpDir, ".pi", "agent")
-		if err := os.MkdirAll(configDir, 0o755); err != nil {
-			t.Fatal(err)
-		}
-		configPath := filepath.Join(configDir, "models.json")
-		if err := os.WriteFile(configPath, []byte("{}"), 0o644); err != nil {
-			t.Fatal(err)
-		}
-
-		paths := pi.Paths()
-		if len(paths) != 1 || paths[0] != configPath {
-			t.Errorf("Paths() = %v, want [%s]", paths, configPath)
-		}
-	})
-}
-
-func TestPiEdit(t *testing.T) {
-	pi := &Pi{}
-	tmpDir := t.TempDir()
-	setTestHome(t, tmpDir)
-
-	configDir := filepath.Join(tmpDir, ".pi", "agent")
-	configPath := filepath.Join(configDir, "models.json")
-
-	cleanup := func() {
-		os.RemoveAll(configDir)
-	}
-
-	readConfig := func() map[string]any {
-		data, _ := os.ReadFile(configPath)
-		var cfg map[string]any
-		json.Unmarshal(data, &cfg)
-		return cfg
-	}
-
-	t.Run("returns nil for empty models", func(t *testing.T) {
-		if err := pi.Edit([]string{}); err != nil {
-			t.Errorf("Edit([]) error = %v, want nil", err)
-		}
-	})
-
-	t.Run("creates config with models", func(t *testing.T) {
-		cleanup()
-
-		models := []string{"llama3.2", "qwen3:8b"}
-		if err := pi.Edit(models); err != nil {
-			t.Fatalf("Edit() error = %v", err)
-		}
-
-		cfg := readConfig()
-
-		providers, ok := cfg["providers"].(map[string]any)
-		if !ok {
-			t.Error("Config missing providers")
-		}
-
-		ollama, ok := providers["ollama"].(map[string]any)
-		if !ok {
-			t.Error("Providers missing ollama")
-		}
-
-		modelsArray, ok := ollama["models"].([]any)
-		if !ok || len(modelsArray) != 2 {
-			t.Errorf("Expected 2 models, got %v", modelsArray)
-		}
-
-		if ollama["baseUrl"] == nil {
-			t.Error("Missing baseUrl")
-		}
-		if ollama["api"] != "openai-completions" {
-			t.Errorf("Expected api=openai-completions, got %v", ollama["api"])
-		}
-		if ollama["apiKey"] != "ollama" {
-			t.Errorf("Expected apiKey=ollama, got %v", ollama["apiKey"])
-		}
-	})
-
-	t.Run("updates existing config preserving ollama provider settings", func(t *testing.T) {
-		cleanup()
-		os.MkdirAll(configDir, 0o755)
-
-		existingConfig := `{
-			"providers": {
-				"ollama": {
-					"baseUrl": "http://custom:8080/v1",
-					"api": "custom-api",
-					"apiKey": "custom-key",
-					"models": [
-						{"id": "old-model", "_launch": true}
-					]
-				}
-			}
-		}`
-		if err := os.WriteFile(configPath, []byte(existingConfig), 0o644); err != nil {
-			t.Fatal(err)
-		}
-
-		models := []string{"new-model"}
-		if err := pi.Edit(models); err != nil {
-			t.Fatalf("Edit() error = %v", err)
-		}
-
-		cfg := readConfig()
-		providers := cfg["providers"].(map[string]any)
-		ollama := providers["ollama"].(map[string]any)
-
-		if ollama["baseUrl"] != "http://custom:8080/v1" {
-			t.Errorf("Custom baseUrl not preserved, got %v", ollama["baseUrl"])
-		}
-		if ollama["api"] != "custom-api" {
-			t.Errorf("Custom api not preserved, got %v", ollama["api"])
-		}
-		if ollama["apiKey"] != "custom-key" {
-			t.Errorf("Custom apiKey not preserved, got %v", ollama["apiKey"])
-		}
-
-		modelsArray := ollama["models"].([]any)
-		if len(modelsArray) != 1 {
-			t.Errorf("Expected 1 model after update, got %d", len(modelsArray))
-		} else {
-			modelEntry := modelsArray[0].(map[string]any)
-			if modelEntry["id"] != "new-model" {
-				t.Errorf("Expected new-model, got %v", modelEntry["id"])
-			}
-			// Verify _launch marker is present
-			if modelEntry["_launch"] != true {
-				t.Errorf("Expected _launch marker to be true")
-			}
-		}
-	})
-
-	t.Run("replaces old models with new ones", func(t *testing.T) {
-		cleanup()
-		os.MkdirAll(configDir, 0o755)
-
-		// Old models must have _launch marker to be managed by us
-		existingConfig := `{
-			"providers": {
-				"ollama": {
-					"baseUrl": "http://localhost:11434/v1",
-					"api": "openai-completions",
-					"apiKey": "ollama",
-					"models": [
-						{"id": "old-model-1", "_launch": true},
-						{"id": "old-model-2", "_launch": true}
-					]
-				}
-			}
-		}`
-		if err := os.WriteFile(configPath, []byte(existingConfig), 0o644); err != nil {
-			t.Fatal(err)
-		}
-
-		newModels := []string{"new-model-1", "new-model-2"}
-		if err := pi.Edit(newModels); err != nil {
-			t.Fatalf("Edit() error = %v", err)
-		}
-
-		cfg := readConfig()
-		providers := cfg["providers"].(map[string]any)
-		ollama := providers["ollama"].(map[string]any)
-		modelsArray := ollama["models"].([]any)
-
-		if len(modelsArray) != 2 {
-			t.Errorf("Expected 2 models, got %d", len(modelsArray))
-		}
-
-		modelIDs := make(map[string]bool)
-		for _, m := range modelsArray {
-			modelObj := m.(map[string]any)
-			id := modelObj["id"].(string)
-			modelIDs[id] = true
-		}
-
-		if !modelIDs["new-model-1"] || !modelIDs["new-model-2"] {
-			t.Errorf("Expected new models, got %v", modelIDs)
-		}
-		if modelIDs["old-model-1"] || modelIDs["old-model-2"] {
-			t.Errorf("Old models should have been removed, got %v", modelIDs)
-		}
-	})
-
-	t.Run("handles partial overlap in model list", func(t *testing.T) {
-		cleanup()
-		os.MkdirAll(configDir, 0o755)
-
-		// Models must have _launch marker to be managed
-		existingConfig := `{
-			"providers": {
-				"ollama": {
-					"baseUrl": "http://localhost:11434/v1",
-					"api": "openai-completions",
-					"apiKey": "ollama",
-					"models": [
-						{"id": "keep-model", "_launch": true},
-						{"id": "remove-model", "_launch": true}
-					]
-				}
-			}
-		}`
-		if err := os.WriteFile(configPath, []byte(existingConfig), 0o644); err != nil {
-			t.Fatal(err)
-		}
-
-		newModels := []string{"keep-model", "add-model"}
-		if err := pi.Edit(newModels); err != nil {
-			t.Fatalf("Edit() error = %v", err)
-		}
-
-		cfg := readConfig()
-		providers := cfg["providers"].(map[string]any)
-		ollama := providers["ollama"].(map[string]any)
-		modelsArray := ollama["models"].([]any)
-
-		if len(modelsArray) != 2 {
-			t.Errorf("Expected 2 models, got %d", len(modelsArray))
-		}
-
-		modelIDs := make(map[string]bool)
-		for _, m := range modelsArray {
-			modelObj := m.(map[string]any)
-			id := modelObj["id"].(string)
-			modelIDs[id] = true
-		}
-
-		if !modelIDs["keep-model"] || !modelIDs["add-model"] {
-			t.Errorf("Expected keep-model and add-model, got %v", modelIDs)
-		}
-		if modelIDs["remove-model"] {
-			t.Errorf("remove-model should have been removed")
-		}
-	})
-
-	t.Run("handles corrupt config gracefully", func(t *testing.T) {
-		cleanup()
-		os.MkdirAll(configDir, 0o755)
-
-		if err := os.WriteFile(configPath, []byte("{invalid json}"), 0o644); err != nil {
-			t.Fatal(err)
-		}
-
-		models := []string{"test-model"}
-		if err := pi.Edit(models); err != nil {
-			t.Fatalf("Edit() should not fail with corrupt config, got %v", err)
-		}
-
-		data, err := os.ReadFile(configPath)
-		if err != nil {
-			t.Fatalf("Failed to read config: %v", err)
-		}
-
-		var cfg map[string]any
-		if err := json.Unmarshal(data, &cfg); err != nil {
-			t.Fatalf("Config should be valid after Edit, got parse error: %v", err)
-		}
-
-		providers := cfg["providers"].(map[string]any)
-		ollama := providers["ollama"].(map[string]any)
-		modelsArray := ollama["models"].([]any)
-
-		if len(modelsArray) != 1 {
-			t.Errorf("Expected 1 model, got %d", len(modelsArray))
-		}
-	})
-
-	// CRITICAL SAFETY TEST: verifies we don't stomp on user configs
-	t.Run("preserves user-managed models without _launch marker", func(t *testing.T) {
-		cleanup()
-		os.MkdirAll(configDir, 0o755)
-
-		// User has manually configured models in ollama provider (no _launch marker)
-		existingConfig := `{
-			"providers": {
-				"ollama": {
-					"baseUrl": "http://localhost:11434/v1",
-					"api": "openai-completions",
-					"apiKey": "ollama",
-					"models": [
-						{"id": "user-model-1"},
-						{"id": "user-model-2", "customField": "preserved"},
-						{"id": "ollama-managed", "_launch": true}
-					]
-				}
-			}
-		}`
-		if err := os.WriteFile(configPath, []byte(existingConfig), 0o644); err != nil {
-			t.Fatal(err)
-		}
-
-		// Add a new ollama-managed model
-		newModels := []string{"new-ollama-model"}
-		if err := pi.Edit(newModels); err != nil {
-			t.Fatalf("Edit() error = %v", err)
-		}
-
-		cfg := readConfig()
-		providers := cfg["providers"].(map[string]any)
-		ollama := providers["ollama"].(map[string]any)
-		modelsArray := ollama["models"].([]any)
-
-		// Should have: new-ollama-model (managed) + 2 user models (preserved)
-		if len(modelsArray) != 3 {
-			t.Errorf("Expected 3 models (1 new managed + 2 preserved user models), got %d", len(modelsArray))
-		}
-
-		modelIDs := make(map[string]map[string]any)
-		for _, m := range modelsArray {
-			modelObj := m.(map[string]any)
-			id := modelObj["id"].(string)
-			modelIDs[id] = modelObj
-		}
-
-		// Verify new model has _launch marker
-		if m, ok := modelIDs["new-ollama-model"]; !ok {
-			t.Errorf("new-ollama-model should be present")
-		} else if m["_launch"] != true {
-			t.Errorf("new-ollama-model should have _launch marker")
-		}
-
-		// Verify user models are preserved
-		if _, ok := modelIDs["user-model-1"]; !ok {
-			t.Errorf("user-model-1 should be preserved")
-		}
-		if _, ok := modelIDs["user-model-2"]; !ok {
-			t.Errorf("user-model-2 should be preserved")
-		} else if modelIDs["user-model-2"]["customField"] != "preserved" {
-			t.Errorf("user-model-2 customField should be preserved")
-		}
-
-		// Verify old ollama-managed model is removed (not in new list)
-		if _, ok := modelIDs["ollama-managed"]; ok {
-			t.Errorf("ollama-managed should be removed (old ollama model not in new selection)")
-		}
-	})
-
-	t.Run("updates settings.json with default provider and model", func(t *testing.T) {
-		cleanup()
-		os.MkdirAll(configDir, 0o755)
-
-		// Create existing settings with other fields
-		settingsPath := filepath.Join(configDir, "settings.json")
-		existingSettings := `{
-			"theme": "dark",
-			"customSetting": "value",
-			"defaultProvider": "anthropic",
-			"defaultModel": "claude-3"
-		}`
-		if err := os.WriteFile(settingsPath, []byte(existingSettings), 0o644); err != nil {
-			t.Fatal(err)
-		}
-
-		models := []string{"llama3.2"}
-		if err := pi.Edit(models); err != nil {
-			t.Fatalf("Edit() error = %v", err)
-		}
-
-		data, err := os.ReadFile(settingsPath)
-		if err != nil {
-			t.Fatalf("Failed to read settings: %v", err)
-		}
-
-		var settings map[string]any
-		if err := json.Unmarshal(data, &settings); err != nil {
-			t.Fatalf("Failed to parse settings: %v", err)
-		}
-
-		// Verify defaultProvider is set to ollama
-		if settings["defaultProvider"] != "ollama" {
-			t.Errorf("defaultProvider = %v, want ollama", settings["defaultProvider"])
-		}
-
-		// Verify defaultModel is set to first model
-		if settings["defaultModel"] != "llama3.2" {
-			t.Errorf("defaultModel = %v, want llama3.2", settings["defaultModel"])
-		}
-
-		// Verify other fields are preserved
-		if settings["theme"] != "dark" {
-			t.Errorf("theme = %v, want dark (preserved)", settings["theme"])
-		}
-		if settings["customSetting"] != "value" {
-			t.Errorf("customSetting = %v, want value (preserved)", settings["customSetting"])
-		}
-	})
-
-	t.Run("creates settings.json if it does not exist", func(t *testing.T) {
-		cleanup()
-		os.MkdirAll(configDir, 0o755)
-
-		models := []string{"qwen3:8b"}
-		if err := pi.Edit(models); err != nil {
-			t.Fatalf("Edit() error = %v", err)
-		}
-
-		settingsPath := filepath.Join(configDir, "settings.json")
-		data, err := os.ReadFile(settingsPath)
-		if err != nil {
-			t.Fatalf("settings.json should be created: %v", err)
-		}
-
-		var settings map[string]any
-		if err := json.Unmarshal(data, &settings); err != nil {
-			t.Fatalf("Failed to parse settings: %v", err)
-		}
-
-		if settings["defaultProvider"] != "ollama" {
-			t.Errorf("defaultProvider = %v, want ollama", settings["defaultProvider"])
-		}
-		if settings["defaultModel"] != "qwen3:8b" {
-			t.Errorf("defaultModel = %v, want qwen3:8b", settings["defaultModel"])
-		}
-	})
-
-	t.Run("handles corrupt settings.json gracefully", func(t *testing.T) {
-		cleanup()
-		os.MkdirAll(configDir, 0o755)
-
-		// Create corrupt settings
-		settingsPath := filepath.Join(configDir, "settings.json")
-		if err := os.WriteFile(settingsPath, []byte("{invalid"), 0o644); err != nil {
-			t.Fatal(err)
-		}
-
-		models := []string{"test-model"}
-		if err := pi.Edit(models); err != nil {
-			t.Fatalf("Edit() should not fail with corrupt settings, got %v", err)
-		}
-
-		data, err := os.ReadFile(settingsPath)
-		if err != nil {
-			t.Fatalf("Failed to read settings: %v", err)
-		}
-
-		var settings map[string]any
-		if err := json.Unmarshal(data, &settings); err != nil {
-			t.Fatalf("settings.json should be valid after Edit, got parse error: %v", err)
-		}
-
-		if settings["defaultProvider"] != "ollama" {
-			t.Errorf("defaultProvider = %v, want ollama", settings["defaultProvider"])
-		}
-		if settings["defaultModel"] != "test-model" {
-			t.Errorf("defaultModel = %v, want test-model", settings["defaultModel"])
-		}
-	})
-}
-
-func TestPiModels(t *testing.T) {
-	pi := &Pi{}
-
-	t.Run("returns nil when no config exists", func(t *testing.T) {
-		tmpDir := t.TempDir()
-		setTestHome(t, tmpDir)
-
-		models := pi.Models()
-		if models != nil {
-			t.Errorf("Models() = %v, want nil", models)
-		}
-	})
-
-	t.Run("returns models from config", func(t *testing.T) {
-		tmpDir := t.TempDir()
-		setTestHome(t, tmpDir)
-
-		configDir := filepath.Join(tmpDir, ".pi", "agent")
-		if err := os.MkdirAll(configDir, 0o755); err != nil {
-			t.Fatal(err)
-		}
-		config := `{
-			"providers": {
-				"ollama": {
-					"models": [
-						{"id": "llama3.2"},
-						{"id": "qwen3:8b"}
-					]
-				}
-			}
-		}`
-		configPath := filepath.Join(configDir, "models.json")
-		if err := os.WriteFile(configPath, []byte(config), 0o644); err != nil {
-			t.Fatal(err)
-		}
-
-		models := pi.Models()
-		if len(models) != 2 {
-			t.Errorf("Models() returned %d models, want 2", len(models))
-		}
-		if models[0] != "llama3.2" || models[1] != "qwen3:8b" {
-			t.Errorf("Models() = %v, want [llama3.2 qwen3:8b] (sorted)", models)
-		}
-	})
-
-	t.Run("returns sorted models", func(t *testing.T) {
-		tmpDir := t.TempDir()
-		setTestHome(t, tmpDir)
-
-		configDir := filepath.Join(tmpDir, ".pi", "agent")
-		if err := os.MkdirAll(configDir, 0o755); err != nil {
-			t.Fatal(err)
-		}
-		config := `{
-			"providers": {
-				"ollama": {
-					"models": [
-						{"id": "z-model"},
-						{"id": "a-model"},
-						{"id": "m-model"}
-					]
-				}
-			}
-		}`
-		configPath := filepath.Join(configDir, "models.json")
-		if err := os.WriteFile(configPath, []byte(config), 0o644); err != nil {
-			t.Fatal(err)
-		}
-
-		models := pi.Models()
-		if models[0] != "a-model" || models[1] != "m-model" || models[2] != "z-model" {
-			t.Errorf("Models() = %v, want [a-model m-model z-model] (sorted)", models)
-		}
-	})
-
-	t.Run("returns nil when models array is missing", func(t *testing.T) {
-		tmpDir := t.TempDir()
-		setTestHome(t, tmpDir)
-
-		configDir := filepath.Join(tmpDir, ".pi", "agent")
-		if err := os.MkdirAll(configDir, 0o755); err != nil {
-			t.Fatal(err)
-		}
-		config := `{
-			"providers": {
-				"ollama": {}
-			}
-		}`
-		configPath := filepath.Join(configDir, "models.json")
-		if err := os.WriteFile(configPath, []byte(config), 0o644); err != nil {
-			t.Fatal(err)
-		}
-
-		models := pi.Models()
-		if models != nil {
-			t.Errorf("Models() = %v, want nil when models array is missing", models)
-		}
-	})
-
-	t.Run("handles corrupt config gracefully", func(t *testing.T) {
-		tmpDir := t.TempDir()
-		setTestHome(t, tmpDir)
-
-		configDir := filepath.Join(tmpDir, ".pi", "agent")
-		if err := os.MkdirAll(configDir, 0o755); err != nil {
-			t.Fatal(err)
-		}
-		configPath := filepath.Join(configDir, "models.json")
-		if err := os.WriteFile(configPath, []byte("{invalid json}"), 0o644); err != nil {
-			t.Fatal(err)
-		}
-
-		models := pi.Models()
-		if models != nil {
-			t.Errorf("Models() = %v, want nil for corrupt config", models)
-		}
-	})
-}

integration/basic_test.go

@@ -144,3 +144,47 @@ func TestUnicodeModelDir(t *testing.T) {
 	}
 	ChatTestHelper(ctx, t, req, blueSkyExpected)
 }
+
+// TestNumPredict verifies that when num_predict is set, the model generates
+// exactly that many tokens. It uses logprobs to count the actual tokens output.
+func TestNumPredict(t *testing.T) {
+	ctx, cancel := context.WithTimeout(context.Background(), 2*time.Minute)
+	defer cancel()
+
+	client, _, cleanup := InitServerConnection(ctx, t)
+	defer cleanup()
+
+	if err := PullIfMissing(ctx, client, "qwen3:0.6b"); err != nil {
+		t.Fatalf("failed to pull model: %v", err)
+	}
+
+	req := api.GenerateRequest{
+		Model:    "qwen3:0.6b",
+		Prompt:   "Write a long story.",
+		Stream:   &stream,
+		Logprobs: true,
+		Options: map[string]any{
+			"num_predict": 10,
+			"temperature": 0,
+			"seed":        123,
+		},
+	}
+
+	logprobCount := 0
+	var finalResponse api.GenerateResponse
+	err := client.Generate(ctx, &req, func(resp api.GenerateResponse) error {
+		logprobCount += len(resp.Logprobs)
+		if resp.Done {
+			finalResponse = resp
+		}
+		return nil
+	})
+	if err != nil {
+		t.Fatalf("generate failed: %v", err)
+	}
+
+	if logprobCount != 10 {
+		t.Errorf("expected 10 tokens (logprobs), got %d (EvalCount=%d, DoneReason=%s)",
+			logprobCount, finalResponse.EvalCount, finalResponse.DoneReason)
+	}
+}

ml/backend.go

@@ -175,6 +175,7 @@ type Tensor interface {
 	SILU(ctx Context, up ...Tensor) Tensor
 	RELU(ctx Context, up ...Tensor) Tensor
 	Sigmoid(ctx Context) Tensor
+	SigmoidOut(ctx Context) Tensor
 
 	// AlphaLimitSILU is a variant of SILU that clamps the input to the range [-limit, limit]
 	SILUAlphaLimit(ctx Context, up Tensor, alpha, limit float32) Tensor

ml/backend/ggml/ggml.go

@@ -1468,6 +1468,13 @@ func (t *Tensor) Sigmoid(ctx ml.Context) ml.Tensor {
 	}
 }
 
+func (t *Tensor) SigmoidOut(ctx ml.Context) ml.Tensor {
+	return &Tensor{
+		b: t.b,
+		t: C.ggml_sigmoid(ctx.(*Context).ctx, t.t),
+	}
+}
+
 func (t *Tensor) View(ctx ml.Context, offset int, shape ...int) ml.Tensor {
 	switch len(shape) {
 	case 1:

model/models/qwen3next/model.go

@@ -135,7 +135,7 @@ func (mlp *sparse) Forward(ctx ml.Context, hiddenStates ml.Tensor, opts *Options
 		// Apply shared expert gating
 		if mlp.SharedGateInp != nil {
 			sharedGateVal := mlp.SharedGateInp.Forward(ctx, hiddenStates2D)
-			sharedGateVal = sharedGateVal.Sigmoid(ctx)
+			sharedGateVal = sharedGateVal.SigmoidOut(ctx)
 			// Broadcast gate to match dimensions
 			sharedGateVal = sharedGateVal.Repeat(ctx, 0, sharedOut.Dim(0))
 			sharedOut = sharedOut.Mul(ctx, sharedGateVal)

runner/ollamarunner/runner.go

@@ -514,13 +514,6 @@ func (s *Server) forwardBatch(pendingBatch batchState) (nextBatch batchState, er
 			continue
 		}
 
-		// if past the num predict limit
-		if seq.numPredict > 0 && seq.numPredicted >= seq.numPredict {
-			s.removeSequence(seqIdx, llm.DoneReasonLength)
-			nextBatch.seqs[seqIdx] = nil
-			continue
-		}
-
 		if !s.cache.enabled {
 			seq.inputs = append(seq.cache.Inputs, seq.inputs...)
 			seq.cache.Inputs = []*input.Input{}
@@ -709,7 +702,6 @@ func (s *Server) computeBatch(activeBatch batchState) {
 			continue
 		}
 
-		seq.numPredicted++
 		nextToken := &input.Input{Token: 0} // placeholder we'll fill in after Compute/Floats
 		seq.inputs = []*input.Input{nextToken}
 		nextBatchTokens[i] = nextToken
@@ -745,7 +737,9 @@ func (s *Server) computeBatch(activeBatch batchState) {
 			logutil.Trace("computeBatch: sequence replaced, discarding its results", "batchID", activeBatch.id, "seqIdx", i)
 			continue
 		}
+
 		seq.lastUpdatedAt = t
+		seq.numPredicted++
 		if seq.numPredicted == 1 {
 			seq.processingDuration = seq.lastUpdatedAt.Sub(seq.startedAt)
 			seq.startedAt = seq.lastUpdatedAt
@@ -791,6 +785,13 @@ func (s *Server) computeBatch(activeBatch batchState) {
 		}
 
 		seq.pendingResponses = append(seq.pendingResponses, piece)
+
+		// if past the num predict limit
+		if seq.numPredict > 0 && seq.numPredicted >= seq.numPredict {
+			s.removeSequence(i, llm.DoneReasonLength)
+			continue
+		}
+
 		sequence := strings.Join(seq.pendingResponses, "")
 
 		if ok, stop := common.FindStop(sequence, seq.stop); ok {