create: select the qwen3.5 parser and renderer for Qwen3.5/Next

Qwen3.5/Qwen3-Next architecture strings contain the substring "qwen3", so the
broad qwen3 match claimed them for the generic parser and qwen3-coder
renderer, whose template doesn't frame the thinking block — an empty
<think></think> leaked into content and think=false was ignored. Match the
family first via isQwen35Family so the parser, renderer, and
thinking-capability checks share one variant list.

LLAMA_CPP_VERSION

@@ -1 +1 @@
-b9781
+b9888

MLX_VERSION

@@ -1 +1 @@
-51b2768da7e1897d3c4258f7ddbb47083d1eef01
+de7b4ed986b6d6f55b8ace5e73c24d1ca0bea89b

agent/compactor.go

@@ -0,0 +1,629 @@
+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
+
+	maxCompactionSummaryBytes  = 16 * 1024
+	compactionSummaryTruncated = "\n\n[summary truncated]"
+
+	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)
+}
+
+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)
+}
+
+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 (c *SimpleCompactor) threshold() float64 {
+	return ResolveCompactionThreshold(c.Options.Threshold)
+}
+
+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 {
+	if len(summary) <= maxCompactionSummaryBytes {
+		return summary
+	}
+	limit := maxCompactionSummaryBytes - len(compactionSummaryTruncated)
+	if limit < 0 {
+		limit = 0
+	}
+	var b strings.Builder
+	for _, r := range summary {
+		if b.Len()+len(string(r)) > limit {
+			break
+		}
+		b.WriteRune(r)
+	}
+	return strings.TrimSpace(b.String()) + compactionSummaryTruncated
+}
+
+func estimateCompactionTokens(text string) int {
+	text = strings.TrimSpace(text)
+	if text == "" {
+		return 0
+	}
+	return max(1, (len([]rune(text))+3)/4)
+}
+
+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
+	}
+}

agent/compactor_test.go

@@ -0,0 +1,773 @@
+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", maxCompactionSummaryBytes+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 len(result.Summary) > maxCompactionSummaryBytes {
+		t.Fatalf("summary bytes = %d, want <= %d", len(result.Summary), maxCompactionSummaryBytes)
+	}
+	if !strings.HasSuffix(result.Summary, compactionSummaryTruncated) {
+		t.Fatalf("summary missing truncation marker")
+	}
+	if !strings.Contains(result.Messages[1].Content, compactionSummaryTruncated) {
+		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)
+	}
+}

agent/events.go

@@ -0,0 +1,79 @@
+package agent
+
+import (
+	"context"
+
+	"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"
+)
+
+type Event struct {
+	Type       EventType      `json:"type"`
+	RunID      string         `json:"runId,omitempty"`
+	ChatID     string         `json:"chatId,omitempty"`
+	Model      string         `json:"model,omitempty"`
+	Status     string         `json:"status,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)
+}
+
+func (s *Session) emit(event Event) error {
+	if s == nil {
+		return nil
+	}
+	var firstErr error
+	for _, sink := range s.EventSinks {
+		if sink == nil {
+			continue
+		}
+		if err := sink.Emit(event); err != nil && firstErr == nil {
+			firstErr = err
+		}
+	}
+	return firstErr
+}
+
+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
+}

agent/registry.go

@@ -0,0 +1,94 @@
+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
+}
+
+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 {
+	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
+}

agent/session.go

@@ -0,0 +1,1024 @@
+package agent
+
+import (
+	"context"
+	"encoding/json"
+	"errors"
+	"fmt"
+	"os"
+	"path/filepath"
+	"strings"
+
+	"github.com/google/uuid"
+
+	"github.com/ollama/ollama/api"
+)
+
+type ChatClient interface {
+	Chat(context.Context, *api.ChatRequest, api.ChatResponseFunc) error
+}
+
+type ApprovalRequest struct {
+	WorkingDir string
+	Calls      []ApprovalToolCall
+}
+
+type ApprovalToolCall struct {
+	ToolCallID string
+	ToolName   string
+	Args       map[string]any
+}
+
+type Approval struct {
+	Allow    bool
+	AllowAll bool
+	Reason   string
+}
+
+type ApprovalPrompter interface {
+	PromptApproval(context.Context, ApprovalRequest) (Approval, error)
+}
+
+type Session struct {
+	Client           ChatClient
+	EventSinks       []EventSink
+	Tools            *Registry
+	ApprovalPrompter ApprovalPrompter
+	AllowAllTools    bool
+	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
+	// MaxToolRounds limits consecutive model/tool cycles.
+	// Zero uses the default guard; negative 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
+}
+
+type toolExecutionStop string
+
+const (
+	toolExecutionDenied   toolExecutionStop = "denied"
+	toolExecutionCanceled toolExecutionStop = "canceled"
+)
+
+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         string
+	ignoreCanceled bool
+	err            error
+}
+
+func (st *runState) finishDone() {
+	st.finish = runFinish{status: "done"}
+	st.phase = runPhaseDone
+}
+
+func (st *runState) finishDenied() {
+	st.finish = runFinish{status: "denied"}
+	st.phase = runPhaseDone
+}
+
+func (st *runState) finishCanceled() {
+	st.finish = runFinish{status: "canceled", 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 s == nil {
+		return nil, errors.New("nil session")
+	}
+	if s.Client == nil {
+		return nil, errors.New("agent session requires a chat client")
+	}
+	if opts.Model == "" {
+		return nil, errors.New("agent session requires a model")
+	}
+	runID := uuid.NewString()
+	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(Event{Type: EventError, RunID: runID, ChatID: opts.ChatID, Model: opts.Model, Error: err.Error()})
+		return nil, err
+	}
+
+	st := runState{
+		runID:         runID,
+		opts:          opts,
+		phase:         runPhaseModel,
+		messages:      messages,
+		maxToolRounds: resolvedMaxToolRounds(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 &RunResult{Messages: st.messages, Latest: st.latest, WorkingDir: s.WorkingDir}, err
+			}
+		case runPhaseDone:
+			return s.finishRun(ctx, &st)
+		}
+	}
+}
+
+func (s *Session) runModelStep(ctx context.Context, st *runState) error {
+	opts := st.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(Event{Type: EventError, RunID: st.runID, ChatID: opts.ChatID, Model: opts.Model, Error: 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(Event{Type: EventError, RunID: st.runID, ChatID: opts.ChatID, Model: opts.Model, Error: skipErr.Error()})
+			return skipErr
+		}
+		st.messages = append(st.messages, skipped...)
+		st.finishCanceled()
+		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(Event{Type: EventError, RunID: st.runID, ChatID: opts.ChatID, Model: opts.Model, Error: 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(Event{Type: EventError, RunID: st.runID, ChatID: opts.ChatID, Model: opts.Model, Error: 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(Event{Type: EventError, RunID: st.runID, ChatID: st.opts.ChatID, Model: st.opts.Model, Error: 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
+	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(Event{Type: EventError, RunID: st.runID, ChatID: opts.ChatID, Model: opts.Model, Error: err.Error()})
+		return err
+	}
+
+	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 := Event{Type: EventRunFinished, RunID: st.runID, ChatID: st.opts.ChatID, Model: st.opts.Model, Status: 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) {
+	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 tools := s.availableTools(); len(tools) > 0 {
+		req.Tools = 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 response.Message.Content == "" && response.Message.Thinking == "" && len(response.Message.ToolCalls) == 0 {
+			*latest = response
+			return nil
+		}
+
+		if response.Message.Thinking != "" {
+			assistant.Thinking += response.Message.Thinking
+			if err := s.emit(Event{Type: EventThinkingDelta, RunID: runID, ChatID: opts.ChatID, Model: opts.Model, Thinking: response.Message.Thinking}); err != nil {
+				return err
+			}
+		}
+
+		if response.Message.Content != "" {
+			assistant.Content += response.Message.Content
+			if err := s.emit(Event{Type: EventMessageDelta, RunID: runID, ChatID: opts.ChatID, Model: opts.Model, Content: 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(Event{Type: EventToolCallDetected, RunID: runID, ChatID: opts.ChatID, Model: opts.Model, ToolCalls: 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 (s *Session) executeToolCalls(ctx context.Context, runID string, opts RunOptions, messages []api.Message, calls []api.ToolCall) (toolBatchResult, error) {
+	batch := toolBatchResult{
+		messages: make([]api.Message, 0, len(calls)),
+	}
+	projectedMessages := append([]api.Message(nil), messages...)
+
+	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 && ToolRequiresApproval(tool, args) {
+			approvalReq.Calls = append(approvalReq.Calls, ApprovalToolCall{
+				ToolCallID: call.ID,
+				ToolName:   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, projectedMessages)
+				batch.messages = append(batch.messages, msg)
+				projectedMessages = append(projectedMessages, msg)
+				deniedContent := msg.Content
+				if emitErr := s.emit(Event{Type: EventToolFinished, RunID: runID, ChatID: opts.ChatID, Model: opts.Model, Status: "denied", ToolCallID: plan.call.ID, ToolName: plan.toolName, Args: plan.args, Content: deniedContent, Error: 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, projectedMessages)
+			batch.messages = append(batch.messages, msg)
+			projectedMessages = append(projectedMessages, 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(Event{Type: EventToolFinished, RunID: runID, ChatID: opts.ChatID, Model: opts.Model, Status: "failed", ToolCallID: call.ID, ToolName: toolName, Args: args, Content: content, Error: fmt.Sprintf("unknown tool: %s", toolName)}); emitErr != nil {
+				return toolBatchResult{}, emitErr
+			}
+			continue
+		}
+
+		if err := s.emit(Event{Type: EventToolStarted, RunID: runID, ChatID: opts.ChatID, Model: opts.Model, Status: "running", ToolCallID: call.ID, ToolName: toolName, WorkingDir: plan.workingDir, Args: 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, projectedMessages)
+			batch.messages = append(batch.messages, msg)
+			projectedMessages = append(projectedMessages, 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, Event{Type: EventToolFinished, RunID: runID, ChatID: opts.ChatID, Model: opts.Model, Status: "failed", ToolCallID: call.ID, ToolName: toolName, Args: args, Content: content, Error: 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, projectedMessages)
+		batch.messages = append(batch.messages, msg)
+		projectedMessages = append(projectedMessages, 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, Event{Type: EventToolFinished, RunID: runID, ChatID: opts.ChatID, Model: opts.Model, Status: "done", ToolCallID: call.ID, ToolName: toolName, WorkingDir: eventWorkingDir, Args: args, Content: 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) authorizeToolCalls(ctx context.Context, req ApprovalRequest) (Approval, error) {
+	if s == nil || s.AllowAllTools || len(req.Calls) == 0 {
+		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
+	}
+	if result.AllowAll {
+		result.Allow = true
+		s.AllowAllTools = true
+	}
+	return result, nil
+}
+
+func (s *Session) skipToolCalls(ctx context.Context, runID string, opts RunOptions, calls []api.ToolCall, content string) ([]api.Message, error) {
+	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, Event{Type: EventToolFinished, RunID: runID, ChatID: opts.ChatID, Model: opts.Model, Status: "skipped", ToolCallID: call.ID, ToolName: toolName, Args: args, Content: msg.Content, Error: 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 = "error"
+			}
+			s.emitCompactionSkipped(runID, opts, trigger, result.Reason)
+			skipNotified = true
+		}
+		return messages, skipNotified, nil
+	}
+	if !result.Compacted {
+		if result.Due && !skipNotified {
+			if trigger == "" {
+				trigger = "due"
+			}
+			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, "tool_output")
+
+	result, err := s.Compactor.MaybeCompact(ctx, req)
+	if err != nil {
+		if result.Due && !skipNotified {
+			s.emitCompactionSkipped(runID, opts, "tool_output", result.Reason)
+			skipNotified = true
+		}
+		return messages, skipNotified, nil
+	}
+	if !result.Compacted {
+		if result.Due && !skipNotified {
+			s.emitCompactionSkipped(runID, opts, "tool_output", 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)
+	}
+
+	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, compacted)
+		compacted = append(compacted, refit)
+	}
+
+	s.emitCompacted(runID, opts, compacted, "tool_output", 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 {
+	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(Event{Type: EventCompactionProgress, RunID: runID, ChatID: opts.ChatID, Model: opts.Model, Tokens: progress.Tokens})
+		},
+	}
+}
+
+func (s *Session) emitCompactionStarted(runID string, opts RunOptions, status string) {
+	_ = s.emit(Event{
+		Type:   EventCompactionStarted,
+		RunID:  runID,
+		ChatID: opts.ChatID,
+		Model:  opts.Model,
+		Status: status,
+	})
+}
+
+func (s *Session) emitCompactionSkipped(runID string, opts RunOptions, status, reason string) {
+	_ = s.emit(Event{
+		Type:    EventCompactionSkipped,
+		RunID:   runID,
+		ChatID:  opts.ChatID,
+		Model:   opts.Model,
+		Status:  status,
+		Content: CompactionSkippedMessage(reason),
+	})
+}
+
+func (s *Session) emitCompacted(runID string, opts RunOptions, messages []api.Message, status, summary string) {
+	_ = s.emit(Event{
+		Type:     EventCompacted,
+		RunID:    runID,
+		ChatID:   opts.ChatID,
+		Model:    opts.Model,
+		Status:   status,
+		Content:  summary,
+		Messages: messages,
+	})
+}
+
+func (s *Session) autoCompactionTrigger(req CompactionRequest) string {
+	if compactor, ok := s.Compactor.(*SimpleCompactor); ok && compactor != nil {
+		if req.Force {
+			return "force"
+		}
+		contextWindow := compactor.contextWindowTokens(req.Options)
+		threshold := int(float64(contextWindow) * compactor.threshold())
+		if threshold <= 0 {
+			return ""
+		}
+		if req.Latest.PromptEvalCount > 0 && req.Latest.PromptEvalCount >= threshold {
+			return "prompt_eval"
+		}
+		if estimateCompactionRequestTokens(req) >= threshold {
+			return "estimate"
+		}
+	}
+	return ""
+}
+
+func CompactionSkippedMessage(reason string) string {
+	reason = strings.TrimSpace(reason)
+	if reason == "" {
+		reason = "compaction could not run"
+	}
+	return reason
+}
+
+func resolvedMaxToolRounds(value int) int {
+	if value == 0 {
+		return defaultMaxToolRounds
+	}
+	return value
+}
+
+func (s *Session) toolMessageForContext(toolName, toolCallID, content string, opts RunOptions, messages []api.Message) api.Message {
+	maxRunes := maxToolResultRunes
+	if limit := smallContextToolResultLimitRunes(s.contextWindowTokens(opts)); limit > 0 {
+		maxRunes = min(maxRunes, limit)
+	}
+
+	msg := toolMessageWithLimit(toolName, toolCallID, content, maxRunes)
+	threshold := s.compactionThresholdTokens(opts)
+	if threshold <= 0 {
+		return msg
+	}
+
+	projected := append(append([]api.Message(nil), messages...), msg)
+	projectedTokens := s.estimateRunPromptTokens(opts, projected)
+	if projectedTokens < threshold {
+		return msg
+	}
+
+	baseTokens := s.estimateRunPromptTokens(opts, messages)
+	overheadTokens := estimateMessagesTokens([]api.Message{{
+		Role:       "tool",
+		ToolName:   toolName,
+		ToolCallID: toolCallID,
+	}})
+	// Keep oversized tool output below the compaction threshold 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 := (threshold - baseTokens - overheadTokens - toolTruncationMarkerReserveTokens) * 4
+	maxRunes = min(maxRunes, max(0, availableRunes))
+	msg.Content = truncateToolResultContentTo(content, maxRunes)
+	return msg
+}
+
+func (s *Session) toolMessageForPostCompactionContext(toolName, toolCallID, content string, opts RunOptions, messages []api.Message) api.Message {
+	maxRunes := maxToolResultRunes
+	if limit := smallContextToolResultLimitRunes(s.contextWindowTokens(opts)); limit > 0 {
+		maxRunes = min(maxRunes, limit)
+	}
+
+	contextWindow := s.contextWindowTokens(opts)
+	if contextWindow <= 0 {
+		return toolMessageWithLimit(toolName, toolCallID, content, maxRunes)
+	}
+
+	baseTokens := s.estimateRunPromptTokens(opts, messages)
+	overheadTokens := estimateMessagesTokens([]api.Message{{
+		Role:       "tool",
+		ToolName:   toolName,
+		ToolCallID: toolCallID,
+	}})
+	availableRunes := (contextWindow - baseTokens - overheadTokens - toolTruncationMarkerReserveTokens) * 4
+	maxRunes = min(maxRunes, max(0, availableRunes))
+	return toolMessageWithLimit(toolName, toolCallID, content, maxRunes)
+}
+
+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 compactor, ok := s.Compactor.(*SimpleCompactor); ok && compactor != nil {
+		configuredThreshold = compactor.Options.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
+	}
+
+	configuredWindow := 0
+	if compactor, ok := s.Compactor.(*SimpleCompactor); ok && compactor != nil {
+		configuredWindow = compactor.Options.ContextWindowTokens
+	}
+
+	return ResolveContextWindowTokens(opts.Options, configuredWindow)
+}
+
+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] = sanitizeMessageForRequest(msg)
+	}
+	return sanitized
+}
+
+func sanitizeMessageForRequest(msg api.Message) api.Message {
+	return sanitizeMessageForRun(msg)
+}
+
+func truncateToolResultContent(content string) string {
+	return truncateToolResultContentTo(content, maxToolResultRunes)
+}
+
+func truncateToolResultContentTo(content string, maxRunes int) string {
+	if maxRunes <= 0 {
+		return fmt.Sprintf("%s omitted ~%d tokens. Use a narrower command, line range, or search query if more detail is needed.]", toolOutputFullOmissionPrefix, approximateTokensFromRunes(len([]rune(content))))
+	}
+	if len(content) <= maxRunes {
+		return content
+	}
+	runes := []rune(content)
+	if len(runes) <= maxRunes {
+		return content
+	}
+	head := maxRunes * 3 / 4
+	tail := maxRunes - head
+	omitted := len(runes) - head - tail
+	marker := fmt.Sprintf(
+		"\n\n[tool output truncated: showing first ~%d tokens and last ~%d tokens; omitted ~%d tokens. Use a narrower command, line range, or search query if more detail is needed.]\n\n",
+		approximateTokensFromRunes(head),
+		approximateTokensFromRunes(tail),
+		approximateTokensFromRunes(omitted),
+	)
+	return string(runes[:head]) + marker + string(runes[len(runes)-tail:])
+}
+
+func toolOutputFullyOmitted(content string) bool {
+	return strings.HasPrefix(content, toolOutputFullOmissionPrefix)
+}
+
+func approximateTokensFromRunes(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
+}

agent/session_test.go

@@ -0,0 +1,1826 @@
+package agent
+
+import (
+	"context"
+	"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 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 *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
+}
+
+type recordingApprovalPrompter struct {
+	requests []ApprovalRequest
+	results  []Approval
+}
+
+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 (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,
+		AllowAllTools: 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 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 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,
+		AllowAllTools: 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))
+	}
+
+	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 TestSessionCancellationAfterToolCallAppendsSkippedToolMessage(t *testing.T) {
+	ctx, cancel := context.WithCancel(context.Background())
+	registry := &Registry{}
+	registry.Register(staticTool{})
+	session := &Session{
+		Client:        cancelAfterToolCallClient{cancel: cancel},
+		Tools:         registry,
+		AllowAllTools: true,
+	}
+
+	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,
+		AllowAllTools: true,
+		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 != "canceled" {
+		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,
+		AllowAllTools: true,
+	}
+
+	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 TestSessionToolRoundLimitAppendsSkippedToolMessages(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,
+		AllowAllTools: true,
+	}
+
+	result, err := session.Run(context.Background(), RunOptions{
+		ChatID:        "chat-1",
+		Model:         "model",
+		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 TestSessionToolLoopStopsAtDefaultRoundCap(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,
+		AllowAllTools: true,
+	}
+
+	_, err := session.Run(context.Background(), RunOptions{
+		Model:       "model",
+		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 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,
+		AllowAllTools: true,
+	}
+
+	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,
+		AllowAllTools: 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(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,
+		AllowAllTools: true,
+		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,
+		AllowAllTools: true,
+		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,
+		AllowAllTools: true,
+		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,
+		AllowAllTools: true,
+		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,
+		AllowAllTools: true,
+		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,
+		AllowAllTools: true,
+		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,
+		AllowAllTools: true,
+		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,
+		AllowAllTools: true,
+		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 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.AllowAllTools {
+		t.Fatal("session did not remember allow all")
+	}
+	if len(prompter.requests) != 1 {
+		t.Fatalf("approval prompts = %d, want 1", len(prompter.requests))
+	}
+}
+
+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,
+		AllowAllTools: true,
+	}
+
+	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)
+	}
+}

agent/tools/bash.go

@@ -0,0 +1,388 @@
+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
+}
+
+func (b *Bash) Execute(ctx context.Context, toolCtx agent.ToolContext, args map[string]any) (agent.ToolResult, error) {
+	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
+}
+
+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 {
+	fields := shellSafetyFields(command)
+	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
+}
+
+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",
+		"/.aws/credentials",
+		"/.config/gcloud/application_default_credentials.json",
+		"/.kube/config",
+		"/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
+		}
+	}
+	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 + fmt.Sprintf("\n\n[%s truncated: omitted ~%d tokens]", label, approximateTokensFromBytes(omitted))
+}
+
+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)
+}
+
+func approximateTokensFromBytes(n int) int {
+	if n <= 0 {
+		return 0
+	}
+	return max(1, (n+3)/4)
+}

agent/tools/bash_test.go

@@ -0,0 +1,246 @@
+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: 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: "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},
+	}
+
+	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)
+	}
+}

agent/tools/bash_unix.go

@@ -0,0 +1,49 @@
+//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
+}

agent/tools/bash_unix_test.go

@@ -0,0 +1,40 @@
+//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)
+	}
+}

agent/tools/bash_windows.go

@@ -0,0 +1,134 @@
+//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)
+	}
+}

agent/tools/bash_windows_test.go

@@ -0,0 +1,15 @@
+//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)
+	}
+}

agent/tools/file.go

@@ -0,0 +1,521 @@
+package tools
+
+import (
+	"bufio"
+	"context"
+	"fmt"
+	"io"
+	"os"
+	"path/filepath"
+	"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) {
+	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)
+	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."
+}
+
+func (e *Edit) Schema() api.ToolFunction {
+	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("old_text", api.ToolProperty{
+		Type:        api.PropertyType{"string"},
+		Description: "Exact text to replace.",
+	})
+	props.Set("new_text", api.ToolProperty{
+		Type:        api.PropertyType{"string"},
+		Description: "Replacement text.",
+	})
+	props.Set("replace_all", api.ToolProperty{
+		Type:        api.PropertyType{"boolean"},
+		Description: "Replace every occurrence. Defaults to false and requires old_text to match exactly once.",
+	})
+	return api.ToolFunction{
+		Name:        e.Name(),
+		Description: e.Description(),
+		Parameters: api.ToolFunctionParameters{
+			Type:       "object",
+			Properties: props,
+			Required:   []string{"path", "old_text", "new_text"},
+		},
+	}
+}
+
+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) {
+	path, ok := args["path"].(string)
+	if !ok || strings.TrimSpace(path) == "" {
+		return agent.ToolResult{}, fmt.Errorf("path parameter is required")
+	}
+
+	oldText, ok := args["old_text"].(string)
+	if !ok || oldText == "" {
+		return agent.ToolResult{}, fmt.Errorf("old_text parameter is required")
+	}
+
+	newText, ok := args["new_text"].(string)
+	if !ok {
+		return agent.ToolResult{}, fmt.Errorf("new_text parameter is required")
+	}
+
+	replaceAll, _ := args["replace_all"].(bool)
+
+	if err := rejectFinalSymlink(toolCtx.WorkingDir, path); err != nil {
+		return agent.ToolResult{}, err
+	}
+
+	file, info, err := openRegularFile(toolCtx.WorkingDir, path)
+	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)
+	matches := strings.Count(content, oldText)
+	if matches == 0 {
+		return agent.ToolResult{}, fmt.Errorf("old_text was not found in %s", path)
+	}
+	if matches > 1 && !replaceAll {
+		return agent.ToolResult{}, fmt.Errorf("old_text matched %d times in %s; set replace_all to true to replace every match", matches, path)
+	}
+
+	var updated string
+	if replaceAll {
+		updated = strings.ReplaceAll(content, oldText, newText)
+	} else {
+		updated = strings.Replace(content, oldText, newText, 1)
+	}
+	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 replacement%s).", path, matches, plural(matches))}, nil
+}
+
+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) (*os.File, os.FileInfo, error) {
+	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)
+	}
+	if info.Mode()&os.ModeSymlink != 0 {
+		info, err = root.Stat(rel)
+		if err != nil {
+			return nil, rootPathError(err)
+		}
+	}
+	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"
+}

agent/tools/file_test.go

@@ -0,0 +1,297 @@
+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 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 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)
+	}
+}

agent/tools/file_unix_test.go

@@ -0,0 +1,75 @@
+//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")
+		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)
+	}
+}

agent/tools/web.go

@@ -0,0 +1,195 @@
+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"
+)
+
+var (
+	ErrWebSearchAuthRequired = errors.New("web search requires authentication")
+	ErrWebFetchAuthRequired  = errors.New("web fetch requires authentication")
+)
+
+const (
+	maxWebFetchContentRunes = 60_000
+	webSearchTimeout        = 15 * time.Second
+	webFetchTimeout         = 30 * time.Second
+)
+
+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) {
+	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{}, ErrWebSearchAuthRequired
+		}
+		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) {
+	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")
+	}
+	if _, err := url.Parse(urlStr); err != nil {
+		return agent.ToolResult{}, fmt.Errorf("invalid URL: %w", err)
+	}
+
+	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{}, ErrWebFetchAuthRequired
+		}
+		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 {
+	runes := []rune(content)
+	if len(runes) <= maxWebFetchContentRunes {
+		return content
+	}
+	omitted := len(runes) - maxWebFetchContentRunes
+	return string(runes[:maxWebFetchContentRunes]) + fmt.Sprintf(
+		"\n\n[tool output truncated: showing first ~%d tokens; omitted ~%d tokens. Use a narrower request or search query if more detail is needed.]",
+		approximateToolTokensFromRunes(maxWebFetchContentRunes),
+		approximateToolTokensFromRunes(omitted),
+	)
+}
+
+func approximateToolTokensFromRunes(n int) int {
+	if n <= 0 {
+		return 0
+	}
+	return max(1, (n+3)/4)
+}

agent/tools/web_test.go

@@ -0,0 +1,59 @@
+package tools
+
+import (
+	"encoding/json"
+	"net/http"
+	"net/http/httptest"
+	"strings"
+	"testing"
+
+	coreagent "github.com/ollama/ollama/agent"
+	"github.com/ollama/ollama/api"
+)
+
+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")
+	}
+}
+
+func TestWebFetchBoundsContentBeforeReturning(t *testing.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)
+	}
+}

api/client.go

@@ -473,6 +473,26 @@ func (c *Client) CloudStatusExperimental(ctx context.Context) (*StatusResponse,
 	return &status, nil
 }
 
+// WebSearchExperimental searches the web through the local server's
+// experimental web search endpoint.
+func (c *Client) WebSearchExperimental(ctx context.Context, req *WebSearchRequest) (*WebSearchResponse, error) {
+	var resp WebSearchResponse
+	if err := c.do(ctx, http.MethodPost, "/api/experimental/web_search", req, &resp); err != nil {
+		return nil, err
+	}
+	return &resp, nil
+}
+
+// WebFetchExperimental fetches web page content through the local server's
+// experimental web fetch endpoint.
+func (c *Client) WebFetchExperimental(ctx context.Context, req *WebFetchRequest) (*WebFetchResponse, error) {
+	var resp WebFetchResponse
+	if err := c.do(ctx, http.MethodPost, "/api/experimental/web_fetch", req, &resp); err != nil {
+		return nil, err
+	}
+	return &resp, nil
+}
+
 // Signout will signout a client for a local ollama server.
 func (c *Client) Signout(ctx context.Context) error {
 	return c.do(ctx, http.MethodPost, "/api/signout", nil, nil)

api/client_test.go

@@ -351,6 +351,82 @@ func TestClientDo(t *testing.T) {
 	}
 }
 
+func TestClientWebSearchExperimentalUsesLocalRoute(t *testing.T) {
+	var gotPath string
+	var gotMethod string
+	var gotRequest WebSearchRequest
+	ts := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		gotPath = r.URL.Path
+		gotMethod = r.Method
+		if err := json.NewDecoder(r.Body).Decode(&gotRequest); err != nil {
+			t.Fatal(err)
+		}
+		if err := json.NewEncoder(w).Encode(WebSearchResponse{
+			Results: []WebSearchResult{{Title: "Ollama", URL: "https://ollama.com", Content: "models"}},
+		}); err != nil {
+			t.Fatal(err)
+		}
+	}))
+	defer ts.Close()
+
+	client := NewClient(&url.URL{Scheme: "http", Host: ts.Listener.Addr().String()}, http.DefaultClient)
+	resp, err := client.WebSearchExperimental(t.Context(), &WebSearchRequest{Query: "ollama", MaxResults: 3})
+	if err != nil {
+		t.Fatal(err)
+	}
+	if gotMethod != http.MethodPost {
+		t.Fatalf("method = %q, want POST", gotMethod)
+	}
+	if gotPath != "/api/experimental/web_search" {
+		t.Fatalf("path = %q, want /api/experimental/web_search", gotPath)
+	}
+	if gotRequest.Query != "ollama" || gotRequest.MaxResults != 3 {
+		t.Fatalf("request = %#v", gotRequest)
+	}
+	if len(resp.Results) != 1 || resp.Results[0].Title != "Ollama" {
+		t.Fatalf("response = %#v", resp)
+	}
+}
+
+func TestClientWebFetchExperimentalUsesLocalRoute(t *testing.T) {
+	var gotPath string
+	var gotMethod string
+	var gotRequest WebFetchRequest
+	ts := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		gotPath = r.URL.Path
+		gotMethod = r.Method
+		if err := json.NewDecoder(r.Body).Decode(&gotRequest); err != nil {
+			t.Fatal(err)
+		}
+		if err := json.NewEncoder(w).Encode(WebFetchResponse{
+			Title:   "Ollama",
+			Content: "models",
+			Links:   []string{"https://ollama.com/library"},
+		}); err != nil {
+			t.Fatal(err)
+		}
+	}))
+	defer ts.Close()
+
+	client := NewClient(&url.URL{Scheme: "http", Host: ts.Listener.Addr().String()}, http.DefaultClient)
+	resp, err := client.WebFetchExperimental(t.Context(), &WebFetchRequest{URL: "https://ollama.com"})
+	if err != nil {
+		t.Fatal(err)
+	}
+	if gotMethod != http.MethodPost {
+		t.Fatalf("method = %q, want POST", gotMethod)
+	}
+	if gotPath != "/api/experimental/web_fetch" {
+		t.Fatalf("path = %q, want /api/experimental/web_fetch", gotPath)
+	}
+	if gotRequest.URL != "https://ollama.com" {
+		t.Fatalf("request = %#v", gotRequest)
+	}
+	if resp.Title != "Ollama" || resp.Content != "models" {
+		t.Fatalf("response = %#v", resp)
+	}
+}
+
 type roundTripFunc func(*http.Request) (*http.Response, error)
 
 func (f roundTripFunc) RoundTrip(req *http.Request) (*http.Response, error) {

api/types.go

@@ -868,6 +868,36 @@ type StatusResponse struct {
 	Cloud CloudStatus `json:"cloud"`
 }
 
+// WebSearchRequest is the request for [Client.WebSearchExperimental].
+type WebSearchRequest struct {
+	Query      string `json:"query"`
+	MaxResults int    `json:"max_results,omitempty"`
+}
+
+// WebSearchResult is a single result from [Client.WebSearchExperimental].
+type WebSearchResult struct {
+	Title   string `json:"title"`
+	URL     string `json:"url"`
+	Content string `json:"content"`
+}
+
+// WebSearchResponse is the response from [Client.WebSearchExperimental].
+type WebSearchResponse struct {
+	Results []WebSearchResult `json:"results"`
+}
+
+// WebFetchRequest is the request for [Client.WebFetchExperimental].
+type WebFetchRequest struct {
+	URL string `json:"url"`
+}
+
+// WebFetchResponse is the response from [Client.WebFetchExperimental].
+type WebFetchResponse struct {
+	Title   string   `json:"title"`
+	Content string   `json:"content"`
+	Links   []string `json:"links,omitempty"`
+}
+
 // GenerateResponse is the response passed into [GenerateResponseFunc].
 type GenerateResponse struct {
 	// Model is the model name that generated the response.

cmake/local.cmake

@@ -192,8 +192,16 @@ if(OLLAMA_MLX_BACKENDS)
     add_custom_target(ollama-mlx-sources DEPENDS ${_mlx_source_targets})
 endif()
 
+set(OLLAMA_BUILD_PARALLEL "" CACHE STRING
+    "Number of parallel jobs for nested native builds (empty = use generator default)")
+
+set(_native_parallel_args --parallel)
+if(NOT OLLAMA_BUILD_PARALLEL STREQUAL "")
+    list(APPEND _native_parallel_args ${OLLAMA_BUILD_PARALLEL})
+endif()
+
 set(OLLAMA_NATIVE_BUILD_TOOL_COMMAND
-    ${CMAKE_COMMAND} --build <BINARY_DIR>)
+    ${CMAKE_COMMAND} --build <BINARY_DIR> ${_native_parallel_args})
 set(OLLAMA_NATIVE_BUILD_TARGET_ARG --target)
 if(CMAKE_GENERATOR MATCHES "Makefiles")
     set(OLLAMA_NATIVE_BUILD_TOOL_COMMAND
@@ -236,6 +244,67 @@ function(ollama_cache_arg_is_set name output)
     endif()
 endfunction()
 
+function(ollama_backend_cuda_major backend output)
+    if("${backend}" MATCHES "^cuda_v([0-9]+)$")
+        set(${output} "${CMAKE_MATCH_1}" PARENT_SCOPE)
+    else()
+        set(${output} "" PARENT_SCOPE)
+    endif()
+endfunction()
+
+function(ollama_find_windows_cuda_root major output)
+    if(NOT WIN32 OR "${major}" STREQUAL "")
+        set(${output} "" PARENT_SCOPE)
+        return()
+    endif()
+
+    execute_process(
+        COMMAND ${CMAKE_COMMAND} -E environment
+        OUTPUT_VARIABLE _environment)
+    string(REPLACE "\r\n" "\n" _environment "${_environment}")
+    string(REPLACE "\r" "\n" _environment "${_environment}")
+    string(REGEX MATCHALL "CUDA_PATH_V${major}_[0-9]+=[^\n]*" _matches "${_environment}")
+
+    set(_best_minor -1)
+    set(_best_root "")
+    foreach(_entry IN LISTS _matches)
+        if(_entry MATCHES "^CUDA_PATH_V${major}_([0-9]+)=(.*)$")
+            set(_minor "${CMAKE_MATCH_1}")
+            set(_root "${CMAKE_MATCH_2}")
+            if(_minor GREATER _best_minor)
+                set(_best_minor ${_minor})
+                set(_best_root "${_root}")
+            endif()
+        endif()
+    endforeach()
+
+    if(_best_root STREQUAL "" AND DEFINED ENV{CUDA_PATH})
+        set(_cuda_path "$ENV{CUDA_PATH}")
+        if(EXISTS "${_cuda_path}/version.json")
+            file(READ "${_cuda_path}/version.json" _version_json)
+            if(_version_json MATCHES "\"cuda\"[ \t\r\n]*:[ \t\r\n]*\"${major}\\.")
+                set(_best_root "${_cuda_path}")
+            endif()
+        endif()
+    endif()
+
+    set(${output} "${_best_root}" PARENT_SCOPE)
+endfunction()
+
+function(ollama_append_cuda_toolkit_args output backend)
+    # If CUDAToolkit_ROOT is already explicitly set, just forward it.
+    ollama_append_cache_arg_if_set(${output} CUDAToolkit_ROOT)
+    if(NOT DEFINED CUDAToolkit_ROOT OR "${CUDAToolkit_ROOT}" STREQUAL "")
+        # Auto-discover CUDA toolkit for the requested backend version on Windows.
+        ollama_backend_cuda_major("${backend}" _cuda_major)
+        ollama_find_windows_cuda_root("${_cuda_major}" _cuda_root)
+        if(NOT "${_cuda_root}" STREQUAL "")
+            ollama_escape_cmake_list("${_cuda_root}" _value)
+            set(${output} ${${output}} "-DCUDAToolkit_ROOT=${_value}" PARENT_SCOPE)
+        endif()
+    endif()
+endfunction()
+
 function(ollama_llama_cuda_preset backend output)
     ollama_cache_arg_is_set(CMAKE_CUDA_ARCHITECTURES _has_cuda_arch)
     if(_has_cuda_arch)
@@ -327,12 +396,28 @@ function(ollama_add_llama_server_build name)
                 -DCMAKE_OSX_DEPLOYMENT_TARGET=${CMAKE_OSX_DEPLOYMENT_TARGET})
         endif()
     endif()
+    # Visual Studio requires -T toolset override to select the correct CUDA toolkit.
+    # MSBuild's CUDA integration ignores -DCUDAToolkit_ROOT for nvcc selection.
+    # Prefer user-specified CUDAToolkit_ROOT before falling back to auto-discovery.
+    set(_generator_args)
+    if(WIN32 AND CMAKE_GENERATOR MATCHES "Visual Studio")
+        set(_cuda_root "${CUDAToolkit_ROOT}")
+        if("${_cuda_root}" STREQUAL "")
+            ollama_backend_cuda_major("${name}" _cuda_major)
+            ollama_find_windows_cuda_root("${_cuda_major}" _cuda_root)
+        endif()
+        if(NOT "${_cuda_root}" STREQUAL "")
+            list(APPEND _generator_args -T cuda=${_cuda_root})
+        endif()
+    endif()
     set(_configure_command ${CMAKE_COMMAND}
+        ${_generator_args}
         -S ${CMAKE_SOURCE_DIR}/llama/server
         -B <BINARY_DIR>
         ${_cmake_args})
     if(ARG_PRESET)
         set(_configure_command ${CMAKE_COMMAND}
+            ${_generator_args}
             -S ${CMAKE_SOURCE_DIR}/llama/server
             --preset ${ARG_PRESET}
             -B <BINARY_DIR>
@@ -544,6 +629,7 @@ if(OLLAMA_HAVE_LLAMA_SERVER)
             set(_cuda_args)
             ollama_append_cache_arg_if_set(_cuda_args CMAKE_CUDA_ARCHITECTURES)
             ollama_append_cache_arg_if_set(_cuda_args CMAKE_CUDA_FLAGS)
+            ollama_append_cuda_toolkit_args(_cuda_args ${_backend})
             ollama_add_llama_server_build(${_backend}
                 PRESET ${_cuda_preset}
                 RUNNER_DIR ${_backend}
@@ -555,6 +641,7 @@ if(OLLAMA_HAVE_LLAMA_SERVER)
             set(_cuda_args)
             ollama_append_cache_arg_if_set(_cuda_args CMAKE_CUDA_ARCHITECTURES)
             ollama_append_cache_arg_if_set(_cuda_args CMAKE_CUDA_FLAGS)
+            ollama_append_cuda_toolkit_args(_cuda_args ${_backend})
             ollama_add_llama_server_build(${_backend}
                 PRESET ${_cuda_preset}
                 RUNNER_DIR ${_backend}

cmd/cmd.go

@@ -55,7 +55,6 @@ import (
 	"github.com/ollama/ollama/types/model"
 	"github.com/ollama/ollama/types/syncmap"
 	"github.com/ollama/ollama/version"
-	xcmd "github.com/ollama/ollama/x/cmd"
 	xcreate "github.com/ollama/ollama/x/create"
 	xcreateclient "github.com/ollama/ollama/x/create/client"
 	"github.com/ollama/ollama/x/imagegen"
@@ -885,11 +884,6 @@ func RunHandler(cmd *cobra.Command, args []string) error {
 		return imagegen.RunCLI(cmd, name, opts.Prompt, interactive, opts.KeepAlive)
 	}
 
-	// Check for experimental flag
-	isExperimental, _ := cmd.Flags().GetBool("experimental")
-	yoloMode, _ := cmd.Flags().GetBool("experimental-yolo")
-	enableWebsearch, _ := cmd.Flags().GetBool("experimental-websearch")
-
 	if interactive {
 		if err := loadOrUnloadModel(cmd, &opts); err != nil {
 			var sErr api.AuthorizationError
@@ -916,11 +910,6 @@ func RunHandler(cmd *cobra.Command, args []string) error {
 			}
 		}
 
-		// Use experimental agent loop with tools
-		if isExperimental {
-			return xcmd.GenerateInteractive(cmd, opts.Model, opts.WordWrap, opts.Options, opts.Think, opts.HideThinking, opts.KeepAlive, yoloMode, enableWebsearch)
-		}
-
 		return generateInteractive(cmd, opts)
 	}
 	if err := generate(cmd, opts); err != nil {
@@ -2413,9 +2402,6 @@ func NewCLI() *cobra.Command {
 	runCmd.Flags().Bool("hidethinking", false, "Hide thinking output (if provided)")
 	runCmd.Flags().Bool("truncate", false, "For embedding models: truncate inputs exceeding context length (default: true). Set --truncate=false to error instead")
 	runCmd.Flags().Int("dimensions", 0, "Truncate output embeddings to specified dimension (embedding models only)")
-	runCmd.Flags().Bool("experimental", false, "Enable experimental agent loop with tools")
-	runCmd.Flags().Bool("experimental-yolo", false, "Skip all tool approval prompts (use with caution)")
-	runCmd.Flags().Bool("experimental-websearch", false, "Enable web search tool in experimental mode")
 
 	// Image generation flags (width, height, steps, seed, etc.)
 	imagegen.RegisterFlags(runCmd)

cmd/launch/claude.go

@@ -60,17 +60,25 @@ func (c *Claude) Run(model string, _ []LaunchModel, args []string) error {
 	cmd.Stdout = os.Stdout
 	cmd.Stderr = os.Stderr
 
-	env := append(os.Environ(),
-		"ANTHROPIC_BASE_URL="+envconfig.Host().String(),
+	cmd.Env = append(os.Environ(), c.envVars(model)...)
+	return cmd.Run()
+}
+
+func (c *Claude) envVars(model string) []string {
+	env := []string{
+		"ANTHROPIC_BASE_URL=" + envconfig.Host().String(),
 		"ANTHROPIC_API_KEY=",
 		"ANTHROPIC_AUTH_TOKEN=ollama",
 		"CLAUDE_CODE_ATTRIBUTION_HEADER=0",
-	)
+		"DISABLE_TELEMETRY=1",
+		"DISABLE_ERROR_REPORTING=1",
+		"DISABLE_FEEDBACK_COMMAND=1",
+		"CLAUDE_CODE_DISABLE_FEEDBACK_SURVEY=1",
+		"CLAUDE_CODE_DISABLE_NONESSENTIAL_TRAFFIC=1",
+	}
 
 	env = append(env, c.modelEnvVars(model)...)
-
-	cmd.Env = env
-	return cmd.Run()
+	return env
 }
 
 func ensureClaudeInstalled() (string, error) {

cmd/launch/claude_test.go

@@ -8,6 +8,8 @@ import (
 	"slices"
 	"strings"
 	"testing"
+
+	"github.com/ollama/ollama/envconfig"
 )
 
 func TestClaudeIntegration(t *testing.T) {
@@ -332,6 +334,40 @@ func TestClaudeArgs(t *testing.T) {
 	}
 }
 
+func TestClaudeEnvVars(t *testing.T) {
+	c := &Claude{}
+
+	envMap := func(envs []string) map[string]string {
+		m := make(map[string]string)
+		for _, e := range envs {
+			k, v, _ := strings.Cut(e, "=")
+			m[k] = v
+		}
+		return m
+	}
+
+	got := envMap(c.envVars("llama3.2"))
+	for key, want := range map[string]string{
+		"ANTHROPIC_BASE_URL":                       envconfig.Host().String(),
+		"ANTHROPIC_API_KEY":                        "",
+		"ANTHROPIC_AUTH_TOKEN":                     "ollama",
+		"CLAUDE_CODE_ATTRIBUTION_HEADER":           "0",
+		"DISABLE_TELEMETRY":                        "1",
+		"DISABLE_ERROR_REPORTING":                  "1",
+		"DISABLE_FEEDBACK_COMMAND":                 "1",
+		"CLAUDE_CODE_DISABLE_FEEDBACK_SURVEY":      "1",
+		"CLAUDE_CODE_DISABLE_NONESSENTIAL_TRAFFIC": "1",
+		"ANTHROPIC_DEFAULT_OPUS_MODEL":             "llama3.2",
+		"ANTHROPIC_DEFAULT_SONNET_MODEL":           "llama3.2",
+		"ANTHROPIC_DEFAULT_HAIKU_MODEL":            "llama3.2",
+		"CLAUDE_CODE_SUBAGENT_MODEL":               "llama3.2",
+	} {
+		if got[key] != want {
+			t.Errorf("%s = %q, want %q", key, got[key], want)
+		}
+	}
+}
+
 func TestClaudeModelEnvVars(t *testing.T) {
 	c := &Claude{}
 

cmd/launch/hermes.go

@@ -14,6 +14,7 @@ import (
 	"strconv"
 	"strings"
 
+	"golang.org/x/mod/semver"
 	"gopkg.in/yaml.v3"
 
 	"github.com/ollama/ollama/api"
@@ -23,8 +24,10 @@ import (
 )
 
 const (
-	hermesInstallScript     = "curl -fsSL https://raw.githubusercontent.com/NousResearch/hermes-agent/main/scripts/install.sh | bash -s -- --skip-setup"
-	hermesWindowsInstallURL = "https://raw.githubusercontent.com/NousResearch/hermes-agent/main/scripts/install.ps1"
+	// https://github.com/NousResearch/hermes-agent/releases/tag/v2026.6.5
+	hermesDesktopMinVersion = "v0.16.0"
+	hermesInstallScript     = "curl -fsSL https://hermes-agent.nousresearch.com/install.sh | bash -s -- --skip-setup"
+	hermesWindowsInstallURL = "https://hermes-agent.nousresearch.com/install.ps1"
 	hermesWindowsInstallCmd = "& ([scriptblock]::Create((irm " + hermesWindowsInstallURL + "))) -SkipSetup"
 	hermesProviderName      = "Ollama"
 	hermesProviderKey       = "ollama-launch"
@@ -94,9 +97,48 @@ func (h *HermesDesktop) Run(_ string, _ []LaunchModel, args []string) error {
 	if err != nil {
 		return err
 	}
+	if err := h.ensureHermesDesktopMinVersion(bin); err != nil {
+		return err
+	}
 	return hermesAttachedCommand(bin, h.launchArgs(args)...).Run()
 }
 
+func (h *HermesDesktop) ensureHermesDesktopMinVersion(bin string) error {
+	if hermesGOOS == "windows" {
+		return nil
+	}
+	version := hermesVersionOf(bin)
+	if version == "" {
+		return nil
+	}
+	if semver.Compare(version, hermesDesktopMinVersion) >= 0 {
+		return nil
+	}
+	fmt.Fprintf(os.Stderr, "%sHermes %s is older than the minimum version (%s) for `hermes desktop`; updating...%s\n", ansiGray, version, hermesDesktopMinVersion, ansiReset)
+	if err := hermesAttachedCommand(bin, "update").Run(); err != nil {
+		return fmt.Errorf("failed to update hermes to %s or newer: %w", hermesDesktopMinVersion, err)
+	}
+	return nil
+}
+
+func hermesVersionOf(bin string) string {
+	out, err := hermesCommand(bin, "--version").Output()
+	if err != nil {
+		return ""
+	}
+	firstLine := strings.SplitN(strings.TrimSpace(string(out)), "\n", 2)[0]
+	return parseHermesVersion(firstLine)
+}
+
+func parseHermesVersion(firstLine string) string {
+	for _, field := range strings.Fields(firstLine) {
+		if semver.IsValid(field) {
+			return field
+		}
+	}
+	return ""
+}
+
 func (h *HermesDesktop) Onboard() error {
 	return config.MarkIntegrationOnboarded("hermes-desktop")
 }
@@ -128,8 +170,8 @@ func (h *HermesDesktop) packagedAppExists() bool {
 }
 
 // These roots mirror Hermes' own install layout:
-// scripts/install.sh uses ~/.hermes/hermes-agent for user installs and
-// /usr/local/lib/hermes-agent for new Linux root installs; scripts/install.ps1
+// install.sh uses ~/.hermes/hermes-agent for user installs and
+// /usr/local/lib/hermes-agent for new Linux root installs; install.ps1
 // and the bootstrap installer use %LOCALAPPDATA%\hermes\hermes-agent on
 // Windows. HERMES_HOME and HERMES_INSTALL_DIR are installer-supported
 // overrides.

cmd/launch/hermes_test.go

@@ -630,7 +630,7 @@ func hermesDesktopTestExecutableRelativePath(goos string) string {
 func writeHermesDesktopTestBinary(t *testing.T, dir string) {
 	t.Helper()
 	bin := filepath.Join(dir, "hermes")
-	if err := os.WriteFile(bin, []byte("#!/bin/sh\nprintf '[%s]\\n' \"$*\" >> \"$HOME/hermes-invocations.log\"\n"), 0o755); err != nil {
+	if err := os.WriteFile(bin, []byte("#!/bin/sh\nif [ \"$1\" = \"--version\" ]; then\n  printf 'Hermes Agent v0.16.0 (2026.6.5)\\n'\n  exit 0\nfi\nprintf '[%s]\\n' \"$*\" >> \"$HOME/hermes-invocations.log\"\n"), 0o755); err != nil {
 		t.Fatal(err)
 	}
 }
@@ -732,6 +732,147 @@ func TestHermesDesktopRun(t *testing.T) {
 		})
 	}
 }
+func writeHermesVersionedTestBinary(t *testing.T, dir, version string) {
+	t.Helper()
+	script := "#!/bin/sh\n" +
+		"case \"$1\" in\n" +
+		"  --version)\n" +
+		"    printf 'Hermes Agent " + version + " (test)\\n'\n" +
+		"    ;;\n" +
+		"  update)\n" +
+		"    printf 'update\\n' >> \"$HOME/hermes-update.log\"\n" +
+		"    ;;\n" +
+		"  *)\n" +
+		"    printf '[%s]\\n' \"$*\" >> \"$HOME/hermes-invocations.log\"\n" +
+		"    ;;\n" +
+		"esac\n"
+	bin := filepath.Join(dir, "hermes")
+	if err := os.WriteFile(bin, []byte(script), 0o755); err != nil {
+		t.Fatal(err)
+	}
+}
+
+func TestParseHermesVersion(t *testing.T) {
+	tests := []struct {
+		name  string
+		input string
+		want  string
+	}{
+		{"standard release", "Hermes Agent v0.16.0 (2026.6.5)", "v0.16.0"},
+		{"newer release", "Hermes Agent v0.17.0 (2026.6.19)", "v0.17.0"},
+		{"older release", "Hermes Agent v0.15.1 (2026.5.29)", "v0.15.1"},
+		{"prerelease", "Hermes Agent v0.16.0-rc1 (2026.6.5)", "v0.16.0-rc1"},
+		{"no version token", "Hermes Agent", ""},
+		{"empty", "", ""},
+	}
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			if got := parseHermesVersion(tt.input); got != tt.want {
+				t.Fatalf("parseHermesVersion(%q) = %q, want %q", tt.input, got, tt.want)
+			}
+		})
+	}
+}
+
+func TestHermesDesktopRun_UpdatesCliOlderThanMinVersion(t *testing.T) {
+	if runtime.GOOS == "windows" {
+		t.Skip("uses a POSIX shell test binary")
+	}
+
+	tmpDir := t.TempDir()
+	setTestHome(t, tmpDir)
+	withLauncherHooks(t)
+	withInteractiveSession(t, true)
+	withHermesPlatform(t, runtime.GOOS)
+	clearHermesMessagingEnvVars(t)
+	clearHermesDesktopPackageEnvVars(t)
+	t.Setenv("PATH", tmpDir+string(os.PathListSeparator)+os.Getenv("PATH"))
+	writeHermesVersionedTestBinary(t, tmpDir, "v0.15.1")
+
+	DefaultConfirmPrompt = func(prompt string, options ConfirmOptions) (bool, error) {
+		t.Fatalf("did not expect messaging prompt during desktop launch: %s", prompt)
+		return false, nil
+	}
+
+	if err := (&HermesDesktop{}).Run("", nil, []string{"--foreground"}); err != nil {
+		t.Fatalf("Run returned error: %v", err)
+	}
+
+	updateLog, err := os.ReadFile(filepath.Join(tmpDir, "hermes-update.log"))
+	if err != nil {
+		t.Fatalf("expected hermes update to run for an older CLI: %v", err)
+	}
+	if strings.TrimSpace(string(updateLog)) != "update" {
+		t.Fatalf("expected update log 'update', got %q", updateLog)
+	}
+	if got := readHermesDesktopInvocations(t, tmpDir); got != "[desktop --foreground]" {
+		t.Fatalf("expected desktop launch after update, got %q", got)
+	}
+}
+
+func TestHermesDesktopRun_SkipsMinVersionCheckOnWindows(t *testing.T) {
+	if runtime.GOOS == "windows" {
+		t.Skip("uses a POSIX shell test binary")
+	}
+
+	tmpDir := t.TempDir()
+	setTestHome(t, tmpDir)
+	withLauncherHooks(t)
+	withInteractiveSession(t, true)
+	withHermesPlatform(t, "windows")
+	clearHermesMessagingEnvVars(t)
+	clearHermesDesktopPackageEnvVars(t)
+	t.Setenv("PATH", tmpDir+string(os.PathListSeparator)+os.Getenv("PATH"))
+	writeHermesVersionedTestBinary(t, tmpDir, "v0.15.1")
+
+	DefaultConfirmPrompt = func(prompt string, options ConfirmOptions) (bool, error) {
+		t.Fatalf("did not expect messaging prompt during desktop launch: %s", prompt)
+		return false, nil
+	}
+
+	if err := (&HermesDesktop{}).Run("", nil, []string{"--foreground"}); err != nil {
+		t.Fatalf("Run returned error: %v", err)
+	}
+
+	if _, err := os.Stat(filepath.Join(tmpDir, "hermes-update.log")); err == nil {
+		t.Fatal("expected hermes update NOT to run on Windows, but hermes-update.log exists")
+	}
+	if got := readHermesDesktopInvocations(t, tmpDir); got != "[desktop --foreground]" {
+		t.Fatalf("expected desktop launch without update, got %q", got)
+	}
+}
+
+func TestHermesDesktopRun_DoesNotUpdateCliAtOrAboveMinVersion(t *testing.T) {
+	if runtime.GOOS == "windows" {
+		t.Skip("uses a POSIX shell test binary")
+	}
+
+	tmpDir := t.TempDir()
+	setTestHome(t, tmpDir)
+	withLauncherHooks(t)
+	withInteractiveSession(t, true)
+	withHermesPlatform(t, runtime.GOOS)
+	clearHermesMessagingEnvVars(t)
+	clearHermesDesktopPackageEnvVars(t)
+	t.Setenv("PATH", tmpDir+string(os.PathListSeparator)+os.Getenv("PATH"))
+	writeHermesVersionedTestBinary(t, tmpDir, "v0.17.0")
+
+	DefaultConfirmPrompt = func(prompt string, options ConfirmOptions) (bool, error) {
+		t.Fatalf("did not expect messaging prompt during desktop launch: %s", prompt)
+		return false, nil
+	}
+
+	if err := (&HermesDesktop{}).Run("", nil, []string{"--foreground"}); err != nil {
+		t.Fatalf("Run returned error: %v", err)
+	}
+
+	if _, err := os.Stat(filepath.Join(tmpDir, "hermes-update.log")); err == nil {
+		t.Fatal("expected hermes update NOT to run for a current CLI, but hermes-update.log exists")
+	}
+	if got := readHermesDesktopInvocations(t, tmpDir); got != "[desktop --foreground]" {
+		t.Fatalf("expected desktop launch without update, got %q", got)
+	}
+}
 
 func TestHermesDesktopRunUsesWindowsLocalAppDataPackage(t *testing.T) {
 	tmpDir := t.TempDir()

discover/runner.go

@@ -67,6 +67,15 @@ func GPUDevices(ctx context.Context, runners []ml.FilteredRunnerDiscovery) []ml.
 		requested := envconfig.LLMLibrary()
 		jetpack := cudaJetpack()
 
+		// If the detected JetPack runner isn't installed, clear the override so
+		// normal discovery can select a standard CUDA build (e.g. cuda_v13,
+		// which supports Orin on JetPack 7).
+		if jetpack != "" {
+			if _, ok := libDirs[filepath.Join(ml.LibOllamaPath, "cuda_"+jetpack)]; !ok {
+				jetpack = ""
+			}
+		}
+
 		// For our initial discovery pass, we gather all the known GPUs through
 		// all the libraries that were detected. This pass may include GPUs that
 		// are enumerated, but not actually supported.

docs/cloud.mdx

@@ -252,3 +252,20 @@ Ollama Cloud model retirement does not affect local models.
 | June 16, 2026 | `qwen3-vl:235b` | `qwen3.5` |
 | June 16, 2026 | `qwen3-vl:235b-instruct` | `qwen3.5` |
 | June 16, 2026 | `cogito-2.1:671b` | `deepseek-v4-flash` |
+| June 30, 2026 | `rnj-1:8b` | |
+| July 15, 2026 | `deepseek-v3.1:671b` | `deepseek-v4-flash` |
+| July 15, 2026 | `deepseek-v3.2` | `deepseek-v4-flash` |
+| July 15, 2026 | `devstral-2:123b` | `mistral-large-3:675b` |
+| July 15, 2026 | `devstral-small-2:24b` | |
+| July 15, 2026 | `ministral-3:14b` | |
+| July 15, 2026 | `ministral-3:3b` | |
+| July 15, 2026 | `ministral-3:8b` | |
+| July 15, 2026 | `gemini-3-flash-preview` | `minimax-m3` |
+| July 15, 2026 | `gemma3:12b` | `gemma4:31b` |
+| July 15, 2026 | `gemma3:27b` | `gemma4:31b` |
+| July 15, 2026 | `gemma3:4b` | `gemma4:31b` |
+| July 15, 2026 | `glm-4.7` | `glm-5.2` |
+| July 15, 2026 | `glm-5` | `glm-5.2` |
+| July 15, 2026 | `minimax-m2.1` | `minimax-m3` |
+| July 15, 2026 | `qwen3-coder-next` | `qwen3.5:397b` |
+| July 15, 2026 | `qwen3-coder:480b` | `qwen3.5:397b` |

docs/development.md

@@ -51,10 +51,10 @@ cmake -B build . -DOLLAMA_LLAMA_BACKENDS=cuda_v13 -DCMAKE_CUDA_ARCHITECTURES=nat
 cmake -B build . -DOLLAMA_LLAMA_BACKENDS=rocm_v7_2 -DCMAKE_HIP_ARCHITECTURES=gfx1100
 ```
 
-You can tune GGML build options by setting `GGML_*` values during configure. For example, to build CUDA v12 for Pascal without flash attention kernels:
+You can tune GGML build options by setting `GGML_*` values during configure. For example, to disable CUDA flash attention kernels for local debugging:
 
 ```shell
-cmake -B build . -DOLLAMA_LLAMA_BACKENDS=cuda_v12 -DCMAKE_CUDA_ARCHITECTURES=61 -DGGML_CUDA_FA=OFF
+cmake -B build . -DOLLAMA_LLAMA_BACKENDS=cuda_v12 -DGGML_CUDA_FA=OFF
 ```
 
 ## macOS (Apple Silicon)

docs/faq.mdx

@@ -343,7 +343,7 @@ When loading a new model, Ollama evaluates the required VRAM for the model again
 
 ## How can I enable Flash Attention?
 
-Flash Attention is a feature of most modern models that can significantly reduce memory usage as the context size grows. To enable Flash Attention, set the `OLLAMA_FLASH_ATTENTION` environment variable to `1` when starting the Ollama server.
+Flash Attention is a feature of most modern models that can significantly reduce memory usage as the context size grows. Ollama uses Flash Attention automatically when the selected backend and devices support it. To force Flash Attention on, set `OLLAMA_FLASH_ATTENTION=1` when starting the Ollama server. To disable it, set `OLLAMA_FLASH_ATTENTION=0`.
 
 ## How can I set the quantization type for the K/V cache?
 

docs/gpu.mdx

@@ -68,7 +68,7 @@ using the `amdgpu-install` utility from
 
 | Family         | Cards and accelerators                                                                                                                                         |
 | -------------- | -------------------------------------------------------------------------------------------------------------------------------------------------------------- |
-| AMD Radeon RX  | `9070 XT` `9070 GRE` `9070` `9060 XT` `9060 XT LP` `9060` `7900 XTX` `7900 XT` `7900 GRE` `7800 XT` `7700 XT` `7700` `7600 XT` `7600` `6950 XT` `6900 XTX` `6900XT` `6800 XT` `6800` `5700 XT` `5700` `5600 XT` `5500 XT` |
+| AMD Radeon RX  | `9070 XT` `9070 GRE` `9070` `9060 XT` `9060 XT LP` `9060` `7900 XTX` `7900 XT` `7900 GRE` `7800 XT` `7700 XT` `7700` `7600 XT` `7600` `6950 XT` `6900 XTX` `6900XT` `6800 XT` `6800` |
 | AMD Radeon AI PRO | `R9700` `R9600D`                                                                                                                                            |
 | AMD Radeon PRO | `W7900` `W7800` `W7700` `W7600` `W7500` `W6900X` `W6800X Duo` `W6800X` `W6800` `V620`                                                                        |
 | AMD Ryzen AI   | `Ryzen AI Max+ 395` `Ryzen AI Max 390` `Ryzen AI Max 385` `Ryzen AI 9 HX 475` `Ryzen AI 9 HX 470` `Ryzen AI 9 465` `Ryzen AI 9 HX 375` `Ryzen AI 9 HX 370` `Ryzen AI 9 365` |
@@ -80,8 +80,8 @@ Ollama requires an AMD ROCm v7 / HIP7-capable driver stack on Windows.
 
 | Family         | Cards and accelerators                                                                                               |
 | -------------- | -------------------------------------------------------------------------------------------------------------------- |
-| AMD Radeon RX  | `7900 XTX` `7900 XT` `7900 GRE` `7800 XT` `7700 XT` `7600 XT` `7600` `6950 XT` `6900 XTX` `6900XT` `6800 XT` `6800`  |
-| AMD Radeon PRO | `W7900` `W7800` `W7700` `W7600` `W7500` `W6900X` `W6800X Duo` `W6800X` `W6800` `V620`                                |
+| AMD Radeon RX  | `7900 XTX` `7900 XT` `7900 GRE` `7800 XT` `7700 XT` `7600 XT` `7600`  |
+| AMD Radeon PRO | `W7900` `W7800` `W7700` `W7600` `W7500`                                |
 
 ### Overrides on Linux
 
@@ -107,13 +107,10 @@ This table shows some example GPUs that map to these LLVM targets:
 | gfx90a | Radeon Instinct MI210/MI250 |
 | gfx942 | Radeon Instinct MI300X/MI300A |
 | gfx950 | Radeon Instinct MI350X |
-| gfx1010 | Radeon RX 5700 XT |
-| gfx1012 | Radeon RX 5500 XT |
 | gfx1030 | Radeon PRO V620 |
 | gfx1100 | Radeon PRO W7900 |
 | gfx1101 | Radeon PRO W7700 |
 | gfx1102 | Radeon RX 7600 |
-| gfx1103 | Radeon 780M |
 | gfx1150 | Ryzen AI 9 HX 375 |
 | gfx1151 | Ryzen AI Max+ 395 |
 | gfx1200 | Radeon RX 9070 |

docs/integrations/hermes.mdx

@@ -56,7 +56,7 @@ hermes setup
 If you'd rather drive Hermes's own wizard instead of `ollama launch hermes`, install it directly:
 
 ```bash
-curl -fsSL https://raw.githubusercontent.com/NousResearch/hermes-agent/main/scripts/install.sh | bash
+curl -fsSL https://hermes-agent.nousresearch.com/install.sh | bash
 ```
 
 Hermes launches the setup wizard automatically. Choose **Quick setup**:

docs/styling.css

@@ -87,6 +87,25 @@ pre, code, .font-mono {
     margin: 1.25rem 0 0.75rem;
 }
 
+.dark .capability-list {
+    border-color: #262626;
+}
+
+.dark .capability-list-title,
+.dark .capability-list-heading,
+.dark .manual-step-title {
+    color: #fafafa;
+}
+
+.dark .capability-list-icon {
+    background: #262626;
+    color: #fafafa;
+}
+
+.dark .capability-list-copy {
+    color: #a3a3a3;
+}
+
 @media (max-width: 920px) {
     .capability-list {
         grid-template-columns: 1fr;

fs/ggml/ggml.go

@@ -520,10 +520,49 @@ func (t Tensor) Elements() uint64 {
 	return count
 }
 
+func (t Tensor) elements() (uint64, bool) {
+	var count uint64 = 1
+	for _, n := range t.Shape {
+		if n != 0 && count > ^uint64(0)/n {
+			return 0, false
+		}
+		count *= n
+	}
+	return count, true
+}
+
 func (t Tensor) Size() uint64 {
 	return t.Elements() * t.typeSize() / t.blockSize()
 }
 
+func (t Tensor) size() (uint64, bool) {
+	elements, ok := t.elements()
+	if !ok {
+		return 0, false
+	}
+
+	typeSize := t.typeSize()
+	blockSize := t.blockSize()
+	if typeSize == 0 || blockSize == 0 {
+		return 0, false
+	}
+
+	rowSize := uint64(1)
+	if len(t.Shape) > 0 {
+		rowSize = t.Shape[0]
+	}
+	if rowSize%blockSize != 0 {
+		return 0, false
+	}
+
+	blocks := elements / blockSize
+	if blocks > ^uint64(0)/typeSize {
+		return 0, false
+	}
+
+	return blocks * typeSize, true
+}
+
 func (t Tensor) Type() string {
 	return TensorType(t.Kind).String()
 }

fs/ggml/gguf.go

@@ -14,9 +14,14 @@ import (
 	"strings"
 
 	"github.com/ollama/ollama/fs"
+	fsgguf "github.com/ollama/ollama/fs/gguf"
 	"golang.org/x/sync/errgroup"
 )
 
+const (
+	maxGGUFWriteConcurrency = 2
+)
+
 type containerGGUF struct {
 	ByteOrder binary.ByteOrder
 
@@ -202,6 +207,9 @@ func (llm *gguf) Decode(rs io.ReadSeeker) error {
 		if err != nil {
 			return fmt.Errorf("failed to read tensor dimensions: %w", err)
 		}
+		if dims > fsgguf.MaxTensorDims {
+			return fmt.Errorf("tensor %q dimensions %d exceeds maximum %d", name, dims, fsgguf.MaxTensorDims)
+		}
 
 		shape := make([]uint64, dims)
 		for i := 0; uint32(i) < dims; i++ {
@@ -229,13 +237,23 @@ func (llm *gguf) Decode(rs io.ReadSeeker) error {
 		}
 
 		llm.tensors = append(llm.tensors, &tensor)
-		llm.parameters += tensor.Elements()
+		elements, ok := tensor.elements()
+		if !ok {
+			return fmt.Errorf("tensor %q elements overflow", tensor.Name)
+		}
+		if llm.parameters > maxUint64()-elements {
+			return fmt.Errorf("parameter count overflow")
+		}
+		llm.parameters += elements
 	}
 
 	// patch KV with parameter count
 	llm.kv["general.parameter_count"] = llm.parameters
 
 	alignment := llm.kv.Uint("general.alignment", 32)
+	if alignment == 0 {
+		return fmt.Errorf("invalid GGUF alignment: 0")
+	}
 
 	offset, err := rs.Seek(0, io.SeekCurrent)
 	if err != nil {
@@ -243,6 +261,9 @@ func (llm *gguf) Decode(rs io.ReadSeeker) error {
 	}
 
 	padding := ggufPadding(offset, int64(alignment))
+	if padding > int64(maxInt64())-offset {
+		return fmt.Errorf("GGUF tensor offset overflow")
+	}
 	llm.tensorOffset = uint64(offset + padding)
 
 	// get file size to validate tensor bounds
@@ -256,7 +277,15 @@ func (llm *gguf) Decode(rs io.ReadSeeker) error {
 	}
 
 	for _, tensor := range llm.tensors {
-		tensorEnd := llm.tensorOffset + tensor.Offset + tensor.Size()
+		tensorSize, ok := tensor.size()
+		if !ok {
+			return fmt.Errorf("tensor %q size overflow", tensor.Name)
+		}
+
+		tensorEnd, ok := checkedAdd(llm.tensorOffset, tensor.Offset, tensorSize)
+		if !ok {
+			return fmt.Errorf("tensor %q offset+size overflows", tensor.Name)
+		}
 		if tensorEnd > uint64(fileSize) {
 			return fmt.Errorf("tensor %q offset+size (%d) exceeds file size (%d)", tensor.Name, tensorEnd, fileSize)
 		}
@@ -271,7 +300,10 @@ func (llm *gguf) Decode(rs io.ReadSeeker) error {
 			return fmt.Errorf("failed to seek to init padding: %w", err)
 		}
 
-		if _, err := rs.Seek(int64(tensor.Size()), io.SeekCurrent); err != nil {
+		if tensorSize > maxInt64() {
+			return fmt.Errorf("tensor %q size %d exceeds maximum %d", tensor.Name, tensorSize, maxInt64())
+		}
+		if _, err := rs.Seek(int64(tensorSize), io.SeekCurrent); err != nil {
 			return fmt.Errorf("failed to seek to tensor: %w", err)
 		}
 	}
@@ -298,11 +330,22 @@ func readGGUFV1String(llm *gguf, r io.Reader) (string, error) {
 	if err := binary.Read(r, llm.ByteOrder, &length); err != nil {
 		return "", err
 	}
+	if length == 0 {
+		return "", fmt.Errorf("invalid GGUF v1 string length: 0")
+	}
+
+	size, err := checkGGUFLength(length, "string")
+	if err != nil {
+		return "", err
+	}
 
 	var b bytes.Buffer
-	if _, err := io.CopyN(&b, r, int64(length)); err != nil {
+	if _, err := io.CopyN(&b, r, int64(size)); err != nil {
 		return "", err
 	}
+	if b.Bytes()[b.Len()-1] != 0 {
+		return "", fmt.Errorf("invalid GGUF v1 string terminator")
+	}
 
 	// gguf v1 strings are null-terminated
 	b.Truncate(b.Len() - 1)
@@ -320,7 +363,9 @@ func readGGUFV1StringsData(llm *gguf, r io.Reader, a *array[string]) (any, error
 
 			a.values[i] = e
 		} else {
-			_ = discardGGUFString(llm, r)
+			if err := discardGGUFString(llm, r); err != nil {
+				return nil, err
+			}
 		}
 	}
 
@@ -334,7 +379,10 @@ func discardGGUFString(llm *gguf, r io.Reader) error {
 		return err
 	}
 
-	size := int(llm.ByteOrder.Uint64(buf))
+	size, err := checkGGUFLength(llm.ByteOrder.Uint64(buf), "string")
+	if err != nil {
+		return err
+	}
 	for size > 0 {
 		n, err := r.Read(llm.scratch[:min(size, cap(llm.scratch))])
 		if err != nil {
@@ -356,7 +404,10 @@ func readGGUFString(llm *gguf, r io.Reader) (string, error) {
 		return "", err
 	}
 
-	length := int(llm.ByteOrder.Uint64(buf))
+	length, err := checkGGUFLength(llm.ByteOrder.Uint64(buf), "string")
+	if err != nil {
+		return "", err
+	}
 	if length > len(llm.scratch) {
 		buf = make([]byte, length)
 	} else {
@@ -394,7 +445,9 @@ func readGGUFStringsData(llm *gguf, r io.Reader, a *array[string]) (any, error)
 
 			a.values[i] = e
 		} else {
-			discardGGUFString(llm, r)
+			if err := discardGGUFString(llm, r); err != nil {
+				return nil, err
+			}
 		}
 	}
 
@@ -413,12 +466,20 @@ func (a *array[T]) MarshalJSON() ([]byte, error) {
 	return json.Marshal(a.values)
 }
 
-func newArray[T any](size, maxSize int) *array[T] {
-	a := array[T]{size: size}
-	if maxSize < 0 || size <= maxSize {
-		a.values = make([]T, size)
+func newArray[T any](size uint64, maxSize int) (*array[T], error) {
+	if size > uint64(maxInt()) {
+		return nil, fmt.Errorf("GGUF array size %d exceeds maximum %d", size, maxInt())
 	}
-	return &a
+	if size > fsgguf.MaxArraySize {
+		return nil, fmt.Errorf("GGUF array size %d exceeds maximum %d", size, fsgguf.MaxArraySize)
+	}
+
+	n := int(size)
+	a := array[T]{size: n}
+	if maxSize < 0 || n <= maxSize {
+		a.values = make([]T, n)
+	}
+	return &a, nil
 }
 
 func readGGUFArray(llm *gguf, r io.Reader) (any, error) {
@@ -434,40 +495,32 @@ func readGGUFArray(llm *gguf, r io.Reader) (any, error) {
 
 	switch t {
 	case ggufTypeUint8:
-		a := newArray[uint8](int(n), llm.maxArraySize)
-		return readGGUFArrayData(llm, r, a)
+		return readGGUFArrayOf[uint8](llm, r, n)
 	case ggufTypeInt8:
-		a := newArray[int8](int(n), llm.maxArraySize)
-		return readGGUFArrayData(llm, r, a)
+		return readGGUFArrayOf[int8](llm, r, n)
 	case ggufTypeUint16:
-		a := newArray[uint16](int(n), llm.maxArraySize)
-		return readGGUFArrayData(llm, r, a)
+		return readGGUFArrayOf[uint16](llm, r, n)
 	case ggufTypeInt16:
-		a := newArray[int16](int(n), llm.maxArraySize)
-		return readGGUFArrayData(llm, r, a)
+		return readGGUFArrayOf[int16](llm, r, n)
 	case ggufTypeUint32:
-		a := newArray[uint32](int(n), llm.maxArraySize)
-		return readGGUFArrayData(llm, r, a)
+		return readGGUFArrayOf[uint32](llm, r, n)
 	case ggufTypeInt32:
-		a := newArray[int32](int(n), llm.maxArraySize)
-		return readGGUFArrayData(llm, r, a)
+		return readGGUFArrayOf[int32](llm, r, n)
 	case ggufTypeUint64:
-		a := newArray[uint64](int(n), llm.maxArraySize)
-		return readGGUFArrayData(llm, r, a)
+		return readGGUFArrayOf[uint64](llm, r, n)
 	case ggufTypeInt64:
-		a := newArray[int64](int(n), llm.maxArraySize)
-		return readGGUFArrayData(llm, r, a)
+		return readGGUFArrayOf[int64](llm, r, n)
 	case ggufTypeFloat32:
-		a := newArray[float32](int(n), llm.maxArraySize)
-		return readGGUFArrayData(llm, r, a)
+		return readGGUFArrayOf[float32](llm, r, n)
 	case ggufTypeFloat64:
-		a := newArray[float64](int(n), llm.maxArraySize)
-		return readGGUFArrayData(llm, r, a)
+		return readGGUFArrayOf[float64](llm, r, n)
 	case ggufTypeBool:
-		a := newArray[bool](int(n), llm.maxArraySize)
-		return readGGUFArrayData(llm, r, a)
+		return readGGUFArrayOf[bool](llm, r, n)
 	case ggufTypeString:
-		a := newArray[string](int(n), llm.maxArraySize)
+		a, err := newArray[string](n, llm.maxArraySize)
+		if err != nil {
+			return nil, err
+		}
 		if llm.Version == 1 {
 			return readGGUFV1StringsData(llm, r, a)
 		}
@@ -478,6 +531,14 @@ func readGGUFArray(llm *gguf, r io.Reader) (any, error) {
 	}
 }
 
+func readGGUFArrayOf[T any](llm *gguf, r io.Reader, n uint64) (any, error) {
+	a, err := newArray[T](n, llm.maxArraySize)
+	if err != nil {
+		return nil, err
+	}
+	return readGGUFArrayData(llm, r, a)
+}
+
 func readGGUFArrayData[T any](llm *gguf, r io.Reader, a *array[T]) (any, error) {
 	for i := range a.size {
 		e, err := readGGUF[T](llm, r)
@@ -493,6 +554,39 @@ func readGGUFArrayData[T any](llm *gguf, r io.Reader, a *array[T]) (any, error)
 	return a, nil
 }
 
+func checkGGUFLength(n uint64, kind string) (int, error) {
+	if n > uint64(maxInt()) {
+		return 0, fmt.Errorf("GGUF %s length %d exceeds maximum %d", kind, n, maxInt())
+	}
+	if n > fsgguf.MaxStringLength {
+		return 0, fmt.Errorf("GGUF %s length %d exceeds maximum %d", kind, n, fsgguf.MaxStringLength)
+	}
+	return int(n), nil
+}
+
+func maxInt() int {
+	return int(^uint(0) >> 1)
+}
+
+func maxInt64() uint64 {
+	return 1<<63 - 1
+}
+
+func maxUint64() uint64 {
+	return ^uint64(0)
+}
+
+func checkedAdd(ns ...uint64) (uint64, bool) {
+	var sum uint64
+	for _, n := range ns {
+		if sum > maxUint64()-n {
+			return 0, false
+		}
+		sum += n
+	}
+	return sum, true
+}
+
 // writeGGUFArray writes a slice s of type E to the write with a gguf type of t
 func writeGGUFArray[S ~[]E, E any](w io.Writer, t uint32, s S) error {
 	if err := binary.Write(w, binary.LittleEndian, ggufTypeArray); err != nil {
@@ -580,8 +674,7 @@ func WriteGGUF(f *os.File, kv fs.Config, ts []*Tensor) error {
 	offset += ggufPadding(offset, int64(alignment))
 
 	var g errgroup.Group
-	g.SetLimit(runtime.GOMAXPROCS(0))
-	// TODO consider reducing if tensors size * gomaxprocs is larger than free memory
+	g.SetLimit(min(runtime.GOMAXPROCS(0), maxGGUFWriteConcurrency))
 	for _, t := range ts {
 		w := io.NewOffsetWriter(f, offset+int64(t.Offset))
 		g.Go(func() error {

fs/ggml/gguf_test.go

@@ -2,12 +2,14 @@ package ggml
 
 import (
 	"bytes"
+	"encoding/binary"
 	"math/rand/v2"
 	"os"
 	"strings"
 	"testing"
 
 	"github.com/google/go-cmp/cmp"
+	fsgguf "github.com/ollama/ollama/fs/gguf"
 )
 
 func TestWriteGGUF(t *testing.T) {
@@ -127,3 +129,141 @@ func TestWriteGGUF(t *testing.T) {
 		}
 	})
 }
+
+func TestDecodeRejectsMalformedGGUFMetadata(t *testing.T) {
+	testCases := []struct {
+		name string
+		data func() []byte
+		err  string
+	}{
+		{
+			name: "v1_zero_length_string",
+			data: func() []byte {
+				var b bytes.Buffer
+				writeRaw(t, &b, []byte("GGUF"))
+				writeRaw(t, &b, uint32(1))
+				writeRaw(t, &b, uint32(0)) // tensors
+				writeRaw(t, &b, uint32(1)) // key-values
+				writeRaw(t, &b, uint64(0))
+				return b.Bytes()
+			},
+			err: "invalid GGUF v1 string length",
+		},
+		{
+			name: "oversized_string",
+			data: func() []byte {
+				var b bytes.Buffer
+				writeRawGGUFHeader(t, &b, 0, 1)
+				writeRaw(t, &b, uint64(fsgguf.MaxStringLength+1))
+				return b.Bytes()
+			},
+			err: "string length",
+		},
+		{
+			name: "oversized_collected_array",
+			data: func() []byte {
+				var b bytes.Buffer
+				writeRawGGUFHeader(t, &b, 0, 1)
+				writeRawGGUFString(t, &b, "tokenizer.ggml.tokens")
+				writeRaw(t, &b, ggufTypeArray)
+				writeRaw(t, &b, ggufTypeString)
+				writeRaw(t, &b, uint64(fsgguf.MaxArraySize+1))
+				return b.Bytes()
+			},
+			err: "array size",
+		},
+		{
+			name: "zero_alignment",
+			data: func() []byte {
+				var b bytes.Buffer
+				writeRawGGUFHeader(t, &b, 0, 1)
+				writeRawGGUFString(t, &b, "general.alignment")
+				writeRaw(t, &b, ggufTypeUint32)
+				writeRaw(t, &b, uint32(0))
+				return b.Bytes()
+			},
+			err: "invalid GGUF alignment",
+		},
+		{
+			name: "tensor_elements_overflow",
+			data: func() []byte {
+				var b bytes.Buffer
+				writeRawGGUFHeader(t, &b, 1, 0)
+				writeRawGGUFString(t, &b, "bad.weight")
+				writeRaw(t, &b, uint32(2))
+				writeRaw(t, &b, ^uint64(0))
+				writeRaw(t, &b, uint64(2))
+				writeRaw(t, &b, uint32(TensorTypeF32))
+				writeRaw(t, &b, uint64(0))
+				return b.Bytes()
+			},
+			err: "elements overflow",
+		},
+		{
+			name: "too_many_tensor_dimensions",
+			data: func() []byte {
+				var b bytes.Buffer
+				writeRawGGUFHeader(t, &b, 1, 0)
+				writeRawGGUFString(t, &b, "bad.weight")
+				writeRaw(t, &b, uint32(fsgguf.MaxTensorDims+1))
+				return b.Bytes()
+			},
+			err: "dimensions",
+		},
+		{
+			name: "tensor_row_not_multiple_of_block_size",
+			data: func() []byte {
+				var b bytes.Buffer
+				writeRawGGUFHeader(t, &b, 1, 0)
+				writeRawGGUFString(t, &b, "bad.weight")
+				writeRaw(t, &b, uint32(1))
+				writeRaw(t, &b, uint64(31))
+				writeRaw(t, &b, uint32(TensorTypeQ4_0))
+				writeRaw(t, &b, uint64(0))
+				return b.Bytes()
+			},
+			err: "size overflow",
+		},
+	}
+
+	for _, tt := range testCases {
+		t.Run(tt.name, func(t *testing.T) {
+			t.Parallel()
+
+			defer func() {
+				if r := recover(); r != nil {
+					t.Fatalf("Decode panicked: %v", r)
+				}
+			}()
+
+			_, err := Decode(bytes.NewReader(tt.data()), -1)
+			if err == nil {
+				t.Fatal("Decode unexpectedly succeeded")
+			}
+			if !strings.Contains(err.Error(), tt.err) {
+				t.Fatalf("Decode error = %q, want containing %q", err, tt.err)
+			}
+		})
+	}
+}
+
+func writeRawGGUFHeader(t *testing.T, b *bytes.Buffer, tensors, kv uint64) {
+	t.Helper()
+	writeRaw(t, b, []byte("GGUF"))
+	writeRaw(t, b, uint32(3))
+	writeRaw(t, b, tensors)
+	writeRaw(t, b, kv)
+}
+
+func writeRawGGUFString(t *testing.T, b *bytes.Buffer, s string) {
+	t.Helper()
+	writeRaw(t, b, uint64(len(s)))
+	writeRaw(t, b, []byte(s))
+}
+
+func writeRaw(t *testing.T, b *bytes.Buffer, v any) {
+	t.Helper()
+	if err := binary.Write(b, binary.LittleEndian, v); err != nil {
+		t.Fatal(err)
+	}
+}

fs/gguf/gguf.go

@@ -13,6 +13,13 @@ import (
 	"strings"
 )
 
+const (
+	MaxStringLength = 16 << 20
+	MaxArraySize    = 64 << 20
+
+	MaxTensorDims = 4
+)
+
 const (
 	typeUint8 uint32 = iota
 	typeInt8
@@ -78,6 +85,9 @@ func Open(path string) (f *File, err error) {
 		offset := f.reader.offset
 
 		alignment := cmp.Or(f.KeyValue("general.alignment").Int(), 32)
+		if alignment <= 0 {
+			return fmt.Errorf("%w alignment %d", ErrUnsupported, alignment)
+		}
 		f.offset = offset + (alignment-offset%alignment)%alignment
 		return nil
 	}
@@ -100,6 +110,9 @@ func (f *File) readTensor() (TensorInfo, error) {
 	if err != nil {
 		return TensorInfo{}, err
 	}
+	if dims > MaxTensorDims {
+		return TensorInfo{}, fmt.Errorf("%w tensor dimensions %d exceeds maximum %d", ErrUnsupported, dims, MaxTensorDims)
+	}
 
 	shape := make([]uint64, dims)
 	for i := range dims {
@@ -119,12 +132,16 @@ func (f *File) readTensor() (TensorInfo, error) {
 		return TensorInfo{}, err
 	}
 
-	return TensorInfo{
+	ti := TensorInfo{
 		Name:   name,
 		Offset: offset,
 		Shape:  shape,
 		Type:   TensorType(type_),
-	}, nil
+	}
+	if _, ok := ti.numBytes(); !ok {
+		return TensorInfo{}, fmt.Errorf("%w tensor %q size overflows", ErrUnsupported, ti.Name)
+	}
+	return ti, nil
 }
 
 func (f *File) readKeyValue() (KeyValue, error) {
@@ -191,11 +208,16 @@ func readString(f *File) (string, error) {
 		return "", err
 	}
 
-	if int(n) > len(f.bts) {
-		f.bts = make([]byte, n)
+	length, err := checkedLength(n, "string", MaxStringLength)
+	if err != nil {
+		return "", err
 	}
 
-	bts := f.bts[:n]
+	if length > len(f.bts) {
+		f.bts = make([]byte, length)
+	}
+
+	bts := f.bts[:length]
 	if _, err := io.ReadFull(f.reader, bts); err != nil {
 		return "", err
 	}
@@ -246,8 +268,13 @@ func readArray(f *File) (any, error) {
 }
 
 func readArrayData[T any](f *File, n uint64) (s []T, err error) {
-	s = make([]T, n)
-	for i := range n {
+	size, err := checkedLength(n, "array size", MaxArraySize)
+	if err != nil {
+		return nil, err
+	}
+
+	s = make([]T, size)
+	for i := range size {
 		e, err := read[T](f)
 		if err != nil {
 			return nil, err
@@ -260,8 +287,13 @@ func readArrayData[T any](f *File, n uint64) (s []T, err error) {
 }
 
 func readArrayString(f *File, n uint64) (s []string, err error) {
-	s = make([]string, n)
-	for i := range n {
+	size, err := checkedLength(n, "array size", MaxArraySize)
+	if err != nil {
+		return nil, err
+	}
+
+	s = make([]string, size)
+	for i := range size {
 		e, err := readString(f)
 		if err != nil {
 			return nil, err
@@ -273,6 +305,24 @@ func readArrayString(f *File, n uint64) (s []string, err error) {
 	return s, nil
 }
 
+func checkedLength(n uint64, kind string, max uint64) (int, error) {
+	if n > uint64(maxInt()) {
+		return 0, fmt.Errorf("%s %d exceeds maximum %d", kind, n, maxInt())
+	}
+	if n > max {
+		return 0, fmt.Errorf("%s %d exceeds maximum %d", kind, n, max)
+	}
+	return int(n), nil
+}
+
+func maxInt() int {
+	return int(^uint(0) >> 1)
+}
+
+func maxInt64() uint64 {
+	return 1<<63 - 1
+}
+
 func (f *File) Close() error {
 	f.keyValues.stop()
 	f.tensors.stop()
@@ -337,11 +387,57 @@ func (f *File) TensorInfos() iter.Seq2[int, TensorInfo] {
 
 func (f *File) TensorReader(name string) (TensorInfo, io.Reader, error) {
 	t := f.TensorInfo(name)
-	if t.NumBytes() == 0 {
+	if err := f.Err(); err != nil {
+		return TensorInfo{}, nil, err
+	}
+	if t.Name == "" {
+		return TensorInfo{}, nil, fmt.Errorf("tensor %s not found", name)
+	}
+	numBytes, ok := t.numBytes()
+	if !ok {
+		return TensorInfo{}, nil, fmt.Errorf("%w tensor %q size overflows", ErrUnsupported, t.Name)
+	}
+	if numBytes == 0 {
 		return TensorInfo{}, nil, fmt.Errorf("tensor %s not found", name)
 	}
 
 	// fast forward through tensor info if we haven't already
-	_ = f.tensors.rest()
-	return t, io.NewSectionReader(f.file, f.offset+int64(t.Offset), t.NumBytes()), nil
+	f.tensors.rest()
+	if err := f.Err(); err != nil {
+		return TensorInfo{}, nil, err
+	}
+
+	if t.Offset > maxInt64() {
+		return TensorInfo{}, nil, fmt.Errorf("%w tensor %q offset %d exceeds maximum %d", ErrUnsupported, t.Name, t.Offset, maxInt64())
+	}
+	offset := f.offset + int64(t.Offset)
+	if offset < f.offset {
+		return TensorInfo{}, nil, fmt.Errorf("%w tensor %q offset overflows", ErrUnsupported, t.Name)
+	}
+
+	fileInfo, err := f.file.Stat()
+	if err != nil {
+		return TensorInfo{}, nil, err
+	}
+	if numBytes > fileInfo.Size()-offset {
+		return TensorInfo{}, nil, fmt.Errorf("%w tensor %q offset+size exceeds file size", ErrUnsupported, t.Name)
+	}
+
+	return t, io.NewSectionReader(f.file, offset, numBytes), nil
+}
+
+func (f *File) Err() error {
+	// Key/value and tensor metadata are read lazily, so parse errors can surface
+	// after Open succeeds.
+	if f.keyValues != nil {
+		if err := f.keyValues.Err(); err != nil {
+			return err
+		}
+	}
+	if f.tensors != nil {
+		if err := f.tensors.Err(); err != nil {
+			return err
+		}
+	}
+	return nil
 }

fs/gguf/gguf_internal_test.go

@@ -0,0 +1,158 @@
+package gguf
+
+import (
+	"bytes"
+	"encoding/binary"
+	"os"
+	"strings"
+	"testing"
+)
+
+func TestReadStringRejectsOversizedLength(t *testing.T) {
+	var b bytes.Buffer
+	writeInternalRaw(t, &b, uint64(MaxStringLength+1))
+
+	_, err := readString(testFile(b.Bytes()))
+	if err == nil {
+		t.Fatal("readString unexpectedly succeeded")
+	}
+	if !strings.Contains(err.Error(), "string") {
+		t.Fatalf("readString error = %q, want string length error", err)
+	}
+}
+
+func TestReadArrayRejectsOversizedCollectedArray(t *testing.T) {
+	var b bytes.Buffer
+	writeInternalRaw(t, &b, typeString)
+	writeInternalRaw(t, &b, uint64(MaxArraySize+1))
+
+	_, err := readArray(testFile(b.Bytes()))
+	if err == nil {
+		t.Fatal("readArray unexpectedly succeeded")
+	}
+	if !strings.Contains(err.Error(), "array size") {
+		t.Fatalf("readArray error = %q, want array size error", err)
+	}
+}
+
+func TestTensorInfoRejectsRowNotMultipleOfBlockSize(t *testing.T) {
+	ti := TensorInfo{
+		Name:  "bad.weight",
+		Shape: []uint64{31},
+		Type:  TensorTypeQ4_0,
+	}
+
+	if _, ok := ti.numBytes(); ok {
+		t.Fatal("numBytes unexpectedly succeeded")
+	}
+}
+
+func TestTensorReaderReturnsLazyParseError(t *testing.T) {
+	var b bytes.Buffer
+	writeInternalRaw(t, &b, []byte("GGUF"))
+	writeInternalRaw(t, &b, uint32(3))
+	writeInternalRaw(t, &b, uint64(1)) // tensors
+	writeInternalRaw(t, &b, uint64(0)) // key-values
+	writeInternalString(t, &b, "bad.weight")
+	writeInternalRaw(t, &b, uint32(MaxTensorDims+1))
+
+	p := writeTempFile(t, b.Bytes())
+	f, err := Open(p)
+	if err != nil {
+		t.Fatal(err)
+	}
+	defer f.Close()
+
+	_, _, err = f.TensorReader("bad.weight")
+	if err == nil {
+		t.Fatal("TensorReader unexpectedly succeeded")
+	}
+	if !strings.Contains(err.Error(), "dimensions") {
+		t.Fatalf("TensorReader error = %q, want dimensions error", err)
+	}
+}
+
+func TestTensorReaderRejectsInvalidOffset(t *testing.T) {
+	var b bytes.Buffer
+	writeInternalRaw(t, &b, []byte("GGUF"))
+	writeInternalRaw(t, &b, uint32(3))
+	writeInternalRaw(t, &b, uint64(1)) // tensors
+	writeInternalRaw(t, &b, uint64(0)) // key-values
+	writeInternalString(t, &b, "bad.weight")
+	writeInternalRaw(t, &b, uint32(1)) // dimensions
+	writeInternalRaw(t, &b, uint64(1))
+	writeInternalRaw(t, &b, uint32(TensorTypeF32))
+	writeInternalRaw(t, &b, ^uint64(0))
+
+	p := writeTempFile(t, b.Bytes())
+	f, err := Open(p)
+	if err != nil {
+		t.Fatal(err)
+	}
+	defer f.Close()
+
+	_, _, err = f.TensorReader("bad.weight")
+	if err == nil {
+		t.Fatal("TensorReader unexpectedly succeeded")
+	}
+	if !strings.Contains(err.Error(), "offset") {
+		t.Fatalf("TensorReader error = %q, want offset error", err)
+	}
+}
+
+func TestTensorReaderRejectsMissingTensor(t *testing.T) {
+	var b bytes.Buffer
+	writeInternalRaw(t, &b, []byte("GGUF"))
+	writeInternalRaw(t, &b, uint32(3))
+	writeInternalRaw(t, &b, uint64(0)) // tensors
+	writeInternalRaw(t, &b, uint64(0)) // key-values
+
+	p := writeTempFile(t, b.Bytes())
+	f, err := Open(p)
+	if err != nil {
+		t.Fatal(err)
+	}
+	defer f.Close()
+
+	_, _, err = f.TensorReader("missing.weight")
+	if err == nil {
+		t.Fatal("TensorReader unexpectedly succeeded")
+	}
+	if !strings.Contains(err.Error(), "not found") {
+		t.Fatalf("TensorReader error = %q, want not found error", err)
+	}
+}
+
+func testFile(data []byte) *File {
+	return &File{
+		reader: newBufferedReader(bytes.NewReader(data), 32<<10),
+		bts:    make([]byte, 4096),
+	}
+}
+
+func writeInternalRaw(t *testing.T, b *bytes.Buffer, v any) {
+	t.Helper()
+	if err := binary.Write(b, binary.LittleEndian, v); err != nil {
+		t.Fatal(err)
+	}
+}
+
+func writeInternalString(t *testing.T, b *bytes.Buffer, s string) {
+	t.Helper()
+	writeInternalRaw(t, b, uint64(len(s)))
+	writeInternalRaw(t, b, []byte(s))
+}
+
+func writeTempFile(t *testing.T, data []byte) string {
+	t.Helper()
+	f, err := os.CreateTemp(t.TempDir(), "")
+	if err != nil {
+		t.Fatal(err)
+	}
+	defer f.Close()
+
+	if _, err := f.Write(data); err != nil {
+		t.Fatal(err)
+	}
+	return f.Name()
+}

fs/gguf/lazy.go

@@ -2,8 +2,8 @@ package gguf
 
 import (
 	"encoding/binary"
+	"fmt"
 	"iter"
-	"log/slog"
 )
 
 type lazy[T any] struct {
@@ -11,6 +11,7 @@ type lazy[T any] struct {
 	next   func() (T, bool)
 	stop   func()
 	values []T
+	err    error
 
 	// successFunc is called when all values have been successfully read.
 	successFunc func() error
@@ -21,13 +22,16 @@ func newLazy[T any](f *File, fn func() (T, error)) (*lazy[T], error) {
 	if err := binary.Read(f.reader, binary.LittleEndian, &it.count); err != nil {
 		return nil, err
 	}
+	if it.count > uint64(maxInt()) {
+		return nil, fmt.Errorf("GGUF item count %d exceeds maximum %d", it.count, maxInt())
+	}
 
 	it.values = make([]T, 0)
 	it.next, it.stop = iter.Pull(func(yield func(T) bool) {
 		for i := range it.count {
 			t, err := fn()
 			if err != nil {
-				slog.Error("error reading tensor", "index", i, "error", err)
+				it.err = fmt.Errorf("error reading GGUF item %d: %w", i, err)
 				return
 			}
 
@@ -38,7 +42,10 @@ func newLazy[T any](f *File, fn func() (T, error)) (*lazy[T], error) {
 		}
 
 		if it.successFunc != nil {
-			it.successFunc()
+			if err := it.successFunc(); err != nil {
+				it.err = err
+				return
+			}
 		}
 	})
 
@@ -57,9 +64,10 @@ func (g *lazy[T]) Values() iter.Seq[T] {
 
 func (g *lazy[T]) All() iter.Seq2[int, T] {
 	return func(yield func(int, T) bool) {
-		for i := range int(g.count) {
-			if i < len(g.values) {
-				if !yield(i, g.values[i]) {
+		for i := range g.count {
+			n := int(i)
+			if n < len(g.values) {
+				if !yield(n, g.values[n]) {
 					break
 				}
 			} else {
@@ -68,7 +76,7 @@ func (g *lazy[T]) All() iter.Seq2[int, T] {
 					break
 				}
 
-				if !yield(i, t) {
+				if !yield(n, t) {
 					break
 				}
 			}
@@ -87,3 +95,7 @@ func (g *lazy[T]) rest() (collected bool) {
 
 	return collected
 }
+
+func (g *lazy[T]) Err() error {
+	return g.err
+}

fs/gguf/tensor.go

@@ -24,11 +24,57 @@ func (ti TensorInfo) NumValues() int64 {
 	return numItems
 }
 
+func (ti TensorInfo) numValues() (int64, bool) {
+	var numItems int64 = 1
+	for _, dim := range ti.Shape {
+		if dim > maxInt64() {
+			return 0, false
+		}
+		n := int64(dim)
+		if n != 0 && numItems > int64(maxInt64())/n {
+			return 0, false
+		}
+		numItems *= n
+	}
+	return numItems, true
+}
+
 // NumBytes returns the number of bytes in the tensor.
 func (ti TensorInfo) NumBytes() int64 {
 	return int64(float64(ti.NumValues()) * ti.Type.NumBytes())
 }
 
+func (ti TensorInfo) numBytes() (int64, bool) {
+	numValues, ok := ti.numValues()
+	if !ok {
+		return 0, false
+	}
+
+	typeSize := ti.Type.typeSize()
+	blockSize := ti.Type.blockSize()
+	if typeSize == 0 || blockSize == 0 {
+		return 0, false
+	}
+
+	rowSize := int64(1)
+	if len(ti.Shape) > 0 {
+		if ti.Shape[0] > maxInt64() {
+			return 0, false
+		}
+		rowSize = int64(ti.Shape[0])
+	}
+	if rowSize%blockSize != 0 {
+		return 0, false
+	}
+
+	blocks := numValues / blockSize
+	if blocks > int64(maxInt64())/typeSize {
+		return 0, false
+	}
+
+	return blocks * typeSize, true
+}
+
 func (ti TensorInfo) LogValue() slog.Value {
 	return slog.GroupValue(
 		slog.String("name", ti.Name),

go.mod

@@ -38,7 +38,6 @@ require (
 	github.com/wk8/go-ordered-map/v2 v2.1.8
 	golang.org/x/image v0.22.0
 	golang.org/x/mod v0.30.0
-	golang.org/x/tools v0.38.0
 	gonum.org/v1/gonum v0.15.0
 )
 

go.sum

@@ -392,8 +392,6 @@ golang.org/x/tools v0.0.0-20200130002326-2f3ba24bd6e7/go.mod h1:TB2adYChydJhpapK
 golang.org/x/tools v0.0.0-20200619180055-7c47624df98f/go.mod h1:EkVYQZoAsY45+roYkvgYkIh4xh/qjgUK9TdY2XT94GE=
 golang.org/x/tools v0.0.0-20210106214847-113979e3529a/go.mod h1:emZCQorbCU4vsT4fOWvOPXz4eW1wZW4PmDk9uLelYpA=
 golang.org/x/tools v0.1.4/go.mod h1:o0xws9oXOQQZyjljx8fwUC0k7L1pTE6eaCbjGeHmOkk=
-golang.org/x/tools v0.38.0 h1:Hx2Xv8hISq8Lm16jvBZ2VQf+RLmbd7wVUsALibYI/IQ=
-golang.org/x/tools v0.38.0/go.mod h1:yEsQ/d/YK8cjh0L6rZlY8tgtlKiBNTL14pGDJPJpYQs=
 golang.org/x/xerrors v0.0.0-20190717185122-a985d3407aa7/go.mod h1:I/5z698sn9Ka8TeJc9MKroUUfqBBauWjQqLJ2OPfmY0=
 golang.org/x/xerrors v0.0.0-20191011141410-1b5146add898/go.mod h1:I/5z698sn9Ka8TeJc9MKroUUfqBBauWjQqLJ2OPfmY0=
 golang.org/x/xerrors v0.0.0-20191204190536-9bdfabe68543/go.mod h1:I/5z698sn9Ka8TeJc9MKroUUfqBBauWjQqLJ2OPfmY0=

kvcache/cache.go

@@ -1,84 +0,0 @@
-package kvcache
-
-import (
-	"errors"
-
-	"github.com/ollama/ollama/ml"
-	"github.com/ollama/ollama/model/input"
-)
-
-var (
-	ErrKvCacheFull  = errors.New("could not find a kv cache slot")
-	ErrNotSupported = errors.New("model does not support operation")
-)
-
-type Cache interface {
-	// ** used by model implementations **
-
-	// SetLayer sets the active layer of the cache
-	SetLayer(layer int)
-
-	// Get returns the history of key and value tensors plus a mask
-	//
-	// The shape of the tensors is documented in the specific
-	// cache implementation used.
-	Get(ctx ml.Context) (ml.Tensor, ml.Tensor, ml.Tensor)
-
-	// Put stores a batch of key and value in the cache
-	//
-	// The shape of the tensors is documented in the specific
-	// cache implementation used.
-	Put(ctx ml.Context, key, value ml.Tensor)
-
-	// SetConfig controls optimizations (mostly backend-specific) that may transform
-	// the output of the cache to work better with specific kernels. If not called,
-	// the backend settings will be used. This works well when calling Attention.
-	//
-	// The config can be overridden by models, especially if they require vanilla
-	// output when implementing their own version of attention. To do this, pass
-	// an empty ml.CacheConfig.
-	//
-	// Most models will not need to use this.
-	SetConfig(ml.CacheConfig)
-
-	// ** cache management **
-
-	// Init sets up runtime parameters.
-	// backend: Used to allocate cache data storage and execute management operations (such as defrag)
-	// dtype: The data type for storing cache entries
-	// maxSequences: The maximum number of sequences stored in the cache - across all batches
-	// capacity: The number of cache entries to store, per sequence
-	// maxBatch: The maximum number of tokens that can occur in a single batch
-	Init(backend ml.Backend, dtype ml.DType, maxSequences, capacity, maxBatch int)
-
-	// Close closes the cache and frees resources associated with it
-	Close()
-
-	// StartForward is called before the start of the model's forward pass.
-	// For each token in the coming batch, there must be a corresponding
-	// entry in positions and seqs. reserve is to preallocate memory
-	// without actually storing data in the cache.
-	StartForward(ctx ml.Context, batch input.Batch, reserve bool) error
-
-	// CopyPrefix copies tokens in the range [0, len) from srcSeq to dstSeq
-	CopyPrefix(srcSeq, dstSeq int, len int32)
-
-	// CanResume returns true if the cache can continue with the next token at
-	// the given position and sequence. Assumes that the caller has already
-	// verified the contents of the cache.
-	CanResume(seq int, pos int32) bool
-
-	// Remove deletes tokens in the range [beginIndex, endIndex) from seq. Set
-	// endIndex to math.MaxInt32 to remove everything starting at beginIndex.
-	//
-	// If an error occurs, the entire context for the sequence should be
-	// removed by calling Remove(seq, 0, math.MaxInt32)
-	Remove(seq int, beginIndex, endIndex int32) error
-}
-
-// CheckpointCache optionally supports restoring recurrent state to a prior
-// position to avoid full prompt reprocessing when a prefix mismatch occurs.
-// The returned position is the number of tokens that can be kept (prefix length).
-type CheckpointCache interface {
-	PrepareRestore(seq int, targetPos int32) (int32, bool)
-}

kvcache/causal.go

@@ -1,666 +0,0 @@
-package kvcache
-
-import (
-	"errors"
-	"fmt"
-	"math"
-	"slices"
-
-	"github.com/ollama/ollama/ml"
-	"github.com/ollama/ollama/model/input"
-)
-
-type shiftFn func(ctx ml.Context, layer int, key, shift ml.Tensor) (ml.Tensor, error)
-
-// Causal cache stores K and V tensors according to their position in the
-// sequence. Returns the history and a mask for attending to past tokens
-//
-// The tensors are of shape embed dim, kv heads, batch size
-// The mask is of shape history size, batch size
-type Causal struct {
-	DType ml.DType
-
-	// swaWindowSize is the number of tokens that will be included in the mask
-	// during attention operations. swaMemorySize is the number of tokens that
-	// will be retained in memory for partial prefix caching. Set to math.MaxInt32
-	// for unlimited or if sliding window attention is not being used.
-	swaWindowSize int32
-	swaMemorySize int32
-
-	chunkSize int32
-
-	opts CausalOptions
-
-	// maxBatch is the largest batch that we might receive
-	maxBatch int
-
-	// config controls mostly backend-specific optimizations
-	config *ml.CacheConfig
-
-	// ** current forward pass **
-
-	// size of the current batch
-	curBatchSize int
-
-	// locations for data storage for this batch
-	curLoc ml.Tensor
-
-	// mask of the cache as used by this batch
-	curMask ml.Tensor
-
-	// the active layer for Get and Put
-	curLayer int
-
-	// locations in the cache that are needed for this batch
-	curCellRange cellRange
-
-	// curSequences is the sequences corresponding to this pass's entries in the cache
-	curSequences []int
-
-	// curPositions is the positions corresponding to this pass's entries in the cache
-	curPositions []int32
-
-	// ** cache metadata **
-
-	// for each possible location in the cache, stores the position and set of sequences
-	// that reference the data there
-	cells []cacheCell
-
-	// maps from sequence to the range of locations where it is stored in the cache
-	cellRanges map[int]cellRange
-
-	// ** cache data storage **
-
-	shiftFn      shiftFn
-	backend      ml.Backend
-	ctxs         map[int]ml.Context
-	keys, values map[int]ml.Tensor
-}
-
-type cacheCell struct {
-	pos       int32
-	sequences []int
-}
-
-type cellRange struct {
-	min int
-	max int
-}
-
-func NewCausalCache(shift shiftFn) *Causal {
-	return &Causal{
-		shiftFn: shift,
-		ctxs:    make(map[int]ml.Context),
-		keys:    make(map[int]ml.Tensor),
-		values:  make(map[int]ml.Tensor),
-	}
-}
-
-func NewSWACache(windowSize int32, shift shiftFn) *Causal {
-	return &Causal{
-		swaWindowSize: windowSize,
-		shiftFn:       shift,
-		ctxs:          make(map[int]ml.Context),
-		keys:          make(map[int]ml.Tensor),
-		values:        make(map[int]ml.Tensor),
-	}
-}
-
-func NewSWAMemCache(windowSize int32, memorySize int32, shift shiftFn) *Causal {
-	return &Causal{
-		swaWindowSize: windowSize,
-		swaMemorySize: memorySize,
-		shiftFn:       shift,
-		ctxs:          make(map[int]ml.Context),
-		keys:          make(map[int]ml.Tensor),
-		values:        make(map[int]ml.Tensor),
-	}
-}
-
-func NewChunkedAttentionCache(chunkSize int32, shift shiftFn) *Causal {
-	return &Causal{
-		chunkSize: chunkSize,
-		shiftFn:   shift,
-		ctxs:      make(map[int]ml.Context),
-		keys:      make(map[int]ml.Tensor),
-		values:    make(map[int]ml.Tensor),
-	}
-}
-
-func (c *Causal) Init(backend ml.Backend, dtype ml.DType, maxSequences, capacity, maxBatch int) {
-	if c.config == nil {
-		var config ml.CacheConfig
-		if cc, ok := backend.(ml.BackendCacheConfig); ok {
-			config = cc.CacheConfig()
-		}
-		c.config = &config
-	}
-
-	if c.config.CachePadding == 0 {
-		c.config.CachePadding = 1
-	}
-
-	if c.config.MaskDType == ml.DTypeOther {
-		c.config.MaskDType = ml.DTypeF32
-	}
-
-	if c.swaWindowSize == 0 {
-		c.swaWindowSize = math.MaxInt32
-	}
-	if c.swaMemorySize == 0 {
-		c.swaMemorySize = c.swaWindowSize
-	}
-	// We will allocate space in the cache for the stop token, which won't be part of a follow on
-	// sequence, so allocate an extra token of storage to ensure that we can jump back without
-	// causing a cache break. As an optimization, only do this when we have parallel sequences
-	// because the extra token will live in the batch buffer and won't get overwritten if we
-	// only have a single sequence.
-	if c.swaMemorySize != math.MaxInt32 && maxSequences > 1 {
-		c.swaMemorySize = max(c.swaMemorySize, c.swaWindowSize+1)
-	}
-	if int(c.swaMemorySize) >= capacity {
-		c.swaMemorySize = math.MaxInt32
-	}
-
-	if c.swaMemorySize < c.swaWindowSize {
-		panic(fmt.Errorf("sliding window memory (%v) must be at least as large as the window (%v)", c.swaMemorySize, c.swaWindowSize))
-	}
-
-	var cacheSize int
-	if c.swaMemorySize == math.MaxInt32 {
-		cacheSize = maxSequences * capacity
-	} else {
-		cacheSize = (maxSequences * int(c.swaMemorySize)) + maxBatch
-	}
-	cacheSize = roundUp(cacheSize, c.config.CachePadding)
-	c.cells = make([]cacheCell, cacheSize)
-
-	c.DType = dtype
-	c.cellRanges = make(map[int]cellRange)
-	c.backend = backend
-	c.maxBatch = maxBatch
-}
-
-func (c *Causal) SetConfig(config ml.CacheConfig) {
-	if c.config != nil {
-		panic("config cannot be changed after being previously set, either by the model or backend")
-	}
-
-	c.config = &config
-}
-
-func (c *Causal) Close() {
-	for _, ctx := range c.ctxs {
-		ctx.Close()
-	}
-}
-
-func (c *Causal) StartForward(ctx ml.Context, batch input.Batch, reserve bool) error {
-	c.curBatchSize = len(batch.Positions)
-	c.curSequences = batch.Sequences
-	c.curPositions = batch.Positions
-	c.opts.Except = nil
-
-	var locs []int32
-	if !reserve {
-		c.updateSlidingWindow()
-
-		var err error
-		locs, err = c.findLocs()
-		if err != nil {
-			return err
-		}
-
-		for i, pos := range batch.Positions {
-			seq := batch.Sequences[i]
-			loc := int(locs[i])
-
-			c.cells[loc] = cacheCell{pos: pos, sequences: []int{seq}}
-
-			seqRange, ok := c.cellRanges[seq]
-			if !ok {
-				seqRange = newRange()
-			}
-
-			seqRange.min = min(seqRange.min, loc)
-			c.curCellRange.min = min(c.curCellRange.min, loc)
-
-			seqRange.max = max(seqRange.max, loc)
-			c.curCellRange.max = max(c.curCellRange.max, loc)
-
-			c.cellRanges[seq] = seqRange
-		}
-	} else {
-		// If we are reserving memory, don't update any of the cache metadata but set the size
-		// to the worst case.
-		locs = make([]int32, c.curBatchSize)
-		for i := range locs {
-			locs[i] = int32(i)
-		}
-		c.curCellRange.min = 0
-		c.curCellRange.max = len(c.cells) - 1
-	}
-
-	c.curLoc = ctx.Input().FromInts(locs, len(locs))
-	c.curMask = c.buildMask(ctx)
-
-	return nil
-}
-
-func newRange() cellRange {
-	return cellRange{
-		min: math.MaxInt,
-		max: 0,
-	}
-}
-
-// Returns a slice of locations where each token in the batch should be stored
-func (c *Causal) findLocs() ([]int32, error) {
-	loc := make([]int32, 0, c.curBatchSize)
-
-	for i := range c.cells {
-		if len(c.cells[i].sequences) == 0 {
-			loc = append(loc, int32(i))
-			if len(loc) >= c.curBatchSize {
-				return loc, nil
-			}
-		}
-	}
-
-	return nil, fmt.Errorf("%w (cache: %v batch: %v)", ErrKvCacheFull, len(c.cells), c.curBatchSize)
-}
-
-func (c *Causal) updateSlidingWindow() {
-	c.curCellRange = newRange()
-
-	if c.swaMemorySize == math.MaxInt32 {
-		for _, seq := range c.curSequences {
-			if seqRange, ok := c.cellRanges[seq]; ok {
-				c.curCellRange.min = min(c.curCellRange.min, seqRange.min)
-				c.curCellRange.max = max(c.curCellRange.max, seqRange.max)
-			}
-		}
-
-		return
-	}
-
-	type lowestPosition struct {
-		pos      int32
-		curBatch bool
-	}
-
-	// create a map of unique sequences to the lowest position in that sequence
-	lowestPos := make(map[int]lowestPosition)
-	for i := range c.curPositions {
-		seq := c.curSequences[i]
-
-		lowest, ok := lowestPos[seq]
-		if !ok {
-			lowest = lowestPosition{pos: c.curPositions[i], curBatch: true}
-		} else if c.curPositions[i] < lowest.pos {
-			lowest.pos = c.curPositions[i]
-		}
-
-		lowestPos[seq] = lowest
-	}
-
-	// for any sequences are not part of this batch, clean up any tokens
-	// that are no longer needed after the processing of the previous
-	// batch
-	for seq, seqRange := range c.cellRanges {
-		if _, ok := lowestPos[seq]; !ok {
-			var last int32
-			for i := seqRange.min; i <= seqRange.max; i++ {
-				if slices.Contains(c.cells[i].sequences, seq) {
-					last = max(last, c.cells[i].pos)
-				}
-			}
-
-			lowestPos[seq] = lowestPosition{pos: last + 1, curBatch: false}
-		}
-	}
-
-	// delete any entries that are beyond the window of the oldest position in the sequence
-	for seq, lowest := range lowestPos {
-		oldRange, ok := c.cellRanges[seq]
-		if !ok {
-			continue
-		}
-
-		newRange := newRange()
-
-		for i := oldRange.min; i <= oldRange.max; i++ {
-			if slices.Contains(c.cells[i].sequences, seq) {
-				if c.cells[i].pos < lowest.pos-c.swaMemorySize {
-					c.cells[i].sequences = slices.DeleteFunc(c.cells[i].sequences, func(s int) bool { return s == seq })
-				} else {
-					newRange.min = min(newRange.min, i)
-					newRange.max = max(newRange.max, i)
-				}
-				if lowest.curBatch && c.cells[i].pos >= lowest.pos-c.swaWindowSize {
-					c.curCellRange.min = min(c.curCellRange.min, i)
-					c.curCellRange.max = max(c.curCellRange.max, i)
-				}
-			}
-		}
-
-		c.cellRanges[seq] = newRange
-	}
-}
-
-func roundDown(length, pad int) int {
-	return (length / pad) * pad
-}
-
-func roundUp(length, pad int) int {
-	return ((length + pad - 1) / pad) * pad
-}
-
-// Builds a mask of history x batch indicating whether for each token in the batch the
-// token in the history should apply. This is based on both the sequence and causality (the
-// position of the history is not ahead of the token in the batch).
-func (c *Causal) buildMask(ctx ml.Context) ml.Tensor {
-	c.curCellRange.min = roundDown(c.curCellRange.min, c.config.CachePadding)
-	c.curCellRange.max = roundUp(c.curCellRange.max+1, c.config.CachePadding) - 1
-
-	length := c.curCellRange.max - c.curCellRange.min + 1
-
-	mask := make([]float32, c.curBatchSize*length)
-
-	for i := range c.curBatchSize {
-		enabled := !slices.Contains(c.opts.Except, i)
-		for j := c.curCellRange.min; j <= c.curCellRange.max; j++ {
-			if !slices.Contains(c.cells[j].sequences, c.curSequences[i]) ||
-				(enabled && c.cells[j].pos > c.curPositions[i]) ||
-				c.chunkSize > 0 && c.cells[j].pos < c.curPositions[i]-c.curPositions[i]%c.chunkSize ||
-				c.cells[j].pos < c.curPositions[i]-c.swaWindowSize {
-				mask[i*length+(j-c.curCellRange.min)] = float32(math.Inf(-1))
-			}
-		}
-	}
-
-	maskTensor := ctx.Input().FromFloats(mask, length, c.curBatchSize)
-
-	if c.config.MaskDType != ml.DTypeF32 {
-		maskTensor = maskTensor.Cast(ctx, c.config.MaskDType)
-	}
-
-	return maskTensor
-}
-
-func (c *Causal) SetLayer(layer int) {
-	c.curLayer = layer
-}
-
-type CausalOptions struct {
-	// Enabled controls whether the causal mask is generated for a particular index in a batch
-	Except []int
-}
-
-// SetCausal disables causal mask generation for a particular range of indicies in
-// the current batch for subsequent calls to Get. The state resets for the next forward pass.
-func (c *Causal) SetCausal(ctx ml.Context, opts CausalOptions) {
-	if !slices.Equal(c.opts.Except, opts.Except) {
-		c.opts = opts
-		if ctx != nil {
-			c.curMask = c.buildMask(ctx)
-		}
-	}
-}
-
-func (c *Causal) Get(ctx ml.Context) (ml.Tensor, ml.Tensor, ml.Tensor) {
-	key := c.keys[c.curLayer]
-	value := c.values[c.curLayer]
-
-	kHeadDim := key.Dim(0)
-	numKVHeads := key.Dim(1)
-	rowSize := key.Stride(2)
-	cachedSize := c.curMask.Dim(0)
-
-	key = key.View(ctx, rowSize*c.curCellRange.min,
-		kHeadDim, key.Stride(1),
-		numKVHeads, key.Stride(2),
-		cachedSize,
-	)
-
-	if c.config.PermutedV {
-		vHeadDim := value.Dim(1)
-		elemSize := value.Stride(0)
-
-		value = value.View(ctx, elemSize*c.curCellRange.min,
-			cachedSize, value.Stride(1),
-			vHeadDim, value.Stride(2),
-			numKVHeads,
-		)
-	} else {
-		vHeadDim := value.Dim(0)
-		rowSize := value.Stride(2)
-
-		value = value.View(ctx, rowSize*c.curCellRange.min,
-			vHeadDim, value.Stride(1),
-			numKVHeads, value.Stride(2),
-			cachedSize,
-		)
-	}
-
-	return key, value, c.curMask
-}
-
-func (c *Causal) Put(ctx ml.Context, key, value ml.Tensor) {
-	kHeadDim := key.Dim(0)
-	vHeadDim := value.Dim(0)
-	numKVHeads := key.Dim(1)
-	batchSize := key.Dim(2)
-
-	if c.curBatchSize != batchSize {
-		panic(fmt.Errorf("inconsistent batch sizes (layer: %v, batch size: %v layer batch size: %v)", c.curLayer, c.curBatchSize, batchSize))
-	}
-
-	if _, ok := c.ctxs[c.curLayer]; !ok {
-		c.ctxs[c.curLayer] = c.backend.NewContextSize(2).Layer(c.curLayer)
-	}
-
-	if _, ok := c.keys[c.curLayer]; !ok {
-		c.keys[c.curLayer] = c.ctxs[c.curLayer].Zeros(c.DType, kHeadDim, numKVHeads, len(c.cells))
-	}
-
-	if _, ok := c.values[c.curLayer]; !ok {
-		if c.config.PermutedV {
-			c.values[c.curLayer] = c.ctxs[c.curLayer].Zeros(c.DType, len(c.cells), vHeadDim, numKVHeads)
-		} else {
-			c.values[c.curLayer] = c.ctxs[c.curLayer].Zeros(c.DType, vHeadDim, numKVHeads, len(c.cells))
-		}
-	}
-
-	key = key.Reshape(ctx, kHeadDim*numKVHeads, batchSize)
-	keyCache := c.keys[c.curLayer]
-	keyCache = keyCache.Reshape(ctx, kHeadDim*numKVHeads, len(c.cells))
-	ctx.Forward(keyCache.SetRows(ctx, key, c.curLoc))
-
-	if c.config.PermutedV {
-		value = value.Reshape(ctx, vHeadDim*numKVHeads, 1, batchSize)
-		value = value.Permute(ctx, 2, 0, 1, 3)
-
-		valueCache := c.values[c.curLayer]
-		valueCache = valueCache.Reshape(ctx, 1, len(c.cells), vHeadDim*numKVHeads)
-
-		ctx.Forward(valueCache.SetRows(ctx, value, c.curLoc))
-	} else {
-		value = value.Reshape(ctx, vHeadDim*numKVHeads, batchSize)
-		valueCache := c.values[c.curLayer]
-		valueCache = valueCache.Reshape(ctx, vHeadDim*numKVHeads, len(c.cells))
-
-		ctx.Forward(valueCache.SetRows(ctx, value, c.curLoc))
-	}
-}
-
-func (c *Causal) CopyPrefix(srcSeq, dstSeq int, len int32) {
-	seqRange := newRange()
-
-	for i := range c.cells {
-		// Remove the contents of dstSeq so that we only have the copied prefix, metadata will be reset at the end
-		if slices.Contains(c.cells[i].sequences, dstSeq) {
-			c.cells[i].sequences = slices.DeleteFunc(c.cells[i].sequences, func(s int) bool { return s == dstSeq })
-		}
-
-		if slices.Contains(c.cells[i].sequences, srcSeq) && c.cells[i].pos < len {
-			c.cells[i].sequences = append(c.cells[i].sequences, dstSeq)
-			if i < seqRange.min {
-				seqRange.min = i
-			}
-			if i > seqRange.max {
-				seqRange.max = i
-			}
-		}
-	}
-
-	c.cellRanges[dstSeq] = seqRange
-}
-
-func (c *Causal) CanResume(seq int, pos int32) bool {
-	if c.swaMemorySize == math.MaxInt32 {
-		return true
-	}
-
-	seqRange, ok := c.cellRanges[seq]
-	if !ok {
-		return false
-	}
-
-	// for sliding window, check that the window of the new sequence is contained in
-	// the window of what we are storing
-	var first int32 = math.MaxInt32
-	var last int32 = -1
-	for i := seqRange.min; i <= seqRange.max; i++ {
-		if slices.Contains(c.cells[i].sequences, seq) {
-			first = min(first, c.cells[i].pos)
-			last = max(last, c.cells[i].pos)
-		}
-	}
-
-	if last == -1 {
-		return false
-	}
-
-	posWindowStart := max(0, pos-c.swaWindowSize)
-	return posWindowStart >= first && pos <= last+1
-}
-
-func (c *Causal) shift(seq int, beginIndex, offset int32) error {
-	if c.shiftFn == nil {
-		return ErrNotSupported
-	}
-
-	seqRange := c.cellRanges[seq]
-
-	for start := seqRange.min; start <= seqRange.max; start += c.maxBatch {
-		size := min(seqRange.max-start+1, c.maxBatch)
-		offsets := make([]int32, size)
-
-		var batchFirst, batchLast int
-
-		batchFirst = -1
-		for i := range offsets {
-			cell := c.cells[start+i]
-
-			if slices.Contains(cell.sequences, seq) && cell.pos >= beginIndex {
-				offsets[i] = offset
-				if batchFirst < 0 {
-					batchFirst = i
-				}
-				batchLast = i
-			}
-		}
-
-		if batchFirst < 0 {
-			continue
-		}
-
-		offsets = offsets[batchFirst : batchLast+1]
-
-		ctx := c.backend.NewContext()
-		kShift := ctx.Input().FromInts(offsets, len(offsets))
-
-		for i, key := range c.keys {
-			if key == nil {
-				continue
-			}
-
-			kHeadDim := key.Dim(0)
-			numKVHeads := key.Dim(1)
-			rowSize := key.Stride(2)
-
-			key = key.View(ctx, rowSize*(start+batchFirst),
-				kHeadDim, key.Stride(1),
-				numKVHeads, key.Stride(2),
-				len(offsets),
-			)
-
-			roped, err := c.shiftFn(ctx, i, key, kShift)
-			if err != nil {
-				ctx.Close()
-				return err
-			}
-
-			ctx.Forward(roped.Copy(ctx, key))
-		}
-
-		ctx.Compute()
-		ctx.Close()
-	}
-
-	return nil
-}
-
-func (c *Causal) Remove(seq int, beginIndex, endIndex int32) error {
-	// TODO(jessegross): We should check to see if removing the middle of the sequence will
-	// cause the sliding window to encompass tokens that we no longer have. If so, then we
-	// should return an error, which will trigger the runner to evaluate the full history and
-	// rebuild the window. However, if we have multimodal inputs in our history, this reuse
-	// results in use after free, so we don't do it for now.
-
-	var offset int32
-	if endIndex != math.MaxInt32 {
-		offset = beginIndex - endIndex
-	}
-
-	seqRange := newRange()
-
-	for i := range c.cells {
-		if slices.Contains(c.cells[i].sequences, seq) {
-			if c.cells[i].pos >= beginIndex && c.cells[i].pos < endIndex {
-				c.cells[i].sequences = slices.DeleteFunc(c.cells[i].sequences, func(s int) bool { return s == seq })
-			} else {
-				if c.cells[i].pos >= endIndex {
-					if slices.ContainsFunc(c.cells[i].sequences, func(s int) bool { return s != seq }) {
-						return errors.New("shifting cells shared by multiple sequences not supported")
-					}
-
-					c.cells[i].pos += offset
-				}
-				if i < seqRange.min {
-					seqRange.min = i
-				}
-				if i > seqRange.max {
-					seqRange.max = i
-				}
-			}
-		}
-	}
-
-	if seqRange == newRange() {
-		delete(c.cellRanges, seq)
-		return nil
-	}
-
-	c.cellRanges[seq] = seqRange
-
-	if endIndex != math.MaxInt32 {
-		err := c.shift(seq, endIndex+offset, offset)
-		if err != nil {
-			return err
-		}
-	}
-
-	return nil
-}

kvcache/causal_test.go

@@ -1,973 +0,0 @@
-package kvcache
-
-import (
-	"fmt"
-	"math"
-	"slices"
-	"testing"
-
-	"github.com/ollama/ollama/ml"
-	"github.com/ollama/ollama/model/input"
-)
-
-type testCase struct {
-	name          string
-	in            []float32
-	inShape       []int
-	seqs          []int
-	pos           []int32
-	expected      []float32
-	expectedShape []int
-	expectedMask  []float32
-}
-
-func runPermutedVariants(t *testing.T, fn func(t *testing.T, backend *testBackend)) {
-	t.Helper()
-	for _, permuted := range []bool{false, true} {
-		t.Run(fmt.Sprintf("PermutedV=%t", permuted), func(t *testing.T) {
-			fn(t, &testBackend{permutedV: permuted})
-		})
-	}
-}
-
-func TestStore(t *testing.T) {
-	runPermutedVariants(t, func(t *testing.T, backend *testBackend) {
-		cache := NewCausalCache(nil)
-		defer cache.Close()
-
-		cache.Init(backend, ml.DTypeF16, 1, 16, 16)
-
-		tests := []testCase{
-			{
-				name:          "FirstBatch",
-				in:            []float32{111, 211, 121, 221, 131, 231, 112, 212, 122, 222, 132, 232, 113, 213, 123, 223, 133, 233, 114, 214, 124, 224, 134, 234},
-				inShape:       []int{2, 3, 4},
-				seqs:          []int{0, 0, 0, 0},
-				pos:           []int32{0, 1, 2, 3},
-				expected:      []float32{111, 211, 121, 221, 131, 231, 112, 212, 122, 222, 132, 232, 113, 213, 123, 223, 133, 233, 114, 214, 124, 224, 134, 234},
-				expectedShape: []int{2, 3, 4},
-				expectedMask:  []float32{0, float32(math.Inf(-1)), float32(math.Inf(-1)), float32(math.Inf(-1)), 0, 0, float32(math.Inf(-1)), float32(math.Inf(-1)), 0, 0, 0, float32(math.Inf(-1)), 0, 0, 0, 0},
-			},
-			{
-				name:          "SecondBatch",
-				in:            []float32{115, 215, 125, 225, 135, 235},
-				inShape:       []int{2, 3, 1},
-				seqs:          []int{0},
-				pos:           []int32{4},
-				expected:      []float32{111, 211, 121, 221, 131, 231, 112, 212, 122, 222, 132, 232, 113, 213, 123, 223, 133, 233, 114, 214, 124, 224, 134, 234, 115, 215, 125, 225, 135, 235},
-				expectedShape: []int{2, 3, 5},
-				expectedMask:  []float32{0, 0, 0, 0, 0},
-			},
-		}
-
-		testCache(t, backend, cache, tests)
-	})
-}
-
-func TestSWA(t *testing.T) {
-	runPermutedVariants(t, func(t *testing.T, backend *testBackend) {
-		cache := NewSWACache(1, nil)
-		defer cache.Close()
-
-		cache.Init(backend, ml.DTypeF16, 1, 16, 16)
-
-		x := float32(math.Inf(-1))
-
-		tests := []testCase{
-			{
-				name:          "FirstBatch",
-				in:            []float32{1, 2, 3, 4},
-				inShape:       []int{1, 1, 4},
-				seqs:          []int{0, 0, 0, 0},
-				pos:           []int32{0, 1, 2, 3},
-				expected:      []float32{1, 2, 3, 4},
-				expectedShape: []int{1, 1, 4},
-				expectedMask: []float32{
-					0, x, x, x,
-					0, 0, x, x,
-					x, 0, 0, x,
-					x, x, 0, 0,
-				},
-			},
-			{
-				name:          "SecondBatch",
-				in:            []float32{5, 6},
-				inShape:       []int{1, 1, 2},
-				seqs:          []int{0, 0},
-				pos:           []int32{4, 5},
-				expected:      []float32{5, 6, 3, 4},
-				expectedShape: []int{1, 1, 4},
-				expectedMask: []float32{
-					0, x, x, 0,
-					0, 0, x, x,
-				},
-			},
-		}
-
-		testCache(t, backend, cache, tests)
-	})
-}
-
-func TestSWASeparateBatches(t *testing.T) {
-	runPermutedVariants(t, func(t *testing.T, backend *testBackend) {
-		cache := NewSWACache(1, nil)
-		defer cache.Close()
-
-		cache.Init(backend, ml.DTypeF16, 2, 16, 2)
-
-		x := float32(math.Inf(-1))
-
-		tests := []testCase{
-			{
-				name:          "First seq 0",
-				in:            []float32{1, 2},
-				inShape:       []int{1, 1, 2},
-				seqs:          []int{0, 0},
-				pos:           []int32{0, 1},
-				expected:      []float32{1, 2},
-				expectedShape: []int{1, 1, 2},
-				expectedMask: []float32{
-					0, x,
-					0, 0,
-				},
-			},
-			{
-				name:          "Second seq 0",
-				in:            []float32{3, 4},
-				inShape:       []int{1, 1, 2},
-				seqs:          []int{0, 0},
-				pos:           []int32{2, 3},
-				expected:      []float32{2, 3, 4},
-				expectedShape: []int{1, 1, 3},
-				expectedMask: []float32{
-					0, 0, x,
-					x, 0, 0,
-				},
-			},
-			{
-				name:          "First seq 1",
-				in:            []float32{5, 6},
-				inShape:       []int{1, 1, 2},
-				seqs:          []int{1, 1},
-				pos:           []int32{0, 1},
-				expected:      []float32{5, 6},
-				expectedShape: []int{1, 1, 2},
-				expectedMask: []float32{
-					0, x,
-					0, 0,
-				},
-			},
-			{
-				name:          "Second seq 1",
-				in:            []float32{7, 8},
-				inShape:       []int{1, 1, 2},
-				seqs:          []int{1, 1},
-				pos:           []int32{2, 3},
-				expected:      []float32{6, 3, 4, 7, 8},
-				expectedShape: []int{1, 1, 5},
-				expectedMask: []float32{
-					0, x, x, 0, x,
-					x, x, x, 0, 0,
-				},
-			},
-			{
-				name:          "Third seq 0",
-				in:            []float32{9, 10},
-				inShape:       []int{1, 1, 2},
-				seqs:          []int{0, 0},
-				pos:           []int32{4, 5},
-				expected:      []float32{9, 10, 3, 4},
-				expectedShape: []int{1, 1, 4},
-				expectedMask: []float32{
-					0, x, x, 0,
-					0, 0, x, x,
-				},
-			},
-		}
-
-		testCache(t, backend, cache, tests)
-	})
-}
-
-func TestSWAMem(t *testing.T) {
-	runPermutedVariants(t, func(t *testing.T, backend *testBackend) {
-		cache := NewSWAMemCache(1, 3, nil)
-		defer cache.Close()
-
-		cache.Init(backend, ml.DTypeF16, 1, 16, 16)
-
-		x := float32(math.Inf(-1))
-
-		tests := []testCase{
-			{
-				name:          "FirstBatch",
-				in:            []float32{1, 2, 3, 4},
-				inShape:       []int{1, 1, 4},
-				seqs:          []int{0, 0, 0, 0},
-				pos:           []int32{0, 1, 2, 3},
-				expected:      []float32{1, 2, 3, 4},
-				expectedShape: []int{1, 1, 4},
-				expectedMask: []float32{
-					0, x, x, x,
-					0, 0, x, x,
-					x, 0, 0, x,
-					x, x, 0, 0,
-				},
-			},
-			{
-				name:          "SecondBatch",
-				in:            []float32{5, 6},
-				inShape:       []int{1, 1, 2},
-				seqs:          []int{0, 0},
-				pos:           []int32{4, 5},
-				expected:      []float32{5, 2, 3, 4, 6},
-				expectedShape: []int{1, 1, 5},
-				expectedMask: []float32{
-					0, x, x, 0, x,
-					0, x, x, x, 0,
-				},
-			},
-		}
-
-		testCache(t, backend, cache, tests)
-	})
-}
-
-func TestChunkedAttention(t *testing.T) {
-	runPermutedVariants(t, func(t *testing.T, backend *testBackend) {
-		cache := NewChunkedAttentionCache(2, nil)
-		defer cache.Close()
-
-		cache.Init(backend, ml.DTypeF16, 1, 16, 16)
-
-		x := float32(math.Inf(-1))
-
-		testCache(
-			t, backend, cache,
-			[]testCase{
-				{
-					name:          "FirstBatch",
-					in:            []float32{1, 2, 3, 4},
-					inShape:       []int{1, 1, 4},
-					seqs:          []int{0, 0, 0, 0},
-					pos:           []int32{0, 1, 2, 3},
-					expected:      []float32{1, 2, 3, 4},
-					expectedShape: []int{1, 1, 4},
-					expectedMask: []float32{
-						0, x, x, x,
-						0, 0, x, x,
-						x, x, 0, x,
-						x, x, 0, 0,
-					},
-				},
-				{
-					name:          "SecondBatch",
-					in:            []float32{5, 6, 7},
-					inShape:       []int{1, 1, 3},
-					seqs:          []int{0, 0, 0},
-					pos:           []int32{4, 5, 6},
-					expected:      []float32{1, 2, 3, 4, 5, 6, 7},
-					expectedShape: []int{1, 1, 7},
-					expectedMask: []float32{
-						x, x, x, x, 0, x, x,
-						x, x, x, x, 0, 0, x,
-						x, x, x, x, x, x, 0,
-					},
-				},
-				{
-					name:          "ThirdBatch",
-					in:            []float32{8, 9},
-					inShape:       []int{1, 1, 2},
-					seqs:          []int{0, 0},
-					pos:           []int32{7, 8},
-					expected:      []float32{1, 2, 3, 4, 5, 6, 7, 8, 9},
-					expectedShape: []int{1, 1, 9},
-					expectedMask: []float32{
-						x, x, x, x, x, x, 0, 0, x,
-						x, x, x, x, x, x, x, x, 0,
-					},
-				},
-			},
-		)
-	})
-}
-
-func TestSequences(t *testing.T) {
-	runPermutedVariants(t, func(t *testing.T, backend *testBackend) {
-		cache := NewCausalCache(nil)
-		defer cache.Close()
-
-		cache.Init(backend, ml.DTypeF16, 1, 16, 16)
-
-		tests := []testCase{
-			{
-				name:          "FirstBatch",
-				in:            []float32{1, 2, 3, 4},
-				inShape:       []int{1, 1, 4},
-				seqs:          []int{0, 0, 1, 1},
-				pos:           []int32{0, 1, 0, 1},
-				expected:      []float32{1, 2, 3, 4},
-				expectedShape: []int{1, 1, 4},
-				expectedMask:  []float32{0, float32(math.Inf(-1)), float32(math.Inf(-1)), float32(math.Inf(-1)), 0, 0, float32(math.Inf(-1)), float32(math.Inf(-1)), float32(math.Inf(-1)), float32(math.Inf(-1)), 0, float32(math.Inf(-1)), float32(math.Inf(-1)), float32(math.Inf(-1)), 0, 0},
-			},
-			{
-				name:          "SecondBatch",
-				in:            []float32{5, 6},
-				inShape:       []int{1, 1, 2},
-				seqs:          []int{0, 1},
-				pos:           []int32{2, 2},
-				expected:      []float32{1, 2, 3, 4, 5, 6},
-				expectedShape: []int{1, 1, 6},
-				expectedMask:  []float32{0, 0, float32(math.Inf(-1)), float32(math.Inf(-1)), 0, float32(math.Inf(-1)), float32(math.Inf(-1)), float32(math.Inf(-1)), 0, 0, float32(math.Inf(-1)), 0},
-			},
-		}
-
-		testCache(t, backend, cache, tests)
-	})
-}
-
-func TestRemove(t *testing.T) {
-	runPermutedVariants(t, func(t *testing.T, backend *testBackend) {
-		cache := NewCausalCache(func(ctx ml.Context, layer int, key, shift ml.Tensor) (ml.Tensor, error) {
-			return key.Add(ctx, shift), nil
-		})
-		defer cache.Close()
-
-		cache.Init(backend, ml.DTypeF16, 1, 16, 16)
-
-		x := float32(math.Inf(-1))
-
-		tests := []testCase{
-			{
-				name:          "FirstBatch",
-				in:            []float32{1, 2, 3, 4},
-				inShape:       []int{1, 1, 4},
-				seqs:          []int{0, 0, 1, 1},
-				pos:           []int32{0, 1, 0, 1},
-				expected:      []float32{1, 2, 3, 4},
-				expectedShape: []int{1, 1, 4},
-				expectedMask: []float32{
-					0, x, x, x,
-					0, 0, x, x,
-					x, x, 0, x,
-					x, x, 0, 0,
-				},
-			},
-		}
-
-		testCache(t, backend, cache, tests)
-
-		err := cache.Remove(0, 1, math.MaxInt32)
-		if err != nil {
-			panic(err)
-		}
-
-		tests = []testCase{
-			{
-				name:          "RemoveEnd",
-				in:            []float32{5, 6},
-				inShape:       []int{1, 1, 2},
-				seqs:          []int{0, 1},
-				pos:           []int32{1, 2},
-				expected:      []float32{1, 5, 3, 4, 6},
-				expectedShape: []int{1, 1, 5},
-				expectedMask: []float32{
-					0, 0, x, x, x,
-					x, x, 0, 0, 0,
-				},
-			},
-		}
-
-		testCache(t, backend, cache, tests)
-
-		err = cache.Remove(0, 0, 1)
-		if err != nil {
-			panic(err)
-		}
-
-		tests = []testCase{
-			{
-				name:          "RemoveMiddle",
-				in:            []float32{7, 8},
-				inShape:       []int{1, 1, 2},
-				seqs:          []int{0, 0},
-				pos:           []int32{1, 2},
-				expected:      []float32{7, 4, 3, 4, 6, 8},
-				expectedShape: []int{1, 1, 6},
-				expectedMask: []float32{
-					0, 0, x, x, x, x,
-					0, 0, x, x, x, 0,
-				},
-			},
-		}
-
-		testCache(t, backend, cache, tests)
-	})
-}
-
-func TestCopy(t *testing.T) {
-	runPermutedVariants(t, func(t *testing.T, backend *testBackend) {
-		cache := NewCausalCache(func(ctx ml.Context, layer int, key, shift ml.Tensor) (ml.Tensor, error) { return key, nil })
-		defer cache.Close()
-
-		cache.Init(backend, ml.DTypeF16, 1, 16, 16)
-
-		tests := []testCase{
-			{
-				name:          "FirstBatch",
-				in:            []float32{1, 2, 3, 4},
-				inShape:       []int{1, 1, 4},
-				seqs:          []int{0, 0, 0, 0},
-				pos:           []int32{0, 1, 2, 3},
-				expected:      []float32{1, 2, 3, 4},
-				expectedShape: []int{1, 1, 4},
-				expectedMask:  []float32{0, float32(math.Inf(-1)), float32(math.Inf(-1)), float32(math.Inf(-1)), 0, 0, float32(math.Inf(-1)), float32(math.Inf(-1)), 0, 0, 0, float32(math.Inf(-1)), 0, 0, 0, 0},
-			},
-		}
-
-		testCache(t, backend, cache, tests)
-
-		cache.CopyPrefix(0, 1, 2)
-
-		tests = []testCase{
-			{
-				name:          "Copy",
-				in:            []float32{5, 6},
-				inShape:       []int{1, 1, 2},
-				seqs:          []int{1, 1},
-				pos:           []int32{3, 4},
-				expected:      []float32{1, 2, 3, 4, 5, 6},
-				expectedShape: []int{1, 1, 6},
-				expectedMask:  []float32{0, 0, float32(math.Inf(-1)), float32(math.Inf(-1)), 0, float32(math.Inf(-1)), 0, 0, float32(math.Inf(-1)), float32(math.Inf(-1)), 0, 0},
-			},
-		}
-
-		testCache(t, backend, cache, tests)
-	})
-}
-
-func testCache(t *testing.T, backend ml.Backend, cache Cache, tests []testCase) {
-	for _, test := range tests {
-		t.Run(test.name, func(t *testing.T) {
-			context := backend.NewContext()
-			defer context.Close()
-
-			err := cache.StartForward(context, input.Batch{Positions: test.pos, Sequences: test.seqs}, false)
-			if err != nil {
-				panic(err)
-			}
-
-			cache.SetLayer(0)
-			tensor := context.FromFloats(test.in, test.inShape...)
-			cache.Put(context, tensor, tensor)
-
-			out, _, mask := cache.Get(context)
-
-			context.Forward(out, mask).Compute(out, mask)
-
-			if !slices.Equal(out.Floats(), test.expected) {
-				t.Errorf("TestCache: have %v; want %v", out.Floats(), test.expected)
-			}
-
-			if !slices.Equal(out.Shape(), test.expectedShape) {
-				t.Errorf("TestCache: has shape %v; want %v", out.Shape(), test.expectedShape)
-			}
-
-			if !slices.Equal(mask.Floats(), test.expectedMask) {
-				t.Errorf("TestCache: have mask: have %v want %v", mask.Floats(), test.expectedMask)
-			}
-		})
-	}
-}
-
-func TestCanResume(t *testing.T) {
-	runPermutedVariants(t, func(t *testing.T, backend *testBackend) {
-		windowSize := int32(4)
-		cache := NewSWACache(windowSize, nil)
-		defer cache.Close()
-
-		cache.Init(backend, ml.DTypeF16, 1, 16, 16)
-
-		context := backend.NewContext()
-		defer context.Close()
-
-		err := cache.StartForward(context, input.Batch{
-			Positions: []int32{0, 1, 2, 3, 4},
-			Sequences: []int{0, 0, 0, 0, 0},
-		}, false)
-		if err != nil {
-			t.Fatalf("StartForward failed: %v", err)
-		}
-
-		cache.SetLayer(0)
-		tensor := context.FromFloats([]float32{1, 2, 3, 4, 5}, 1, 1, 5)
-		cache.Put(context, tensor, tensor)
-
-		// with window size 4, nothing has slid out of the window yet
-		if !cache.CanResume(0, 0) {
-			t.Errorf("CanResume(0, 0) = false, want true (within window)")
-		}
-		if !cache.CanResume(0, 1) {
-			t.Errorf("CanResume(0, 1) = false, want true (within window)")
-		}
-		if !cache.CanResume(0, 2) {
-			t.Errorf("CanResume(0, 2) = false, want true (within window)")
-		}
-		if !cache.CanResume(0, 3) {
-			t.Errorf("CanResume(0, 3) = false, want true (latest position)")
-		}
-		if !cache.CanResume(0, 4) {
-			t.Errorf("CanResume(0, 4) = false, want true (latest position)")
-		}
-
-		// shift window by adding position 5
-		err = cache.StartForward(context, input.Batch{
-			Positions: []int32{5},
-			Sequences: []int{0},
-		}, false)
-		if err != nil {
-			t.Fatalf("StartForward failed: %v", err)
-		}
-
-		cache.SetLayer(0)
-		tensor = context.FromFloats([]float32{6}, 1, 1, 1)
-		cache.Put(context, tensor, tensor)
-
-		// only the latest position has overlapping windows
-		if cache.CanResume(0, 0) {
-			t.Errorf("after shift: CanResume(0, 0) = true, want false (outside window)")
-		}
-		if cache.CanResume(0, 1) {
-			t.Errorf("after shift: CanResume(0, 1) = true, want false (outside window)")
-		}
-		if cache.CanResume(0, 2) {
-			t.Errorf("after shift: CanResume(0, 2) = true, want false (outside window)")
-		}
-		if cache.CanResume(0, 3) {
-			t.Errorf("after shift: CanResume(0, 3) = true, want false (outside window)")
-		}
-		if cache.CanResume(0, 4) {
-			t.Errorf("after shift: CanResume(0, 4) = true, want false (outside window)")
-		}
-		if !cache.CanResume(0, 5) {
-			t.Errorf("after shift: CanResume(0, 5) = false, want true (latest position)")
-		}
-	})
-}
-
-func TestCanResumeSWAMem(t *testing.T) {
-	runPermutedVariants(t, func(t *testing.T, backend *testBackend) {
-		windowSize := int32(4)
-		memSize := int32(5)
-		cache := NewSWAMemCache(windowSize, memSize, nil)
-		defer cache.Close()
-
-		cache.Init(backend, ml.DTypeF16, 1, 16, 16)
-
-		context := backend.NewContext()
-		defer context.Close()
-
-		err := cache.StartForward(context, input.Batch{
-			Positions: []int32{0, 1, 2, 3, 4, 5, 6},
-			Sequences: []int{0, 0, 0, 0, 0, 0, 0},
-		}, false)
-		if err != nil {
-			t.Fatalf("StartForward failed: %v", err)
-		}
-
-		cache.SetLayer(0)
-		tensor := context.FromFloats([]float32{1, 2, 3, 4, 5, 6, 7}, 1, 1, 7)
-		cache.Put(context, tensor, tensor)
-
-		// shift window by adding position 7
-		err = cache.StartForward(context, input.Batch{
-			Positions: []int32{7},
-			Sequences: []int{0},
-		}, false)
-		if err != nil {
-			t.Fatalf("StartForward failed: %v", err)
-		}
-
-		cache.SetLayer(0)
-		tensor = context.FromFloats([]float32{8}, 1, 1, 1)
-		cache.Put(context, tensor, tensor)
-
-		// only the latest position has overlapping windows
-		if cache.CanResume(0, 0) {
-			t.Errorf("after shift: CanResume(0, 0) = true, want false (outside window)")
-		}
-		if cache.CanResume(0, 1) {
-			t.Errorf("after shift: CanResume(0, 1) = true, want false (outside window)")
-		}
-		if cache.CanResume(0, 2) {
-			t.Errorf("after shift: CanResume(0, 2) = true, want false (outside window)")
-		}
-		if cache.CanResume(0, 3) {
-			t.Errorf("after shift: CanResume(0, 3) = true, want false (outside window)")
-		}
-		if cache.CanResume(0, 4) {
-			t.Errorf("after shift: CanResume(0, 4) = true, want false (outside window)")
-		}
-		if cache.CanResume(0, 5) {
-			t.Errorf("after shift: CanResume(0, 5) = true, want false (outside window)")
-		}
-		if !cache.CanResume(0, 6) {
-			t.Errorf("after shift: CanResume(0, 6) = false, want true (inside window)")
-		}
-		if !cache.CanResume(0, 7) {
-			t.Errorf("after shift: CanResume(0, 7) = false, want true (latest position)")
-		}
-	})
-}
-
-type testBackend struct {
-	ml.Backend
-	permutedV bool
-}
-
-func (b *testBackend) NewContext() ml.Context {
-	return &testContext{}
-}
-
-func (b *testBackend) NewContextSize(int) ml.Context {
-	return &testContext{}
-}
-
-func (b *testBackend) CacheConfig() ml.CacheConfig {
-	return ml.CacheConfig{PermutedV: b.permutedV}
-}
-
-type testContext struct {
-	ml.Context
-}
-
-func (c *testContext) Empty(dtype ml.DType, shape ...int) ml.Tensor {
-	total := 0
-
-	if len(shape) > 0 {
-		total = 1
-		for _, s := range shape {
-			total *= s
-		}
-	}
-
-	return &testTensor{dtype: dtype, elementSize: 4, data: make([]float32, total), shape: shape}
-}
-
-func (c *testContext) Zeros(dtype ml.DType, shape ...int) ml.Tensor {
-	return c.Empty(dtype, shape...)
-}
-
-func (c *testContext) FromFloats(s []float32, shape ...int) ml.Tensor {
-	t := c.Empty(ml.DTypeF32, shape...).(*testTensor)
-
-	copy(t.data, s)
-
-	return t
-}
-
-func (c *testContext) FromInts(s []int32, shape ...int) ml.Tensor {
-	f := make([]float32, len(s))
-	for i := range f {
-		f[i] = float32(s[i])
-	}
-
-	out := c.FromFloats(f, shape...)
-	out.(*testTensor).dtype = ml.DTypeI32
-
-	return out
-}
-
-func (c *testContext) Arange(start, stop, step float32, dtype ml.DType) ml.Tensor {
-	s := make([]float32, 0, int((stop-start)/step))
-	for i := start; i < stop; i += step {
-		s = append(s, i)
-	}
-
-	out := c.FromFloats(s, len(s))
-	out.(*testTensor).dtype = dtype
-	return out
-}
-
-func (c *testContext) Input() ml.Context    { return c }
-func (c *testContext) Layer(int) ml.Context { return c }
-
-func (c *testContext) Forward(...ml.Tensor) ml.Context { return c }
-
-func (c *testContext) Compute(...ml.Tensor) {}
-
-func (c *testContext) Reserve() {}
-
-func (c *testContext) MaxGraphNodes() int {
-	return 10
-}
-
-func (c *testContext) Close() {}
-
-type testTensor struct {
-	ml.Tensor
-
-	dtype       ml.DType
-	elementSize int
-	data        []float32
-	shape       []int
-}
-
-func (t *testTensor) Dim(n int) int {
-	return t.shape[n]
-}
-
-func (t *testTensor) Stride(n int) int {
-	stride := t.elementSize
-	for i := range n {
-		stride *= t.shape[i]
-	}
-
-	return stride
-}
-
-func (t *testTensor) Shape() []int {
-	return t.shape
-}
-
-func (t *testTensor) DType() ml.DType {
-	return t.dtype
-}
-
-func (t *testTensor) Floats() []float32 {
-	out := make([]float32, len(t.data))
-	copy(out, t.data)
-	return out
-}
-
-func (t *testTensor) Neg(ctx ml.Context) ml.Tensor {
-	out := ctx.Empty(t.DType(), t.Shape()...).(*testTensor)
-	for i := range out.data {
-		out.data[i] = -t.data[i]
-	}
-	return out
-}
-
-func (t *testTensor) Add(ctx ml.Context, t2 ml.Tensor) ml.Tensor {
-	out := ctx.Empty(t.DType(), t.Shape()...).(*testTensor)
-
-	for i := range out.data {
-		out.data[i] = t.data[i] + t2.(*testTensor).data[i]
-	}
-
-	return out
-}
-
-func (t *testTensor) Reshape(ctx ml.Context, shape ...int) ml.Tensor {
-	return &testTensor{
-		dtype:       t.dtype,
-		elementSize: t.elementSize,
-		data:        t.data,
-		shape:       shape,
-	}
-}
-
-func (t *testTensor) View(ctx ml.Context, offset int, shape ...int) ml.Tensor {
-	offset /= t.elementSize
-
-	var s []int
-
-	switch len(shape) {
-	case 1:
-		s = []int{shape[0]}
-	case 3:
-		s = []int{shape[0], shape[2]}
-	case 5:
-		s = []int{shape[0], shape[2], shape[4]}
-	default:
-		panic("unsupported number of dimensions")
-	}
-
-	context := &testContext{}
-
-	view := context.Empty(t.dtype, s...).(*testTensor)
-	view.data = t.data[offset : offset+len(view.data)]
-
-	return view
-}
-
-func (t *testTensor) Permute(ctx ml.Context, order ...int) ml.Tensor {
-	if len(t.shape) > 4 || len(order) > 4 {
-		panic("permute only supports up to 4 dimensions")
-	}
-
-	if len(order) != len(t.shape) && len(order) != 4 {
-		panic("invalid number of dimensions for permute")
-	}
-
-	// ggml_permute expects 4 axes, so fill in any missing dimensions.
-	orderFull := append(make([]int, 0, 4), order...)
-	for len(orderFull) < 4 {
-		orderFull = append(orderFull, len(orderFull))
-	}
-
-	seen := [4]bool{}
-
-	shape4 := [4]int{1, 1, 1, 1}
-	for i := 0; i < len(t.shape) && i < 4; i++ {
-		shape4[i] = t.shape[i]
-	}
-
-	newShape4 := [4]int{1, 1, 1, 1}
-	for axis := range 4 {
-		dst := orderFull[axis]
-		if dst < 0 || dst >= 4 {
-			panic("invalid axis for permute")
-		}
-		if seen[dst] {
-			panic("duplicate axis for permute")
-		}
-		seen[dst] = true
-		newShape4[dst] = shape4[axis]
-	}
-
-	total := len(t.data)
-	newData := make([]float32, total)
-
-	if total > 0 {
-		oldDims := shape4
-		newDims := newShape4
-
-		oldStride := [4]int{1, 1, 1, 1}
-		newStride := [4]int{1, 1, 1, 1}
-		for i := 1; i < 4; i++ {
-			oldStride[i] = oldStride[i-1] * oldDims[i-1]
-			newStride[i] = newStride[i-1] * newDims[i-1]
-		}
-
-		var coords [4]int
-		var newCoords [4]int
-
-		for idx := range total {
-			remainder := idx
-			for axis := range 4 {
-				dim := oldDims[axis]
-				if dim == 0 {
-					coords[axis] = 0
-					continue
-				}
-				coords[axis] = remainder % dim
-				remainder /= dim
-			}
-
-			for axis := range 4 {
-				newCoords[orderFull[axis]] = coords[axis]
-			}
-
-			newIndex := 0
-			for axis := range 4 {
-				if newDims[axis] == 0 {
-					continue
-				}
-				newIndex += newCoords[axis] * newStride[axis]
-			}
-
-			newData[newIndex] = t.data[idx]
-		}
-	}
-
-	numDims := 4
-	for numDims > 1 && newShape4[numDims-1] <= 1 {
-		numDims--
-	}
-
-	newShape := make([]int, numDims)
-	copy(newShape, newShape4[:numDims])
-
-	return &testTensor{
-		dtype:       t.dtype,
-		elementSize: t.elementSize,
-		data:        newData,
-		shape:       newShape,
-	}
-}
-
-func (t *testTensor) SetRows(ctx ml.Context, src ml.Tensor, idxs ml.Tensor) ml.Tensor {
-	dst := t
-	srcTensor := src.(*testTensor)
-	idxTensor := idxs.(*testTensor)
-
-	shapeTo4D := func(shape []int) [4]int {
-		out := [4]int{1, 1, 1, 1}
-		for i := 0; i < len(shape) && i < 4; i++ {
-			out[i] = shape[i]
-		}
-		return out
-	}
-
-	computeStrides := func(shape [4]int) [4]int {
-		out := [4]int{1, 1, 1, 1}
-		for i := 1; i < 4; i++ {
-			out[i] = out[i-1] * shape[i-1]
-		}
-		return out
-	}
-
-	dstShape4D := shapeTo4D(dst.shape)
-	srcShape4D := shapeTo4D(srcTensor.shape)
-	idxShape4D := shapeTo4D(idxTensor.shape)
-
-	if dstShape4D[0] != srcShape4D[0] || dstShape4D[2] != srcShape4D[2] || dstShape4D[3] != srcShape4D[3] {
-		panic("SetRows requires matching tensor shapes")
-	}
-
-	if srcShape4D[1] != idxShape4D[0] {
-		panic("SetRows rows/index mismatch")
-	}
-
-	if srcShape4D[2]%idxShape4D[1] != 0 || srcShape4D[3]%idxShape4D[2] != 0 {
-		panic("SetRows cannot broadcast indices")
-	}
-
-	if idxShape4D[3] != 1 {
-		panic("SetRows expects 1D or 2D index tensors")
-	}
-
-	dstStride := computeStrides(dstShape4D)
-	srcStride := computeStrides(srcShape4D)
-	idxStride := computeStrides(idxShape4D)
-
-	numColumns := srcShape4D[0]
-	numRows := srcShape4D[1]
-
-	for dim3Index := range dstShape4D[3] {
-		for dim2Index := range dstShape4D[2] {
-			idxDim2 := 0
-			idxDim3 := 0
-			if idxShape4D[1] > 0 {
-				idxDim2 = dim2Index % idxShape4D[1]
-			}
-			if idxShape4D[2] > 0 {
-				idxDim3 = dim3Index % idxShape4D[2]
-			}
-
-			idxBase := idxDim3*idxStride[2] + idxDim2*idxStride[1]
-			srcBase := dim3Index*srcStride[3] + dim2Index*srcStride[2]
-			dstBase := dim3Index*dstStride[3] + dim2Index*dstStride[2]
-
-			for row := range numRows {
-				idx := int(idxTensor.data[idxBase+row*idxStride[0]])
-				if idx < 0 || idx >= dstShape4D[1] {
-					panic("SetRows index out of range")
-				}
-
-				srcOffset := srcBase + row*srcStride[1]
-				dstOffset := dstBase + idx*dstStride[1]
-
-				copy(dst.data[dstOffset:dstOffset+numColumns], srcTensor.data[srcOffset:srcOffset+numColumns])
-			}
-		}
-	}
-
-	return dst
-}
-
-func (t *testTensor) Copy(ctx ml.Context, t2 ml.Tensor) ml.Tensor {
-	copy(t2.(*testTensor).data, t.data)
-	return nil
-}

kvcache/encoder.go

@@ -1,156 +0,0 @@
-package kvcache
-
-import (
-	"fmt"
-
-	"github.com/ollama/ollama/ml"
-	"github.com/ollama/ollama/model/input"
-)
-
-// Encoder cache stores K and V tensors that are position independent
-//
-// The tensors can be of any shape and will be returned as they were stored
-// The mask is currently always nil
-//
-// Not currently safe for multiple sequences
-type EncoderCache struct {
-	// config controls mostly backend-specific optimizations
-	config *ml.CacheConfig
-
-	// ** current forward pass **
-
-	// the active layer for Get and Put
-	curLayer int
-
-	// if something is stored during this pass, this
-	// will be the position (but there is no guarantee
-	// anything will be stored)
-	curPos int32
-
-	// curReserve indicates that this forward pass is only for
-	// memory reservation and we should not update our metadata
-	// based on it.
-	curReserve bool
-
-	// ** cache metadata **
-
-	// was something stored in the cache?
-	encoderCached bool
-
-	// position of the cached data
-	encoderPos int32
-
-	// ** cache data storage **
-	backend      ml.Backend
-	ctxs         map[int]ml.Context
-	keys, values map[int]ml.Tensor
-}
-
-func NewEncoderCache() *EncoderCache {
-	return &EncoderCache{
-		ctxs:   make(map[int]ml.Context),
-		keys:   make(map[int]ml.Tensor),
-		values: make(map[int]ml.Tensor),
-	}
-}
-
-func (c *EncoderCache) Init(backend ml.Backend, dtype ml.DType, maxSequences, capacity, maxBatch int) {
-	if c.config == nil {
-		var config ml.CacheConfig
-		if cc, ok := backend.(ml.BackendCacheConfig); ok {
-			config = cc.CacheConfig()
-		}
-		c.config = &config
-	}
-
-	if maxSequences > 1 {
-		panic(fmt.Errorf("encoder cache does not support multiple sequences; requested: %v", maxSequences))
-	}
-
-	if c.config.CachePadding != 0 && c.config.CachePadding != 1 {
-		panic(fmt.Errorf("encoder cache is unable to enforce requested CachePadding (%v)", c.config.CachePadding))
-	}
-
-	c.backend = backend
-}
-
-func (c *EncoderCache) SetConfig(config ml.CacheConfig) {
-	if c.config != nil {
-		panic("config cannot be changed after being previously set, either by the model or backend")
-	}
-
-	c.config = &config
-}
-
-func (c *EncoderCache) Close() {
-	for _, ctx := range c.ctxs {
-		ctx.Close()
-	}
-}
-
-func (c *EncoderCache) StartForward(ctx ml.Context, batch input.Batch, reserve bool) error {
-	// We work with the most recent image
-	if len(batch.Multimodal) > 0 {
-		c.curPos = batch.Positions[batch.Multimodal[len(batch.Multimodal)-1].Index]
-	}
-
-	c.curReserve = reserve
-
-	return nil
-}
-
-func (c *EncoderCache) SetLayer(layer int) {
-	c.curLayer = layer
-}
-
-func (c *EncoderCache) EncoderCached() bool {
-	return c.encoderCached
-}
-
-func (c *EncoderCache) Get(ctx ml.Context) (ml.Tensor, ml.Tensor, ml.Tensor) {
-	return c.keys[c.curLayer], c.values[c.curLayer], nil
-}
-
-func (c *EncoderCache) Put(ctx ml.Context, key, value ml.Tensor) {
-	if !c.curReserve {
-		c.encoderPos = c.curPos
-		c.encoderCached = true
-	}
-
-	if c.config.PermutedV {
-		value = value.Permute(ctx, 1, 2, 0, 3)
-	}
-
-	if _, ok := c.ctxs[c.curLayer]; !ok {
-		c.ctxs[c.curLayer] = c.backend.NewContextSize(2).Layer(c.curLayer)
-	}
-
-	if _, ok := c.keys[c.curLayer]; !ok {
-		c.keys[c.curLayer] = c.ctxs[c.curLayer].Empty(key.DType(), key.Shape()...)
-	}
-
-	if _, ok := c.values[c.curLayer]; !ok {
-		c.values[c.curLayer] = c.ctxs[c.curLayer].Empty(value.DType(), value.Shape()...)
-	}
-
-	ctx.Forward(
-		key.Copy(ctx, c.keys[c.curLayer]),
-		value.Copy(ctx, c.values[c.curLayer]),
-	)
-}
-
-func (c *EncoderCache) CopyPrefix(srcSeq, dstSeq int, len int32) {
-	panic("encoder cache does not support multiple sequences")
-}
-
-func (c *EncoderCache) CanResume(seq int, pos int32) bool {
-	return true
-}
-
-func (c *EncoderCache) Remove(seq int, beginIndex, endIndex int32) error {
-	if c.encoderPos >= beginIndex && c.encoderPos < endIndex {
-		c.encoderCached = false
-	}
-
-	return nil
-}

kvcache/recurrent.go

@@ -1,752 +0,0 @@
-package kvcache
-
-import (
-	"errors"
-	"fmt"
-	"math"
-	"slices"
-
-	"github.com/ollama/ollama/ml"
-	"github.com/ollama/ollama/model/input"
-)
-
-const (
-	DefaultCheckpointCount    = 24
-	DefaultCheckpointMinPos   = int32(16)
-	DefaultCheckpointInterval = int32(1664)
-)
-
-var ErrInvalidRecurrentShape = errors.New("kvcache: invalid recurrent state shape")
-
-// Config configures a shared hybrid recurrent cache.
-type RecurrentConfig struct {
-	Shift               func(ctx ml.Context, layer int, key, shift ml.Tensor) (ml.Tensor, error)
-	ConvDim             int
-	ConvChannels        int
-	RecurrentStateSize  int
-	CheckpointLogPrefix string
-}
-
-var (
-	_ Cache           = (*Recurrent)(nil)
-	_ CheckpointCache = (*Recurrent)(nil)
-)
-
-// Cache stores:
-// - a standard causal KV cache
-// - per-sequence conv state for recurrent operators
-// - per-sequence recurrent state for recurrent operators
-//
-// Conv state shape (per layer, per sequence): [convDim, convChannels]
-// Recurrent state shape (per layer, per sequence): [recurrentStateSize]
-type Recurrent struct {
-	kv *Causal
-
-	backend      ml.Backend
-	dtype        ml.DType
-	maxSequences int
-
-	// Conv state dimensions
-	convDim      int
-	convChannels int
-
-	// Recurrent state dimensions
-	recurrentStateSize int
-
-	logPrefix string
-
-	// slot mapping for recurrent state (copy-on-write)
-	slotForSeq  map[int]int
-	refCount    []int
-	freeSlots   []int
-	seqCounts   map[int]int
-	slotScratch [1]int32
-
-	// per-layer conv state buffers (allocated lazily)
-	convCtxs   map[int]ml.Context
-	convStates map[int]ml.Tensor // [convDim*convChannels, maxSlots]
-
-	// per-layer recurrent state buffers (allocated lazily)
-	recurrentCtxs   map[int]ml.Context
-	recurrentStates map[int]ml.Tensor // [recurrentStateSize, maxSlots]
-
-	// recurrent checkpoints (per slot)
-	checkpointCount     int
-	checkpointMinPos    int32
-	checkpointInterval  int32
-	checkpointCtxSize   int
-	checkpoints         map[int]*slotCheckpointStore
-	pendingRestore      map[int]checkpointRestore
-	curCheckpointPos    []int32
-	curCheckpointSlots  map[int]int
-	reserveCheckpoints  bool
-	checkpointConvCtxs  map[int]ml.Context
-	checkpointRecurCtxs map[int]ml.Context
-	checkpointReserved  map[int]struct{}
-
-	// current forward batch (derived in StartForward)
-	curSeqs       []int
-	curSlots      []int
-	curSlotsInput ml.Tensor
-	curSeqTokens  int
-
-	// track if EnsureWritable has been called for this forward pass
-	writableEnsured bool
-	writableError   error
-}
-
-func NewRecurrentCache(config RecurrentConfig) *Recurrent {
-	return &Recurrent{
-		kv:                  NewCausalCache(config.Shift),
-		convDim:             config.ConvDim,
-		convChannels:        config.ConvChannels,
-		recurrentStateSize:  config.RecurrentStateSize,
-		logPrefix:           config.CheckpointLogPrefix,
-		slotForSeq:          make(map[int]int),
-		seqCounts:           make(map[int]int),
-		convCtxs:            make(map[int]ml.Context),
-		convStates:          make(map[int]ml.Tensor),
-		recurrentCtxs:       make(map[int]ml.Context),
-		recurrentStates:     make(map[int]ml.Tensor),
-		checkpointCount:     DefaultCheckpointCount,
-		checkpointMinPos:    DefaultCheckpointMinPos,
-		checkpointInterval:  DefaultCheckpointInterval,
-		checkpoints:         make(map[int]*slotCheckpointStore),
-		pendingRestore:      make(map[int]checkpointRestore),
-		curCheckpointSlots:  make(map[int]int),
-		checkpointConvCtxs:  make(map[int]ml.Context),
-		checkpointRecurCtxs: make(map[int]ml.Context),
-		checkpointReserved:  make(map[int]struct{}),
-	}
-}
-
-func (c *Recurrent) Init(backend ml.Backend, dtype ml.DType, maxSequences, capacity, maxBatch int) {
-	c.backend = backend
-	c.dtype = dtype
-	c.maxSequences = maxSequences
-	c.checkpoints = make(map[int]*slotCheckpointStore)
-	c.pendingRestore = make(map[int]checkpointRestore)
-	c.curCheckpointPos = c.curCheckpointPos[:0]
-	c.curCheckpointSlots = make(map[int]int)
-	c.checkpointReserved = make(map[int]struct{})
-	c.checkpointCtxSize = c.checkpointCount * c.maxSequences
-	if c.checkpointCtxSize < 8 {
-		c.checkpointCtxSize = 8
-	}
-
-	// initialize slot allocator
-	c.refCount = make([]int, maxSequences)
-	c.freeSlots = c.freeSlots[:0]
-	for i := maxSequences - 1; i >= 0; i-- {
-		c.freeSlots = append(c.freeSlots, i)
-	}
-
-	c.kv.Init(backend, dtype, maxSequences, capacity, maxBatch)
-}
-
-func (c *Recurrent) Close() {
-	for _, ctx := range c.convCtxs {
-		ctx.Close()
-	}
-	for _, ctx := range c.recurrentCtxs {
-		ctx.Close()
-	}
-	for _, ctx := range c.checkpointConvCtxs {
-		ctx.Close()
-	}
-	for _, ctx := range c.checkpointRecurCtxs {
-		ctx.Close()
-	}
-	c.kv.Close()
-}
-
-func (c *Recurrent) SetConfig(config ml.CacheConfig) {
-	c.kv.SetConfig(config)
-}
-
-func (c *Recurrent) SetLayer(layer int) {
-	c.kv.SetLayer(layer)
-}
-
-func (c *Recurrent) Get(ctx ml.Context) (ml.Tensor, ml.Tensor, ml.Tensor) {
-	return c.kv.Get(ctx)
-}
-
-func (c *Recurrent) Put(ctx ml.Context, key, value ml.Tensor) {
-	c.kv.Put(ctx, key, value)
-}
-
-func (c *Recurrent) StartForward(ctx ml.Context, batch input.Batch, reserve bool) error {
-	if err := c.kv.StartForward(ctx, batch, reserve); err != nil {
-		return err
-	}
-
-	nTokens := len(batch.Sequences)
-	if nTokens == 0 {
-		c.curSeqs = c.curSeqs[:0]
-		c.curSlots = c.curSlots[:0]
-		c.curSlotsInput = nil
-		c.curSeqTokens = 0
-		c.reserveCheckpoints = false
-		c.writableEnsured = false
-		c.writableError = nil
-		return nil
-	}
-
-	// Fast path for single-sequence batches (common during decode and prefill).
-	firstSeq := batch.Sequences[0]
-	singleSeq := true
-	for _, s := range batch.Sequences[1:] {
-		if s != firstSeq {
-			singleSeq = false
-			break
-		}
-	}
-	if singleSeq {
-		return c.startForwardSingleSeq(ctx, firstSeq, nTokens, batch, reserve)
-	}
-
-	// Derive equal-length sequence layout for recurrent layers.
-	seqCounts := c.seqCounts
-	for s := range seqCounts {
-		delete(seqCounts, s)
-	}
-
-	c.curSeqs = c.curSeqs[:0]
-	for _, s := range batch.Sequences {
-		if seqCounts[s] == 0 {
-			c.curSeqs = append(c.curSeqs, s)
-		}
-		seqCounts[s]++
-	}
-
-	nSeqs := len(c.curSeqs)
-	want := nTokens / nSeqs
-	for _, s := range c.curSeqs {
-		if seqCounts[s] != want {
-			return ErrNotSupported
-		}
-	}
-
-	c.curSeqTokens = want
-
-	if reserve {
-		c.curSlots = c.curSlots[:0]
-		for i := range nSeqs {
-			c.curSlots = append(c.curSlots, i)
-		}
-		c.finalizeStartForward(ctx, batch, true)
-		return nil
-	}
-
-	// Ensure slots exist for sequences in this batch.
-	c.curSlots = c.curSlots[:0]
-	var newSlots []int
-	for _, s := range c.curSeqs {
-		slot, ok := c.slotForSeq[s]
-		if !ok {
-			var err error
-			slot, err = c.allocSlot()
-			if err != nil {
-				return err
-			}
-			c.slotForSeq[s] = slot
-			c.refCount[slot] = 1
-			newSlots = append(newSlots, slot)
-		}
-		c.curSlots = append(c.curSlots, slot)
-	}
-
-	if len(newSlots) > 0 {
-		c.zeroSlots(ctx, newSlots)
-	}
-
-	c.finalizeStartForward(ctx, batch, false)
-
-	return nil
-}
-
-func (c *Recurrent) startForwardSingleSeq(ctx ml.Context, seq, seqTokens int, batch input.Batch, reserve bool) error {
-	c.curSeqs = append(c.curSeqs[:0], seq)
-	c.curSeqTokens = seqTokens
-
-	if reserve {
-		c.curSlots = append(c.curSlots[:0], 0)
-		c.finalizeStartForward(ctx, batch, true)
-		return nil
-	}
-
-	slot, ok := c.slotForSeq[seq]
-	if !ok {
-		var err error
-		slot, err = c.allocSlot()
-		if err != nil {
-			return err
-		}
-
-		c.slotForSeq[seq] = slot
-		c.refCount[slot] = 1
-		slotList := [1]int{slot}
-		c.zeroSlots(ctx, slotList[:])
-	}
-
-	c.curSlots = append(c.curSlots[:0], slot)
-	c.finalizeStartForward(ctx, batch, false)
-
-	return nil
-}
-
-func (c *Recurrent) finalizeStartForward(ctx ml.Context, batch input.Batch, reserve bool) {
-	c.setCurSlotsInput(ctx)
-	c.writableEnsured = false
-	c.writableError = nil
-	c.reserveCheckpoints = reserve
-	c.planCheckpoints(batch)
-}
-
-func (c *Recurrent) setCurSlotsInput(ctx ml.Context) {
-	c.curSlotsInput = c.slotsInput(ctx, c.curSlots)
-}
-
-func (c *Recurrent) slotsInput(ctx ml.Context, slots []int) ml.Tensor {
-	switch len(slots) {
-	case 0:
-		return nil
-	case 1:
-		c.slotScratch[0] = int32(slots[0])
-		return ctx.Input().FromInts(c.slotScratch[:], 1)
-	default:
-		slotIndices := make([]int32, len(slots))
-		for i, v := range slots {
-			slotIndices[i] = int32(v)
-		}
-		return ctx.Input().FromInts(slotIndices, len(slotIndices))
-	}
-}
-
-func (c *Recurrent) allocSlot() (int, error) {
-	if len(c.freeSlots) == 0 {
-		return 0, ErrKvCacheFull
-	}
-	slot := c.freeSlots[len(c.freeSlots)-1]
-	c.freeSlots = c.freeSlots[:len(c.freeSlots)-1]
-	return slot, nil
-}
-
-func (c *Recurrent) freeSlot(slot int) {
-	if slot >= 0 && slot < c.maxSequences {
-		c.freeSlots = append(c.freeSlots, slot)
-	}
-}
-
-// zeroSlots zeros recurrent state for the given slots across all cached layers.
-func (c *Recurrent) zeroSlots(ctx ml.Context, slots []int) {
-	if len(slots) == 0 {
-		return
-	}
-
-	inputCtx := ctx.Input()
-	slotsTensor := c.slotsInput(ctx, slots)
-
-	if len(c.convStates) > 0 {
-		zeros := inputCtx.Zeros(ml.DTypeF32, c.convDim*c.convChannels, len(slots))
-		for _, buf := range c.convStates {
-			ctx.Forward(buf.SetRows(ctx, zeros, slotsTensor))
-		}
-	}
-
-	if len(c.recurrentStates) > 0 {
-		zeros := inputCtx.Zeros(ml.DTypeF32, c.recurrentStateSize, len(slots))
-		for _, buf := range c.recurrentStates {
-			ctx.Forward(buf.SetRows(ctx, zeros, slotsTensor))
-		}
-	}
-}
-
-// EnsureWritable ensures sequences have private slots (copy-on-write).
-func (c *Recurrent) EnsureWritable(ctx ml.Context) error {
-	for i, seq := range c.curSeqs {
-		slot, ok := c.slotForSeq[seq]
-		if !ok {
-			continue
-		}
-
-		if slot < 0 || slot >= len(c.refCount) {
-			continue
-		}
-
-		if c.refCount[slot] <= 1 {
-			continue
-		}
-
-		newSlot, err := c.allocSlot()
-		if err != nil {
-			return err
-		}
-		c.refCount[slot]--
-		c.refCount[newSlot] = 1
-		c.slotForSeq[seq] = newSlot
-		c.curSlots[i] = newSlot
-
-		c.copyRecurrentState(ctx, slot, newSlot)
-		c.copyCheckpoints(ctx, slot, newSlot)
-	}
-
-	c.setCurSlotsInput(ctx)
-
-	return nil
-}
-
-func (c *Recurrent) copyRecurrentState(ctx ml.Context, srcSlot, dstSlot int) {
-	src := ctx.Input().FromInts([]int32{int32(srcSlot)}, 1)
-	dst := ctx.Input().FromInts([]int32{int32(dstSlot)}, 1)
-
-	for _, buf := range c.convStates {
-		rows := buf.Rows(ctx, src)
-		if rows.DType() != ml.DTypeF32 {
-			rows = rows.Cast(ctx, ml.DTypeF32)
-		}
-		ctx.Forward(buf.SetRows(ctx, rows, dst))
-	}
-
-	for _, buf := range c.recurrentStates {
-		rows := buf.Rows(ctx, src)
-		if rows.DType() != ml.DTypeF32 {
-			rows = rows.Cast(ctx, ml.DTypeF32)
-		}
-		ctx.Forward(buf.SetRows(ctx, rows, dst))
-	}
-}
-
-func (c *Recurrent) CopyPrefix(srcSeq, dstSeq int, prefixLen int32) {
-	c.kv.CopyPrefix(srcSeq, dstSeq, prefixLen)
-
-	if dstSlot, ok := c.slotForSeq[dstSeq]; ok {
-		if c.validSlot(dstSlot) {
-			c.refCount[dstSlot]--
-			if c.refCount[dstSlot] <= 0 {
-				c.refCount[dstSlot] = 0
-				c.freeSlot(dstSlot)
-			}
-		}
-		delete(c.slotForSeq, dstSeq)
-	}
-
-	srcSlot, ok := c.slotForSeq[srcSeq]
-	if !ok {
-		return
-	}
-
-	if c.validSlot(srcSlot) {
-		c.slotForSeq[dstSeq] = srcSlot
-		c.refCount[srcSlot]++
-	}
-}
-
-func (c *Recurrent) CanResume(seq int, pos int32) bool {
-	if !c.kv.CanResume(seq, pos) {
-		return false
-	}
-	if pos == 0 {
-		return true
-	}
-	return c.hasCheckpoint(seq, pos)
-}
-
-func (c *Recurrent) Remove(seq int, beginIndex, endIndex int32) error {
-	if beginIndex > 0 && endIndex != math.MaxInt32 {
-		if err := c.kv.Remove(seq, beginIndex, endIndex); err != nil {
-			return err
-		}
-		delete(c.pendingRestore, seq)
-
-		slot, ok := c.slotForSeq[seq]
-		if !ok || !c.validSlot(slot) {
-			return nil
-		}
-
-		// Detach shared recurrent state/checkpoints before mutating checkpoint positions.
-		if c.refCount[slot] > 1 {
-			newSlot, err := c.allocSlot()
-			if err != nil {
-				return err
-			}
-			ctx := c.backend.NewContext()
-			c.copyRecurrentState(ctx, slot, newSlot)
-			c.copyCheckpoints(ctx, slot, newSlot)
-			if len(c.convStates) > 0 || len(c.recurrentStates) > 0 {
-				ctx.Compute()
-			}
-			ctx.Close()
-
-			c.refCount[slot]--
-			c.refCount[newSlot] = 1
-			c.slotForSeq[seq] = newSlot
-			slot = newSlot
-		}
-
-		c.shiftCheckpoints(slot, beginIndex, endIndex)
-		return nil
-	}
-
-	if beginIndex > 0 {
-		restore, ok := c.pendingRestore[seq]
-		if !ok || restore.pos+1 != beginIndex {
-			return ErrNotSupported
-		}
-		if !c.restoreComplete(restore) {
-			return ErrNotSupported
-		}
-		if slot, ok := c.slotForSeq[seq]; ok && c.validSlot(slot) && c.refCount[slot] > 1 {
-			newSlot, err := c.allocSlot()
-			if err != nil {
-				return err
-			}
-			ctx := c.backend.NewContext()
-			c.copyRecurrentState(ctx, slot, newSlot)
-			c.copyCheckpoints(ctx, slot, newSlot)
-			if len(c.convStates) > 0 || len(c.recurrentStates) > 0 {
-				ctx.Compute()
-			}
-			ctx.Close()
-
-			c.refCount[slot]--
-			c.refCount[newSlot] = 1
-			c.slotForSeq[seq] = newSlot
-
-			restore.slot = newSlot
-			c.pendingRestore[seq] = restore
-		}
-	}
-
-	if err := c.kv.Remove(seq, beginIndex, endIndex); err != nil {
-		return err
-	}
-
-	if beginIndex > 0 {
-		restore := c.pendingRestore[seq]
-		delete(c.pendingRestore, seq)
-		return c.applyCheckpointRestore(restore)
-	}
-
-	slot, ok := c.slotForSeq[seq]
-	delete(c.pendingRestore, seq)
-	if !ok {
-		return nil
-	}
-
-	if !c.validSlot(slot) {
-		delete(c.slotForSeq, seq)
-		return nil
-	}
-
-	c.refCount[slot]--
-	if c.refCount[slot] <= 0 {
-		c.refCount[slot] = 0
-		c.clearCheckpoints(slot)
-		c.freeSlot(slot)
-	}
-	delete(c.slotForSeq, seq)
-
-	return nil
-}
-
-func (c *Recurrent) validSlot(slot int) bool {
-	return slot >= 0 && slot < len(c.refCount)
-}
-
-func (c *Recurrent) SlotsTensor() ml.Tensor {
-	return c.curSlotsInput
-}
-
-// contiguousSlots returns the starting slot if current slots are contiguous and ordered.
-func (c *Recurrent) contiguousSlots() (int, bool) {
-	if len(c.curSlots) == 0 {
-		return 0, false
-	}
-	start := c.curSlots[0]
-	for i, s := range c.curSlots {
-		if s != start+i {
-			return 0, false
-		}
-	}
-	return start, true
-}
-
-func (c *Recurrent) SeqTokens() int {
-	return c.curSeqTokens
-}
-
-func (c *Recurrent) NumSeqs() int {
-	return len(c.curSeqs)
-}
-
-func (c *Recurrent) convBuffer(layer int) ml.Tensor {
-	if buf, ok := c.convStates[layer]; ok {
-		return buf
-	}
-
-	if _, ok := c.convCtxs[layer]; !ok {
-		c.convCtxs[layer] = c.backend.NewContextSize(1).Layer(layer)
-	}
-
-	buf := c.convCtxs[layer].Zeros(ml.DTypeF32, c.convDim*c.convChannels, c.maxSequences)
-	c.convStates[layer] = buf
-	return buf
-}
-
-func (c *Recurrent) recurrentBuffer(layer int) ml.Tensor {
-	if buf, ok := c.recurrentStates[layer]; ok {
-		return buf
-	}
-
-	if _, ok := c.recurrentCtxs[layer]; !ok {
-		c.recurrentCtxs[layer] = c.backend.NewContextSize(1).Layer(layer)
-	}
-
-	buf := c.recurrentCtxs[layer].Zeros(ml.DTypeF32, c.recurrentStateSize, c.maxSequences)
-	c.recurrentStates[layer] = buf
-	return buf
-}
-
-func (c *Recurrent) ensureWritable(ctx ml.Context) error {
-	c.ensureWritableOnce(ctx)
-	return c.writableError
-}
-
-func (c *Recurrent) currentSlotRows(ctx ml.Context, buf ml.Tensor, rowSize int) ml.Tensor {
-	if start, ok := c.contiguousSlots(); ok {
-		offset := start * buf.Stride(1)
-		return buf.View(ctx, offset, rowSize, buf.Stride(1), c.NumSeqs())
-	}
-
-	return buf.Rows(ctx, c.SlotsTensor())
-}
-
-func (c *Recurrent) writeCurrentSlotRows(ctx ml.Context, buf ml.Tensor, rowSize int, src ml.Tensor) {
-	if start, ok := c.contiguousSlots(); ok {
-		offset := start * buf.Stride(1)
-		view := buf.View(ctx, offset, rowSize, buf.Stride(1), c.NumSeqs())
-		ctx.Forward(src.Copy(ctx, view))
-		return
-	}
-
-	ctx.Forward(buf.SetRows(ctx, src, c.SlotsTensor()))
-}
-
-func (c *Recurrent) ensureWritableOnce(ctx ml.Context) {
-	if !c.writableEnsured {
-		needsWritable := false
-		for _, seq := range c.curSeqs {
-			slot, ok := c.slotForSeq[seq]
-			if !ok {
-				continue
-			}
-			if slot >= 0 && slot < len(c.refCount) && c.refCount[slot] > 1 {
-				needsWritable = true
-				break
-			}
-		}
-
-		if needsWritable {
-			if err := c.EnsureWritable(ctx); err != nil {
-				c.writableError = err
-			}
-		}
-		c.writableEnsured = true
-	}
-}
-
-// ConvState returns conv state for current batch sequences as [convDim, convChannels, nSeqs].
-func (c *Recurrent) ConvState(ctx ml.Context, layer int) (ml.Tensor, error) {
-	if err := c.ensureWritable(ctx); err != nil {
-		return nil, err
-	}
-
-	buf := c.convBuffer(layer)
-	cur := c.currentSlotRows(ctx, buf, c.convDim*c.convChannels)
-	return cur.Reshape(ctx, c.convDim, c.convChannels, c.NumSeqs()), nil
-}
-
-// UpdateConvState writes new conv state for current batch sequences.
-func (c *Recurrent) UpdateConvState(ctx ml.Context, layer int, newState ml.Tensor) {
-	buf := c.convBuffer(layer)
-	src := newState.Reshape(ctx, c.convDim*c.convChannels, c.NumSeqs())
-	srcF32 := src
-	if src.DType() != ml.DTypeF32 {
-		srcF32 = src.Cast(ctx, ml.DTypeF32)
-	}
-	c.writeCurrentSlotRows(ctx, buf, c.convDim*c.convChannels, srcF32)
-
-	c.captureConvCheckpoint(ctx, layer, srcF32)
-}
-
-// RecurrentState returns recurrent state for current batch sequences with shape [dims..., nSeqs].
-func (c *Recurrent) RecurrentState(ctx ml.Context, layer int, dims ...int) (ml.Tensor, error) {
-	if err := c.ensureWritable(ctx); err != nil {
-		return nil, err
-	}
-	if len(dims) == 0 {
-		return nil, ErrInvalidRecurrentShape
-	}
-
-	size := 1
-	for _, d := range dims {
-		if d <= 0 {
-			return nil, ErrInvalidRecurrentShape
-		}
-		size *= d
-	}
-	if size != c.recurrentStateSize {
-		return nil, fmt.Errorf("%w: got %v (size %d), want size %d", ErrInvalidRecurrentShape, dims, size, c.recurrentStateSize)
-	}
-
-	buf := c.recurrentBuffer(layer)
-	cur := c.currentSlotRows(ctx, buf, c.recurrentStateSize)
-	shape := make([]int, 0, len(dims)+1)
-	shape = append(shape, dims...)
-	shape = append(shape, c.NumSeqs())
-	return cur.Reshape(ctx, shape...), nil
-}
-
-// RecurrentState4D returns recurrent state as [dim0, dim1, dim2, nSeqs].
-func (c *Recurrent) RecurrentState4D(ctx ml.Context, layer int, dim0, dim1, dim2 int) (ml.Tensor, error) {
-	if err := c.ensureWritable(ctx); err != nil {
-		return nil, err
-	}
-	if dim0 <= 0 || dim1 <= 0 || dim2 <= 0 {
-		return nil, ErrInvalidRecurrentShape
-	}
-
-	size := dim0 * dim1 * dim2
-	if size != c.recurrentStateSize {
-		return nil, fmt.Errorf("%w: got [%d %d %d] (size %d), want size %d", ErrInvalidRecurrentShape, dim0, dim1, dim2, size, c.recurrentStateSize)
-	}
-
-	buf := c.recurrentBuffer(layer)
-	cur := c.currentSlotRows(ctx, buf, c.recurrentStateSize)
-	return cur.Reshape(ctx, dim0, dim1, dim2, c.NumSeqs()), nil
-}
-
-// UpdateRecurrentState writes new recurrent state for current batch sequences.
-func (c *Recurrent) UpdateRecurrentState(ctx ml.Context, layer int, newState ml.Tensor) {
-	buf := c.recurrentBuffer(layer)
-	src := newState.Reshape(ctx, c.recurrentStateSize, c.NumSeqs())
-	srcF32 := src
-	if src.DType() != ml.DTypeF32 {
-		srcF32 = src.Cast(ctx, ml.DTypeF32)
-	}
-	c.writeCurrentSlotRows(ctx, buf, c.recurrentStateSize, srcF32)
-
-	c.captureRecurrentCheckpoint(ctx, layer, srcF32)
-}
-
-// IsSupportedForBatch returns true if the current batch layout supports recurrent layers.
-func (c *Recurrent) IsSupportedForBatch() bool {
-	return c.curSeqTokens > 0 && len(c.curSeqs) > 0
-}
-
-// Seqs returns the ordered unique sequences for the current forward pass.
-func (c *Recurrent) Seqs() []int {
-	return slices.Clone(c.curSeqs)
-}

kvcache/recurrent_checkpoints.go

@@ -1,561 +0,0 @@
-package kvcache
-
-import (
-	"log/slog"
-	"math"
-
-	"github.com/ollama/ollama/ml"
-	"github.com/ollama/ollama/model/input"
-)
-
-// TODO(jmorganca): Add byte-serialized host-RAM checkpoints to reduce GPU
-// memory usage while preserving prefix reuse for recurrent state.
-
-type checkpointEntry struct {
-	pos       int32
-	conv      map[int]ml.Tensor
-	recurrent map[int]ml.Tensor
-}
-
-type slotCheckpointStore struct {
-	entries []checkpointEntry
-	size    int
-	next    int
-	lastPos int32
-}
-
-type checkpointRestore struct {
-	slot int
-	idx  int
-	pos  int32
-}
-
-func newSlotCheckpointStore(n int) *slotCheckpointStore {
-	entries := make([]checkpointEntry, n)
-	for i := range entries {
-		entries[i].pos = -1
-	}
-	return &slotCheckpointStore{
-		entries: entries,
-		lastPos: -1,
-	}
-}
-
-func (s *slotCheckpointStore) reset() {
-	s.size = 0
-	s.next = 0
-	s.lastPos = -1
-	for i := range s.entries {
-		s.entries[i].pos = -1
-	}
-}
-
-func (s *slotCheckpointStore) record(pos int32) int {
-	if len(s.entries) == 0 {
-		return -1
-	}
-	idx := s.next
-	s.next = (s.next + 1) % len(s.entries)
-	if s.size < len(s.entries) {
-		s.size++
-	}
-	s.entries[idx].pos = pos
-	s.lastPos = pos
-	return idx
-}
-
-func (s *slotCheckpointStore) bestIndex(targetPos int32) (int, int32, bool) {
-	bestIdx := -1
-	bestPos := int32(-1)
-	for i := range s.entries {
-		pos := s.entries[i].pos
-		if pos < 0 || pos >= targetPos {
-			continue
-		}
-		if pos > bestPos {
-			bestPos = pos
-			bestIdx = i
-		}
-	}
-	if bestIdx < 0 {
-		return -1, -1, false
-	}
-	return bestIdx, bestPos, true
-}
-
-func (s *slotCheckpointStore) pruneAfter(pos int32) {
-	if len(s.entries) == 0 {
-		s.size = 0
-		s.next = 0
-		s.lastPos = -1
-		return
-	}
-
-	size := 0
-	next := -1
-	minPos := int32(math.MaxInt32)
-	minIdx := 0
-	for i := range s.entries {
-		if s.entries[i].pos > pos {
-			s.entries[i].pos = -1
-		}
-		if s.entries[i].pos >= 0 {
-			size++
-			if s.entries[i].pos < minPos {
-				minPos = s.entries[i].pos
-				minIdx = i
-			}
-		} else if next == -1 {
-			next = i
-		}
-	}
-
-	s.size = size
-	if size == 0 {
-		s.next = 0
-		s.lastPos = -1
-		return
-	}
-	if next != -1 {
-		s.next = next
-	} else {
-		// Full ring: overwrite the oldest checkpoint next.
-		s.next = minIdx
-	}
-	s.lastPos = pos
-}
-
-func (s *slotCheckpointStore) shiftRange(beginIndex, endIndex int32) {
-	if len(s.entries) == 0 {
-		s.size = 0
-		s.next = 0
-		s.lastPos = -1
-		return
-	}
-
-	offset := beginIndex - endIndex
-
-	size := 0
-	next := -1
-	minPos := int32(math.MaxInt32)
-	maxPos := int32(-1)
-	minIdx := 0
-
-	for i := range s.entries {
-		pos := s.entries[i].pos
-		if pos >= 0 {
-			if pos >= beginIndex && pos < endIndex {
-				s.entries[i].pos = -1
-			} else if pos >= endIndex {
-				s.entries[i].pos = pos + offset
-			}
-		}
-
-		pos = s.entries[i].pos
-		if pos >= 0 {
-			size++
-			if pos < minPos {
-				minPos = pos
-				minIdx = i
-			}
-			if pos > maxPos {
-				maxPos = pos
-			}
-		} else if next == -1 {
-			next = i
-		}
-	}
-
-	s.size = size
-	if size == 0 {
-		s.next = 0
-		s.lastPos = -1
-		return
-	}
-
-	if next != -1 {
-		s.next = next
-	} else {
-		// Full ring: overwrite the oldest checkpoint next.
-		s.next = minIdx
-	}
-	s.lastPos = maxPos
-}
-
-func (s *slotCheckpointStore) window() (size int, minPos, maxPos, lastPos int32) {
-	minPos = int32(math.MaxInt32)
-	maxPos = int32(-1)
-	for i := range s.entries {
-		pos := s.entries[i].pos
-		if pos < 0 {
-			continue
-		}
-		size++
-		if pos < minPos {
-			minPos = pos
-		}
-		if pos > maxPos {
-			maxPos = pos
-		}
-	}
-	if size == 0 {
-		minPos = -1
-		maxPos = -1
-	}
-	return size, minPos, maxPos, s.lastPos
-}
-
-func (c *Recurrent) checkpointTag() string {
-	if c.logPrefix == "" {
-		return "kvcache.recurrent"
-	}
-	return c.logPrefix
-}
-
-func (c *Recurrent) planCheckpoints(batch input.Batch) {
-	if c.checkpointCount == 0 || len(c.curSeqs) == 0 {
-		c.curCheckpointPos = c.curCheckpointPos[:0]
-		for k := range c.curCheckpointSlots {
-			delete(c.curCheckpointSlots, k)
-		}
-		return
-	}
-
-	if cap(c.curCheckpointPos) < len(c.curSeqs) {
-		c.curCheckpointPos = make([]int32, len(c.curSeqs))
-	} else {
-		c.curCheckpointPos = c.curCheckpointPos[:len(c.curSeqs)]
-	}
-	for i := range c.curCheckpointPos {
-		c.curCheckpointPos[i] = -1
-	}
-	for k := range c.curCheckpointSlots {
-		delete(c.curCheckpointSlots, k)
-	}
-
-	posMax := make(map[int]int32, len(c.curSeqs))
-	for i, seq := range batch.Sequences {
-		pos := batch.Positions[i]
-		if cur, ok := posMax[seq]; !ok || pos > cur {
-			posMax[seq] = pos
-		}
-	}
-
-	for i, seq := range c.curSeqs {
-		pos, ok := posMax[seq]
-		if !ok {
-			continue
-		}
-		if pos < c.checkpointMinPos {
-			continue
-		}
-		slot := c.curSlots[i]
-		store := c.checkpointStore(slot)
-		lastPos := store.lastPos
-		if lastPos < 0 || pos-lastPos >= c.checkpointInterval {
-			c.curCheckpointPos[i] = pos
-		}
-	}
-}
-
-func (c *Recurrent) checkpointStore(slot int) *slotCheckpointStore {
-	store, ok := c.checkpoints[slot]
-	if ok {
-		return store
-	}
-	store = newSlotCheckpointStore(c.checkpointCount)
-	c.checkpoints[slot] = store
-	return store
-}
-
-func (c *Recurrent) checkpointIndexForSlot(slot int, pos int32) int {
-	if c.checkpointCount == 0 {
-		return -1
-	}
-	if idx, ok := c.curCheckpointSlots[slot]; ok {
-		return idx
-	}
-	store := c.checkpointStore(slot)
-	idx := store.record(pos)
-	if idx >= 0 {
-		c.curCheckpointSlots[slot] = idx
-	}
-	return idx
-}
-
-func (c *Recurrent) hasCheckpoint(seq int, pos int32) bool {
-	if pos <= 0 {
-		return false
-	}
-	slot, ok := c.slotForSeq[seq]
-	if !ok {
-		return false
-	}
-	store, ok := c.checkpoints[slot]
-	if !ok {
-		return false
-	}
-	_, _, ok = store.bestIndex(pos)
-	return ok
-}
-
-func (c *Recurrent) PrepareRestore(seq int, targetPos int32) (int32, bool) {
-	if targetPos <= 0 {
-		return 0, false
-	}
-	slot, ok := c.slotForSeq[seq]
-	if !ok {
-		return 0, false
-	}
-	store, ok := c.checkpoints[slot]
-	if !ok {
-		slog.Debug(c.checkpointTag()+": checkpoint miss", "seq", seq, "slot", slot, "target", targetPos, "size", 0)
-		return 0, false
-	}
-	idx, pos, ok := store.bestIndex(targetPos)
-	if !ok {
-		size, minPos, maxPos, lastPos := store.window()
-		slog.Debug(c.checkpointTag()+": checkpoint miss", "seq", seq, "slot", slot, "target", targetPos, "size", size,
-			"min", minPos, "max", maxPos, "last", lastPos)
-		return 0, false
-	}
-	c.pendingRestore[seq] = checkpointRestore{
-		slot: slot,
-		idx:  idx,
-		pos:  pos,
-	}
-	return pos + 1, true
-}
-
-func (c *Recurrent) applyCheckpointRestore(restore checkpointRestore) error {
-	entry, ok := c.restoreEntry(restore)
-	if !ok {
-		return ErrNotSupported
-	}
-
-	ctx := c.backend.NewContext()
-	defer ctx.Close()
-
-	slotIdx := ctx.Input().FromInts([]int32{int32(restore.slot)}, 1)
-	for layer, src := range entry.conv {
-		buf := c.convBuffer(layer)
-		ctx.Forward(buf.SetRows(ctx, src, slotIdx))
-	}
-	for layer, src := range entry.recurrent {
-		buf := c.recurrentBuffer(layer)
-		ctx.Forward(buf.SetRows(ctx, src, slotIdx))
-	}
-
-	if len(entry.conv) > 0 || len(entry.recurrent) > 0 {
-		ctx.Compute()
-	}
-	store := c.checkpoints[restore.slot]
-	store.pruneAfter(restore.pos)
-	return nil
-}
-
-func (c *Recurrent) restoreComplete(restore checkpointRestore) bool {
-	_, ok := c.restoreEntry(restore)
-	return ok
-}
-
-func (c *Recurrent) restoreEntry(restore checkpointRestore) (*checkpointEntry, bool) {
-	store, ok := c.checkpoints[restore.slot]
-	if !ok || restore.idx < 0 || restore.idx >= len(store.entries) {
-		return nil, false
-	}
-	entry := &store.entries[restore.idx]
-	if entry.pos < 0 {
-		return nil, false
-	}
-	if !c.entryComplete(entry) {
-		return nil, false
-	}
-	return entry, true
-}
-
-func (c *Recurrent) entryComplete(entry *checkpointEntry) bool {
-	for layer := range c.convStates {
-		if entry.conv == nil || entry.conv[layer] == nil {
-			return false
-		}
-	}
-	for layer := range c.recurrentStates {
-		if entry.recurrent == nil || entry.recurrent[layer] == nil {
-			return false
-		}
-	}
-	return true
-}
-
-func (c *Recurrent) clearCheckpoints(slot int) {
-	if store, ok := c.checkpoints[slot]; ok {
-		store.reset()
-	}
-}
-
-func (c *Recurrent) shiftCheckpoints(slot int, beginIndex, endIndex int32) {
-	if store, ok := c.checkpoints[slot]; ok {
-		store.shiftRange(beginIndex, endIndex)
-	}
-}
-
-func (c *Recurrent) copyCheckpoints(ctx ml.Context, srcSlot, dstSlot int) {
-	if c.checkpointCount == 0 {
-		return
-	}
-	srcStore, ok := c.checkpoints[srcSlot]
-	if !ok || srcStore.size == 0 {
-		return
-	}
-	dstStore := c.checkpointStore(dstSlot)
-	dstStore.size = srcStore.size
-	dstStore.next = srcStore.next
-	dstStore.lastPos = srcStore.lastPos
-
-	for i := range srcStore.entries {
-		srcEntry := &srcStore.entries[i]
-		dstEntry := &dstStore.entries[i]
-		dstEntry.pos = srcEntry.pos
-		if srcEntry.conv != nil {
-			if dstEntry.conv == nil {
-				dstEntry.conv = make(map[int]ml.Tensor)
-			}
-			for layer, src := range srcEntry.conv {
-				dst := c.ensureCheckpointConv(layer, dstEntry)
-				ctx.Forward(src.Copy(ctx, dst))
-			}
-		}
-		if srcEntry.recurrent != nil {
-			if dstEntry.recurrent == nil {
-				dstEntry.recurrent = make(map[int]ml.Tensor)
-			}
-			for layer, src := range srcEntry.recurrent {
-				dst := c.ensureCheckpointRecurrent(layer, dstEntry)
-				ctx.Forward(src.Copy(ctx, dst))
-			}
-		}
-	}
-}
-
-func (c *Recurrent) captureConvCheckpoint(ctx ml.Context, layer int, src ml.Tensor) {
-	if c.checkpointCount == 0 {
-		return
-	}
-	if c.reserveCheckpoints {
-		c.reserveCheckpointConv(layer)
-		return
-	}
-	if len(c.curCheckpointPos) == 0 {
-		return
-	}
-	for i, pos := range c.curCheckpointPos {
-		if pos < 0 {
-			continue
-		}
-		slot := c.curSlots[i]
-		idx := c.checkpointIndexForSlot(slot, pos)
-		if idx < 0 {
-			continue
-		}
-		entry := &c.checkpoints[slot].entries[idx]
-		dst := c.ensureCheckpointConv(layer, entry)
-		seqSlice := src.Slice(ctx, 1, i, i+1, 1)
-		ctx.Forward(seqSlice.Copy(ctx, dst))
-	}
-}
-
-func (c *Recurrent) captureRecurrentCheckpoint(ctx ml.Context, layer int, src ml.Tensor) {
-	if c.checkpointCount == 0 {
-		return
-	}
-	if c.reserveCheckpoints {
-		c.reserveCheckpointRecurrent(layer)
-		return
-	}
-	if len(c.curCheckpointPos) == 0 {
-		return
-	}
-	for i, pos := range c.curCheckpointPos {
-		if pos < 0 {
-			continue
-		}
-		slot := c.curSlots[i]
-		idx := c.checkpointIndexForSlot(slot, pos)
-		if idx < 0 {
-			continue
-		}
-		entry := &c.checkpoints[slot].entries[idx]
-		dst := c.ensureCheckpointRecurrent(layer, entry)
-		seqSlice := src.Slice(ctx, 1, i, i+1, 1)
-		ctx.Forward(seqSlice.Copy(ctx, dst))
-	}
-}
-
-func (c *Recurrent) ensureCheckpointConv(layer int, entry *checkpointEntry) ml.Tensor {
-	if entry.conv == nil {
-		entry.conv = make(map[int]ml.Tensor)
-	}
-	if t, ok := entry.conv[layer]; ok {
-		return t
-	}
-	ctx, ok := c.checkpointConvCtxs[layer]
-	if !ok {
-		ctx = c.backend.NewContextSize(c.checkpointCtxSize).Layer(layer)
-		c.checkpointConvCtxs[layer] = ctx
-	}
-	t := ctx.Zeros(ml.DTypeF32, c.convDim*c.convChannels, 1)
-	entry.conv[layer] = t
-	return t
-}
-
-func (c *Recurrent) ensureCheckpointRecurrent(layer int, entry *checkpointEntry) ml.Tensor {
-	if entry.recurrent == nil {
-		entry.recurrent = make(map[int]ml.Tensor)
-	}
-	if t, ok := entry.recurrent[layer]; ok {
-		return t
-	}
-	ctx, ok := c.checkpointRecurCtxs[layer]
-	if !ok {
-		ctx = c.backend.NewContextSize(c.checkpointCtxSize).Layer(layer)
-		c.checkpointRecurCtxs[layer] = ctx
-	}
-	t := ctx.Zeros(ml.DTypeF32, c.recurrentStateSize, 1)
-	entry.recurrent[layer] = t
-	return t
-}
-
-func (c *Recurrent) reserveCheckpointConv(layer int) {
-	key := checkpointReserveKey(layer, 0)
-	if _, ok := c.checkpointReserved[key]; ok {
-		return
-	}
-	for slot := range c.maxSequences {
-		store := c.checkpointStore(slot)
-		for i := range store.entries {
-			entry := &store.entries[i]
-			_ = c.ensureCheckpointConv(layer, entry)
-		}
-	}
-	c.checkpointReserved[key] = struct{}{}
-}
-
-func (c *Recurrent) reserveCheckpointRecurrent(layer int) {
-	key := checkpointReserveKey(layer, 1)
-	if _, ok := c.checkpointReserved[key]; ok {
-		return
-	}
-	for slot := range c.maxSequences {
-		store := c.checkpointStore(slot)
-		for i := range store.entries {
-			entry := &store.entries[i]
-			_ = c.ensureCheckpointRecurrent(layer, entry)
-		}
-	}
-	c.checkpointReserved[key] = struct{}{}
-}
-
-func checkpointReserveKey(layer int, kind int) int {
-	return layer*2 + kind
-}

kvcache/recurrent_checkpoints_test.go

@@ -1,288 +0,0 @@
-package kvcache
-
-import (
-	"errors"
-	"math"
-	"slices"
-	"testing"
-
-	"github.com/ollama/ollama/ml"
-)
-
-func newTestCache() *Recurrent {
-	return NewRecurrentCache(RecurrentConfig{ConvDim: 1, ConvChannels: 2, RecurrentStateSize: 2})
-}
-
-func TestSlotCheckpointStoreBestIndex(t *testing.T) {
-	store := newSlotCheckpointStore(2)
-	store.record(10)
-	store.record(20)
-
-	_, pos, ok := store.bestIndex(15)
-	if !ok || pos != 10 {
-		t.Fatalf("expected best pos 10, got pos=%d ok=%v", pos, ok)
-	}
-
-	store.record(30) // overwrite oldest (10)
-
-	if _, _, ok := store.bestIndex(15); ok {
-		t.Fatalf("expected no checkpoint for targetPos=15 after overwrite")
-	}
-
-	_, pos, ok = store.bestIndex(40)
-	if !ok || pos != 30 {
-		t.Fatalf("expected best pos 30, got pos=%d ok=%v", pos, ok)
-	}
-}
-
-func TestCachePrepareRestore(t *testing.T) {
-	cache := newTestCache()
-	cache.checkpointCount = 3
-	cache.checkpoints = make(map[int]*slotCheckpointStore)
-	cache.pendingRestore = make(map[int]checkpointRestore)
-
-	cache.slotForSeq[1] = 0
-	store := cache.checkpointStore(0)
-	store.record(5)
-	store.record(9)
-	store.record(15)
-
-	restorePos, ok := cache.PrepareRestore(1, 12)
-	if !ok {
-		t.Fatalf("expected restore ok")
-	}
-	if restorePos != 10 {
-		t.Fatalf("expected restorePos 10, got %d", restorePos)
-	}
-	rest, ok := cache.pendingRestore[1]
-	if !ok {
-		t.Fatalf("expected pending restore entry")
-	}
-	if rest.pos != 9 {
-		t.Fatalf("expected pending restore pos 9, got %d", rest.pos)
-	}
-}
-
-func TestSlotCheckpointStorePruneAfter(t *testing.T) {
-	store := newSlotCheckpointStore(3)
-	store.record(10)
-	store.record(20)
-	store.record(30)
-
-	store.pruneAfter(20)
-
-	if store.lastPos != 20 {
-		t.Fatalf("expected lastPos 20, got %d", store.lastPos)
-	}
-
-	_, pos, ok := store.bestIndex(25)
-	if !ok || pos != 20 {
-		t.Fatalf("expected best pos 20 after prune, got pos=%d ok=%v", pos, ok)
-	}
-
-	_, pos, ok = store.bestIndex(35)
-	if !ok || pos != 20 {
-		t.Fatalf("expected pruned best pos 20 for targetPos=35, got pos=%d ok=%v", pos, ok)
-	}
-}
-
-func TestCacheRestoreRejectsIncompleteCheckpoint(t *testing.T) {
-	cache := newTestCache()
-	cache.checkpointCount = 3
-	cache.checkpoints = make(map[int]*slotCheckpointStore)
-	cache.pendingRestore = make(map[int]checkpointRestore)
-
-	cache.slotForSeq[1] = 0
-	cache.refCount = []int{1}
-	cache.freeSlots = nil
-
-	// Simulate layer 0 requires both conv and recurrent checkpoints.
-	cache.convStates[0] = nil
-	cache.recurrentStates[0] = nil
-
-	store := cache.checkpointStore(0)
-	idx := store.record(9)
-	entry := &store.entries[idx]
-	entry.conv = map[int]ml.Tensor{0: nil}
-	// entry.recurrent intentionally missing
-
-	cache.pendingRestore[1] = checkpointRestore{slot: 0, idx: idx, pos: 9}
-
-	err := cache.Remove(1, 10, math.MaxInt32)
-	if !errors.Is(err, ErrNotSupported) {
-		t.Fatalf("expected ErrNotSupported for incomplete checkpoint, got %v", err)
-	}
-}
-
-func TestCacheRestoreAcceptsCompleteCheckpoint(t *testing.T) {
-	cache := newTestCache()
-	cache.checkpointCount = 3
-	cache.checkpoints = make(map[int]*slotCheckpointStore)
-	cache.pendingRestore = make(map[int]checkpointRestore)
-
-	cache.slotForSeq[1] = 0
-	cache.refCount = []int{1}
-	cache.freeSlots = nil
-
-	store := cache.checkpointStore(0)
-	idx := store.record(9)
-
-	cache.pendingRestore[1] = checkpointRestore{slot: 0, idx: idx, pos: 9}
-
-	restore := cache.pendingRestore[1]
-	if !cache.restoreComplete(restore) {
-		t.Fatalf("expected restoreComplete to return true for complete checkpoint")
-	}
-}
-
-func TestCacheRecurrentStateShapeValidation(t *testing.T) {
-	cache := newTestCache()
-	_, err := cache.RecurrentState(nil, 0, 3)
-	if !errors.Is(err, ErrInvalidRecurrentShape) {
-		t.Fatalf("expected ErrInvalidRecurrentShape, got %v", err)
-	}
-}
-
-func TestSlotCheckpointStoreShiftRange(t *testing.T) {
-	store := newSlotCheckpointStore(5)
-	store.record(1)
-	store.record(4)
-	store.record(7)
-	store.record(10)
-
-	store.shiftRange(2, 6)
-
-	var positions []int32
-	for i := range store.entries {
-		if store.entries[i].pos >= 0 {
-			positions = append(positions, store.entries[i].pos)
-		}
-	}
-	slices.Sort(positions)
-
-	want := []int32{1, 3, 6}
-	if !slices.Equal(positions, want) {
-		t.Fatalf("unexpected shifted positions: got=%v want=%v", positions, want)
-	}
-	if store.lastPos != 6 {
-		t.Fatalf("expected lastPos 6, got %d", store.lastPos)
-	}
-}
-
-func TestCacheRemoveMiddleShiftsCheckpoints(t *testing.T) {
-	cache := newTestCache()
-	cache.slotForSeq[1] = 0
-	cache.refCount = []int{1}
-	cache.pendingRestore[1] = checkpointRestore{slot: 0, idx: 0, pos: 1}
-
-	store := cache.checkpointStore(0)
-	store.record(1)
-	store.record(4)
-	store.record(7)
-	store.record(10)
-
-	if err := cache.Remove(1, 2, 6); err != nil {
-		t.Fatalf("expected middle remove to succeed, got %v", err)
-	}
-
-	if _, ok := cache.pendingRestore[1]; ok {
-		t.Fatalf("expected pending restore to be cleared after middle remove")
-	}
-
-	var positions []int32
-	for i := range store.entries {
-		if store.entries[i].pos >= 0 {
-			positions = append(positions, store.entries[i].pos)
-		}
-	}
-	slices.Sort(positions)
-
-	want := []int32{1, 3, 6}
-	if !slices.Equal(positions, want) {
-		t.Fatalf("unexpected checkpoint positions after remove: got=%v want=%v", positions, want)
-	}
-}
-
-func TestSlotCheckpointStoreRingBufferWrapAround(t *testing.T) {
-	store := newSlotCheckpointStore(3)
-
-	store.record(10)
-	store.record(20)
-	store.record(30)
-
-	store.entries[0].conv = make(map[int]ml.Tensor)
-	store.entries[0].conv[0] = nil
-	store.entries[0].recurrent = make(map[int]ml.Tensor)
-	store.entries[0].recurrent[0] = nil
-
-	store.record(40)
-
-	if store.entries[0].conv == nil {
-		t.Fatalf("expected conv map to be preserved on reuse")
-	}
-	if store.entries[0].recurrent == nil {
-		t.Fatalf("expected recurrent map to be preserved on reuse")
-	}
-	if store.entries[0].pos != 40 {
-		t.Fatalf("expected entry 0 pos to be 40, got %d", store.entries[0].pos)
-	}
-}
-
-func TestSlotCheckpointStoreFullCapacity(t *testing.T) {
-	store := newSlotCheckpointStore(2)
-
-	idx1 := store.record(10)
-	idx2 := store.record(20)
-
-	if idx1 != 0 || idx2 != 1 {
-		t.Fatalf("expected indices 0, 1, got %d, %d", idx1, idx2)
-	}
-	if store.size != 2 {
-		t.Fatalf("expected size 2, got %d", store.size)
-	}
-
-	_, pos1, ok1 := store.bestIndex(15)
-	_, pos2, ok2 := store.bestIndex(25)
-
-	if !ok1 || pos1 != 10 {
-		t.Fatalf("expected best pos 10 for target 15, got pos=%d ok=%v", pos1, ok1)
-	}
-	if !ok2 || pos2 != 20 {
-		t.Fatalf("expected best pos 20 for target 25, got pos=%d ok=%v", pos2, ok2)
-	}
-}
-
-func TestSlotCheckpointStoreEmptyBuffer(t *testing.T) {
-	store := newSlotCheckpointStore(0)
-
-	idx := store.record(10)
-	if idx != -1 {
-		t.Fatalf("expected record to return -1 for empty buffer, got %d", idx)
-	}
-
-	_, _, ok := store.bestIndex(15)
-	if ok {
-		t.Fatalf("expected no checkpoint for empty buffer")
-	}
-}
-
-func TestSlotCheckpointStorePruneAfterAll(t *testing.T) {
-	store := newSlotCheckpointStore(3)
-	store.record(10)
-	store.record(20)
-	store.record(30)
-
-	store.pruneAfter(5)
-
-	if store.size != 0 {
-		t.Fatalf("expected size 0 after pruning all, got %d", store.size)
-	}
-	if store.lastPos != -1 {
-		t.Fatalf("expected lastPos -1 after pruning all, got %d", store.lastPos)
-	}
-
-	_, _, ok := store.bestIndex(100)
-	if ok {
-		t.Fatalf("expected no checkpoint after pruning all")
-	}
-}

kvcache/wrapper.go

@@ -1,110 +0,0 @@
-package kvcache
-
-import (
-	"math"
-
-	"github.com/ollama/ollama/ml"
-	"github.com/ollama/ollama/model/input"
-)
-
-// Wrapper cache is a container for multiple types of caches,
-// such as for the encoding and decoding portions of a model.
-type WrapperCache struct {
-	// caches we are wrapping
-	caches []Cache
-
-	// cache to be used for this layer
-	curType int
-}
-
-func NewWrapperCache(caches ...Cache) *WrapperCache {
-	return &WrapperCache{
-		caches: caches,
-	}
-}
-
-func (c *WrapperCache) Init(backend ml.Backend, dtype ml.DType, maxSequences, capacity, maxBatch int) {
-	for _, cache := range c.caches {
-		cache.Init(backend, dtype, maxSequences, capacity, maxBatch)
-	}
-}
-
-func (c *WrapperCache) SetConfig(config ml.CacheConfig) {
-	for _, cache := range c.caches {
-		cache.SetConfig(config)
-	}
-}
-
-func (c *WrapperCache) Close() {
-	for _, cache := range c.caches {
-		cache.Close()
-	}
-}
-
-func (c *WrapperCache) StartForward(ctx ml.Context, batch input.Batch, reserve bool) error {
-	for i, cache := range c.caches {
-		err := cache.StartForward(ctx, batch, reserve)
-		if err != nil {
-			// unwind on error - Remove with endIndex set to math.MaxInt32 does not fail
-			for j := i - 1; j >= 0; j-- {
-				for k := range batch.Positions {
-					_ = c.caches[j].Remove(batch.Sequences[k], batch.Positions[k], math.MaxInt32)
-				}
-			}
-			return err
-		}
-	}
-
-	c.curType = 0
-	return nil
-}
-
-func (c *WrapperCache) SetLayer(layer int) {
-	for _, cache := range c.caches {
-		cache.SetLayer(layer)
-	}
-}
-
-func (c *WrapperCache) SetLayerType(layerType int) {
-	c.curType = layerType
-}
-
-func (c *WrapperCache) UnderlyingCache() Cache {
-	return c.caches[c.curType]
-}
-
-func (c *WrapperCache) Get(ctx ml.Context) (ml.Tensor, ml.Tensor, ml.Tensor) {
-	return c.caches[c.curType].Get(ctx)
-}
-
-func (c *WrapperCache) Put(ctx ml.Context, key, value ml.Tensor) {
-	c.caches[c.curType].Put(ctx, key, value)
-}
-
-func (c *WrapperCache) CopyPrefix(srcSeq, dstSeq int, len int32) {
-	for _, cache := range c.caches {
-		cache.CopyPrefix(srcSeq, dstSeq, len)
-	}
-}
-
-func (c *WrapperCache) CanResume(seq int, pos int32) bool {
-	for _, cache := range c.caches {
-		if !cache.CanResume(seq, pos) {
-			return false
-		}
-	}
-
-	return true
-}
-
-func (c *WrapperCache) Remove(seq int, beginIndex, endIndex int32) error {
-	// If the one of these fails, the caller is supposed to retry with endIndex set to math.MaxInt32, which should not fail
-	for _, cache := range c.caches {
-		err := cache.Remove(seq, beginIndex, endIndex)
-		if err != nil {
-			return err
-		}
-	}
-
-	return nil
-}

llama/compat/llama-ollama-compat-util.cpp

@@ -275,7 +275,7 @@ bool take_load_op(const char * dest_name, LoadOp & out) {
 }
 
 bool read_at(const char * path, size_t offset, void * dst, size_t size) {
-    FILE * f = std::fopen(path, "rb");
+    FILE * f = ggml_fopen(path, "rb");
     if (!f) {
         std::fprintf(stderr, "%s: open failed path=%s offset=%zu size=%zu errno=%d (%s)\n",
                      __func__, path, offset, size, errno, std::strerror(errno));

llama/compat/models/003-llama-cpp-laguna.patch

@@ -5,7 +5,7 @@ diff --git a/src/llama-arch.cpp b/src/llama-arch.cpp
      { LLM_ARCH_MELLUM,           "mellum"           },
 +    { LLM_ARCH_LAGUNA,           "laguna"           },
      { LLM_ARCH_UNKNOWN,          "(unknown)"        },
-@@ -380,2 +381,3 @@ static const std::map<llm_tensor, const char *> LLM_TENSOR_NAMES = {
+@@ -398,2 +399,3 @@ static const std::map<llm_tensor, const char *> LLM_TENSOR_NAMES = {
      { LLM_TENSOR_ATTN_GATE,                              "blk.%d.attn_gate" },
 +    { LLM_TENSOR_ATTN_GATE_LAGUNA,                       "blk.%d.attn_g" },
      { LLM_TENSOR_FFN_POST_NORM,                          "blk.%d.post_ffw_norm" },
@@ -17,7 +17,7 @@ diff --git a/src/llama-arch.h b/src/llama-arch.h
 --- a/src/llama-arch.h
 +++ b/src/llama-arch.h
 @@ -145,2 +145,3 @@ enum llm_arch {
-     LLM_ARCH_EAGLE3,
+     LLM_ARCH_DFLASH,
 +    LLM_ARCH_LAGUNA,
      LLM_ARCH_UNKNOWN,
 @@ -382,2 +383,3 @@ enum llm_tensor {
@@ -34,10 +34,10 @@ diff --git a/src/llama-model.cpp b/src/llama-model.cpp
 +            return new llama_model_laguna(params);
          case LLM_ARCH_DECI:
              return new llama_model_deci(params);
-@@ -2409,2 +2411,3 @@ llama_rope_type llama_model_rope_type(const llama_model * model) {
-         case LLM_ARCH_TALKIE:
-+        case LLM_ARCH_LAGUNA:
+@@ -2525,2 +2527,3 @@ llama_rope_type llama_model_rope_type(const llama_model * model) {
          case LLM_ARCH_MELLUM:
++        case LLM_ARCH_LAGUNA:
+         case LLM_ARCH_DFLASH:
 diff --git a/src/llama-model-loader.cpp b/src/llama-model-loader.cpp
 --- a/src/llama-model-loader.cpp
 +++ b/src/llama-model-loader.cpp
@@ -76,7 +76,7 @@ diff --git a/src/llama-vocab.h b/src/llama-vocab.h
 diff --git a/src/models/models.h b/src/models/models.h
 --- a/src/models/models.h
 +++ b/src/models/models.h
-@@ -1481,1 +1481,14 @@ struct llama_model_dots1 : public llama_model_base {
+@@ -1657,1 +1657,14 @@ struct llama_model_arcee : public llama_model_base {
 +struct llama_model_laguna : public llama_model_base {
 +    llama_model_laguna(const struct llama_model_params & params) : llama_model_base(params) {}
 +    void load_arch_hparams(llama_model_loader & ml) override;

llama/server/CMakePresets.json

@@ -68,7 +68,7 @@
       "inherits": ["llama_cuda_v12_base"],
       "binaryDir": "${sourceDir}/../../build/llama-server-cuda_v12",
       "cacheVariables": {
-        "CMAKE_CUDA_ARCHITECTURES": "50-virtual;52-virtual;60-virtual;61-virtual;70;75;80;86;89;90;90a;120"
+        "CMAKE_CUDA_ARCHITECTURES": "50-virtual;52-virtual;60;61;70;75;80;86;89;90;90a;120"
       }
     },
     {
@@ -178,7 +178,7 @@
       "inherits": ["rocm_v7_1_base"],
       "binaryDir": "${sourceDir}/../../build/llama-server-rocm_v7_1",
       "cacheVariables": {
-        "AMDGPU_TARGETS": "gfx942;gfx950;gfx1010;gfx1012;gfx1030;gfx1100;gfx1101;gfx1102;gfx1103;gfx1150;gfx1151;gfx1200;gfx1201"
+        "AMDGPU_TARGETS": "gfx1030;gfx1100;gfx1101;gfx1102;gfx1150;gfx1151;gfx1200;gfx1201"
       }
     },
     {
@@ -202,7 +202,7 @@
       "inherits": ["rocm_v7_2_base"],
       "binaryDir": "${sourceDir}/../../build/llama-server-rocm_v7_2",
       "cacheVariables": {
-        "AMDGPU_TARGETS": "gfx908:xnack-;gfx90a:xnack+;gfx90a:xnack-;gfx942;gfx950;gfx1010;gfx1012;gfx1030;gfx1100;gfx1101;gfx1102;gfx1103;gfx1150;gfx1151;gfx1200;gfx1201"
+        "AMDGPU_TARGETS": "gfx908:xnack-;gfx90a:xnack+;gfx90a:xnack-;gfx942;gfx950;gfx1030;gfx1100;gfx1101;gfx1102;gfx1150;gfx1151;gfx1200;gfx1201"
       }
     },
     {

llm/llama_server.go

@@ -77,6 +77,14 @@ const (
 	openEndedGenerationContextMultiplier = 10
 )
 
+const (
+	llamaArgFitTargetEnv = "LLAMA_ARG_FIT_TARGET"
+	bytesPerMiB          = 1 << 20
+
+	// mmprojOffloadHeadroom leaves 1 GiB for backend buffers beyond projector weights.
+	mmprojOffloadHeadroom = 1 << 30
+)
+
 // DefaultEmbeddingNumBatchForContext caps the embedding batch default to the
 // active context length before it is passed to llama-server.
 func DefaultEmbeddingNumBatchForContext(numCtx int) int {
@@ -420,7 +428,7 @@ func startLlamaServer(launch llamaServerLaunchConfig, out io.Writer) (cmd *exec.
 		cmd.Stderr = out
 	}
 	cmd.SysProcAttr = LlamaServerSysProcAttr
-	SetupLlamaServerCommandEnv(cmd, exe, launch.gpuLibs, launch.extraEnvs)
+	SetupLlamaServerCommandEnv(cmd, exe, launch.gpuLibs, launch.extraEnvsForStart())
 
 	slog.Info("starting llama-server", "cmd", cmd)
 	slog.Debug("subprocess", "", filteredEnv(cmd.Env))
@@ -621,11 +629,6 @@ func appendMainGPUArgs(params []string, opts api.Options) []string {
 	return append(params, "--split-mode", "none", "--main-gpu", strconv.Itoa(*opts.MainGPU))
 }
 
-const (
-	// mmprojOffloadHeadroom leaves 1 GiB for backend buffers beyond projector weights.
-	mmprojOffloadHeadroom = 1 << 30
-)
-
 func appendMMProjArgs(params []string, launch llamaServerLaunchConfig) []string {
 	if len(launch.projectors) == 0 {
 		return params
@@ -658,9 +661,6 @@ func shouldDisableMMProjOffload(opts api.Options, gpus []ml.DeviceInfo, modelLay
 	requiredMemory := mmprojMemory + mmprojOffloadHeadroom
 
 	for _, gpu := range gpus {
-		if gpu.Integrated && gpu.Library != "Metal" {
-			return true, "shared-memory-gpu"
-		}
 		memory := gpu.FreeMemory
 		if memory == 0 || (gpu.TotalMemory > 0 && gpu.TotalMemory < memory) {
 			memory = gpu.TotalMemory
@@ -673,6 +673,47 @@ func shouldDisableMMProjOffload(opts api.Options, gpus []ml.DeviceInfo, modelLay
 	return false, ""
 }
 
+func (launch llamaServerLaunchConfig) extraEnvsForStart() map[string]string {
+	pad, ok := launch.mmprojFitTargetMiB()
+	if !ok {
+		return launch.extraEnvs
+	}
+
+	if existing, ok := launch.extraEnvs[llamaArgFitTargetEnv]; ok {
+		existingTarget, err := strconv.ParseUint(existing, 10, 64)
+		if err != nil {
+			slog.Warn("invalid llama-server fit target", "env", llamaArgFitTargetEnv, "value", existing, "error", err)
+			return launch.extraEnvs
+		}
+
+		envs := cloneStringMap(launch.extraEnvs)
+		envs[llamaArgFitTargetEnv] = strconv.FormatUint(existingTarget+pad, 10)
+		return envs
+	}
+
+	if _, ok := os.LookupEnv(llamaArgFitTargetEnv); ok {
+		// Preserve an inherited user override. SetupLlamaServerCommandEnv
+		// will pass it through unless extraEnvs overrides it.
+		return launch.extraEnvs
+	}
+
+	envs := cloneStringMap(launch.extraEnvs)
+	envs[llamaArgFitTargetEnv] = strconv.FormatUint(pad, 10)
+	return envs
+}
+
+func (launch llamaServerLaunchConfig) mmprojFitTargetMiB() (uint64, bool) {
+	if len(launch.projectors) == 0 || launch.mmprojMemory == 0 {
+		return 0, false
+	}
+	if disable, _ := launch.mmprojOffloadDisabled(); disable {
+		return 0, false
+	}
+
+	requiredMemory := launch.mmprojMemory + mmprojOffloadHeadroom
+	return (requiredMemory + bytesPerMiB - 1) / bytesPerMiB, true
+}
+
 // mmprojMemoryRequirement is a stopgap until fit accounts for mmproj memory directly.
 func mmprojMemoryRequirement(modelPath string, f *ggml.GGML, projectors []string) (uint64, error) {
 	if len(projectors) == 0 {

llm/llama_server_test.go

@@ -1955,7 +1955,7 @@ func TestAppendFlashAttentionArgs(t *testing.T) {
 	supportedGPU := []ml.DeviceInfo{{DeviceID: ml.DeviceID{Library: "CUDA"}, DriverMajor: 13, ComputeMajor: 8, ComputeMinor: 9}}
 	oldGPU := []ml.DeviceInfo{
 		{DeviceID: ml.DeviceID{Library: "CUDA"}, DriverMajor: 12, ComputeMajor: 8, ComputeMinor: 9},
-		{DeviceID: ml.DeviceID{Library: "CUDA"}, DriverMajor: 12, ComputeMajor: 6, ComputeMinor: 2},
+		{DeviceID: ml.DeviceID{Library: "CUDA"}, DriverMajor: 12, ComputeMajor: 5, ComputeMinor: 0},
 	}
 
 	tests := []struct {
@@ -2121,13 +2121,22 @@ func TestAppendMMProjArgs(t *testing.T) {
 			want:         []string{"base", "--mmproj", "model.gguf", "--no-mmproj-offload"},
 		},
 		{
-			name:         "integrated rocm gpu disables projector offload",
+			name:         "integrated rocm gpu keeps projector offload when projector fits",
 			projectors:   []string{"model.gguf"},
 			opts:         defaultOpts,
 			gpus:         []ml.DeviceInfo{{DeviceID: ml.DeviceID{Library: "ROCm"}, Integrated: true, FreeMemory: 32 << 30}},
 			mmprojMemory: 933 << 20,
 			modelLayers:  81,
-			want:         []string{"base", "--mmproj", "model.gguf", "--no-mmproj-offload"},
+			want:         []string{"base", "--mmproj", "model.gguf"},
+		},
+		{
+			name:         "integrated cuda gpu keeps projector offload",
+			projectors:   []string{"model.gguf"},
+			opts:         defaultOpts,
+			gpus:         []ml.DeviceInfo{{DeviceID: ml.DeviceID{Library: "CUDA"}, Integrated: true, FreeMemory: 32 << 30}},
+			mmprojMemory: 933 << 20,
+			modelLayers:  81,
+			want:         []string{"base", "--mmproj", "model.gguf"},
 		},
 		{
 			name:         "integrated metal gpu keeps projector offload",
@@ -2195,6 +2204,75 @@ func TestAppendMMProjArgs(t *testing.T) {
 	}
 }
 
+func TestMMProjFitTargetExtraEnvs(t *testing.T) {
+	t.Setenv(llamaArgFitTargetEnv, "")
+	_ = os.Unsetenv(llamaArgFitTargetEnv)
+
+	const (
+		projectorMemoryMiB = uint64(933)
+		projectorPadMiB    = projectorMemoryMiB + mmprojOffloadHeadroom/bytesPerMiB
+	)
+
+	fitTargetValue := func(mib uint64) string {
+		return fmt.Sprint(mib)
+	}
+
+	assertFitTarget := func(t *testing.T, got map[string]string, wantMiB uint64) {
+		t.Helper()
+		if got[llamaArgFitTargetEnv] != fitTargetValue(wantMiB) {
+			t.Fatalf("fit target = %q, want %d", got[llamaArgFitTargetEnv], wantMiB)
+		}
+	}
+
+	newLaunch := func(extraEnvs map[string]string) llamaServerLaunchConfig {
+		return llamaServerLaunchConfig{
+			projectors:   []string{"model.gguf"},
+			mmprojMemory: projectorMemoryMiB * bytesPerMiB,
+			opts:         api.DefaultOptions(),
+			gpus:         []ml.DeviceInfo{{DeviceID: ml.DeviceID{Library: "CUDA"}, Integrated: true, FreeMemory: 32 << 30}},
+			modelLayers:  81,
+			extraEnvs:    extraEnvs,
+		}
+	}
+
+	t.Run("sets projector pad when no fit target exists", func(t *testing.T) {
+		launch := newLaunch(map[string]string{"KEEP": "1"})
+
+		got := launch.extraEnvsForStart()
+		assertFitTarget(t, got, projectorPadMiB)
+		if _, ok := launch.extraEnvs[llamaArgFitTargetEnv]; ok {
+			t.Fatal("extraEnvsForStart mutated launch.extraEnvs")
+		}
+	})
+
+	for _, tt := range []struct {
+		name               string
+		launchFitTargetMiB uint64
+	}{
+		{name: "adds projector pad to existing launch fit target", launchFitTargetMiB: 2048},
+		{name: "adds projector pad to smaller launch fit target", launchFitTargetMiB: 512},
+	} {
+		t.Run(tt.name, func(t *testing.T) {
+			launch := newLaunch(map[string]string{
+				llamaArgFitTargetEnv: fitTargetValue(tt.launchFitTargetMiB),
+			})
+
+			got := launch.extraEnvsForStart()
+			assertFitTarget(t, got, tt.launchFitTargetMiB+projectorPadMiB)
+		})
+	}
+
+	t.Run("preserves inherited user fit target", func(t *testing.T) {
+		t.Setenv(llamaArgFitTargetEnv, fitTargetValue(256))
+		launch := newLaunch(map[string]string{})
+
+		got := launch.extraEnvsForStart()
+		if _, ok := got[llamaArgFitTargetEnv]; ok {
+			t.Fatalf("user env should not be overridden, got extra env %q", got[llamaArgFitTargetEnv])
+		}
+	})
+}
+
 func TestMMProjMemoryRequirement(t *testing.T) {
 	if got, err := mmprojMemoryRequirement("model.gguf", nil, nil); err != nil || got != 0 {
 		t.Fatalf("no projector memory = %d, %v; want 0, nil", got, err)

ml/device.go

@@ -585,7 +585,7 @@ func FlashAttentionSupported(l []DeviceInfo) bool {
 
 func cudaFlashAttentionSupported(gpu DeviceInfo) bool {
 	if gpu.Library != "CUDA" ||
-		gpu.ComputeMajor < 7 ||
+		gpu.ComputeMajor < 6 ||
 		(gpu.ComputeMajor == 7 && gpu.ComputeMinor == 2) {
 		return false
 	}

ml/device_test.go

@@ -186,8 +186,19 @@ func TestFlashAttentionSupported(t *testing.T) {
 			gpus: []DeviceInfo{{DeviceID: DeviceID{Library: "CUDA"}, DriverMajor: 12, ComputeMajor: 5, ComputeMinor: 0}},
 		},
 		{
-			name: "cuda compute 6.2 unsupported",
+			name: "cuda compute 6.0 supported",
+			gpus: []DeviceInfo{{DeviceID: DeviceID{Library: "CUDA"}, DriverMajor: 12, ComputeMajor: 6, ComputeMinor: 0}},
+			want: true,
+		},
+		{
+			name: "cuda compute 6.1 supported",
+			gpus: []DeviceInfo{{DeviceID: DeviceID{Library: "CUDA"}, DriverMajor: 12, ComputeMajor: 6, ComputeMinor: 1}},
+			want: true,
+		},
+		{
+			name: "cuda compute 6.2 supported",
 			gpus: []DeviceInfo{{DeviceID: DeviceID{Library: "CUDA"}, DriverMajor: 12, ComputeMajor: 6, ComputeMinor: 2}},
+			want: true,
 		},
 		{
 			name: "cuda compute 7.2 unsupported",
@@ -216,7 +227,7 @@ func TestFlashAttentionSupported(t *testing.T) {
 			name: "mixed cuda unsupported",
 			gpus: []DeviceInfo{
 				{DeviceID: DeviceID{Library: "CUDA"}, DriverMajor: 12, ComputeMajor: 8, ComputeMinor: 9},
-				{DeviceID: DeviceID{Library: "CUDA"}, DriverMajor: 12, ComputeMajor: 6, ComputeMinor: 2},
+				{DeviceID: DeviceID{Library: "CUDA"}, DriverMajor: 12, ComputeMajor: 5, ComputeMinor: 0},
 			},
 		},
 		{

ml/nn/attention.go

@@ -1,84 +0,0 @@
-package nn
-
-import (
-	"fmt"
-
-	"github.com/ollama/ollama/kvcache"
-	"github.com/ollama/ollama/ml"
-)
-
-// Attention implements scaled dot-product attention for transformer models:
-// Attention(Q, K, V) = softmax(QK^T/√d_k)V
-//
-// Parameters:
-//   - ctx: Context for tensor operations
-//   - query: Query tensor (Q) with shape [d_k, heads, seq_len_q]
-//   - key: Key tensor (K) with shape [d_k, kv_heads, seq_len_k], can be nil to read from cache only
-//   - value: Value tensor (V) with shape [d_v, kv_heads, seq_len_k], can be nil to read from cache only
-//   - scale: Scaling factor, typically 1/√d_k where d_k is the key dimension
-//   - cache: KV cache to store key/value and get past history, can be nil to only use provided key/value
-//
-// Returns:
-//
-//	Attention output with shape [d_v, heads, seq_len_q]
-func Attention(ctx ml.Context, query, key, value ml.Tensor, scale float64, cache kvcache.Cache) ml.Tensor {
-	return AttentionWithVMLA(ctx, query, key, value, nil, nil, scale, cache)
-}
-
-func AttentionWithSinks(ctx ml.Context, query, key, value, sinks ml.Tensor, scale float64, cache kvcache.Cache) ml.Tensor {
-	return AttentionWithVMLA(ctx, query, key, value, sinks, nil, scale, cache)
-}
-
-func AttentionWithVMLA(ctx ml.Context, query, key, value, sinks ml.Tensor, vmla ml.Tensor, scale float64, cache kvcache.Cache) ml.Tensor {
-	ctx.Forward(query)
-	if key != nil && value != nil {
-		if query.Dim(0) != key.Dim(0) {
-			panic(fmt.Errorf("d_k in attention operation does not match between query(%v) and key(%v)", query.Dim(0), key.Dim(0)))
-		}
-
-		if key.Dim(1) != value.Dim(1) {
-			panic(fmt.Errorf("kv_heads in attention operation does not match between key(%v) and value(%v)", key.Dim(1), value.Dim(1)))
-		}
-
-		if key.Dim(2) != value.Dim(2) {
-			panic(fmt.Errorf("seq_len_k in attention operation does not match between key(%v) and value(%v)", key.Dim(2), value.Dim(2)))
-		}
-
-		ctx.Forward(key, value)
-		if cache != nil {
-			cache.Put(ctx, key, value)
-		}
-	} else if cache == nil {
-		panic("key & value tensors must be provided if cache is nil")
-	}
-
-	var mask ml.Tensor
-	if cache != nil {
-		key, value, mask = cache.Get(ctx)
-	}
-
-	if sdpa, ok := query.(ml.ScaledDotProductAttention); ok {
-		cacheConfigApplied := cache != nil
-		return sdpa.ScaledDotProductAttention(ctx, key, value, mask, sinks, vmla, scale, cacheConfigApplied)
-	} else {
-		query = query.Permute(ctx, 0, 2, 1, 3)
-		key = key.Permute(ctx, 0, 2, 1, 3)
-		value = value.Permute(ctx, 1, 2, 0, 3).Contiguous(ctx)
-
-		kq := key.MulmatFullPrec(ctx, query)
-
-		kq = kq.Scale(ctx, scale)
-		if mask != nil {
-			kq = kq.Add(ctx, mask)
-		}
-		kq = kq.Softmax(ctx)
-
-		kqv := value.Mulmat(ctx, kq)
-
-		if vmla != nil {
-			kqv = vmla.Mulmat(ctx, kqv)
-		}
-
-		return kqv.Permute(ctx, 0, 2, 1, 3).Contiguous(ctx)
-	}
-}

ml/nn/convolution.go

@@ -1,30 +0,0 @@
-package nn
-
-import "github.com/ollama/ollama/ml"
-
-type Conv2D struct {
-	Weight ml.Tensor `gguf:"weight"`
-	Bias   ml.Tensor `gguf:"bias"`
-}
-
-func (m *Conv2D) Forward(ctx ml.Context, t ml.Tensor, s0, s1, p0, p1, d0, d1 int) ml.Tensor {
-	t = m.Weight.Conv2D(ctx, t, s0, s1, p0, p1, d0, d1)
-	if m.Bias != nil {
-		// Bias shape is (out_channels,) while t shape is (width, height, out_channels, batch)
-		t = t.Add(ctx, m.Bias.Reshape(ctx, 1, 1, -1))
-	}
-	return t
-}
-
-type Conv3D struct {
-	Weight ml.Tensor `gguf:"weight"`
-	Bias   ml.Tensor `gguf:"bias"`
-}
-
-func (m *Conv3D) Forward(ctx ml.Context, t ml.Tensor, c, s0, s1, s2, p0, p1, p2, d0, d1, d2 int) ml.Tensor {
-	t = m.Weight.Conv3D(ctx, t, c, s0, s1, s2, p0, p1, p2, d0, d1, d2)
-	if m.Bias != nil {
-		t = t.Add(ctx, m.Bias)
-	}
-	return t
-}

ml/nn/embedding.go

@@ -1,11 +0,0 @@
-package nn
-
-import "github.com/ollama/ollama/ml"
-
-type Embedding struct {
-	Weight ml.Tensor `gguf:"weight"`
-}
-
-func (m *Embedding) Forward(ctx ml.Context, hiddenState ml.Tensor) ml.Tensor {
-	return m.Weight.Rows(ctx, hiddenState)
-}

ml/nn/linear.go

@@ -1,31 +0,0 @@
-package nn
-
-import "github.com/ollama/ollama/ml"
-
-type Linear struct {
-	Weight ml.Tensor `gguf:"weight"`
-	Bias   ml.Tensor `gguf:"bias"`
-}
-
-func (m *Linear) Forward(ctx ml.Context, t ml.Tensor) ml.Tensor {
-	t = m.Weight.Mulmat(ctx, t)
-	if m.Bias != nil {
-		t = t.Add(ctx, m.Bias)
-	}
-
-	return t
-}
-
-type LinearBatch struct {
-	Weight ml.Tensor `gguf:"weight"`
-	Bias   ml.Tensor `gguf:"bias"`
-}
-
-func (m *LinearBatch) Forward(ctx ml.Context, t, indices ml.Tensor) ml.Tensor {
-	t = m.Weight.MulmatID(ctx, t, indices)
-	if m.Bias != nil {
-		t = t.AddID(ctx, m.Bias, indices)
-	}
-
-	return t
-}

ml/nn/normalization.go

@@ -1,22 +0,0 @@
-package nn
-
-import (
-	"github.com/ollama/ollama/ml"
-)
-
-type LayerNorm struct {
-	Weight ml.Tensor `gguf:"weight"`
-	Bias   ml.Tensor `gguf:"bias"`
-}
-
-func (m *LayerNorm) Forward(ctx ml.Context, t ml.Tensor, eps float32) ml.Tensor {
-	return t.LayerNorm(ctx, m.Weight, m.Bias, eps)
-}
-
-type RMSNorm struct {
-	Weight ml.Tensor `gguf:"weight"`
-}
-
-func (m *RMSNorm) Forward(ctx ml.Context, t ml.Tensor, eps float32) ml.Tensor {
-	return t.RMSNorm(ctx, m.Weight, eps)
-}

ml/nn/pooling/pooling.go

@@ -1,41 +0,0 @@
-package pooling
-
-import (
-	"github.com/ollama/ollama/ml"
-)
-
-type Type uint32
-
-const (
-	TypeNone Type = iota
-	TypeMean
-	TypeCLS
-	TypeLast
-)
-
-func (t Type) String() string {
-	switch t {
-	case TypeMean:
-		return "Mean"
-	case TypeCLS:
-		return "CLS"
-	case TypeLast:
-		return "Last"
-	default:
-		return "Unknown"
-	}
-}
-
-func (t Type) Forward(ctx ml.Context, hiddenStates ml.Tensor) ml.Tensor {
-	switch t {
-	case TypeMean:
-		hiddenStates = hiddenStates.Permute(ctx, 1, 0, 2, 3).Contiguous(ctx).Mean(ctx)
-		return hiddenStates.Permute(ctx, 1, 0, 2, 3).Contiguous(ctx)
-	case TypeCLS:
-		return hiddenStates.Slice(ctx, 1, 0, 1, 1)
-	case TypeLast:
-		return hiddenStates.Slice(ctx, 1, hiddenStates.Dim(1)-1, hiddenStates.Dim(1), 1)
-	default:
-		panic("unknown pooling type")
-	}
-}

ml/nn/rope.go

@@ -1,20 +0,0 @@
-package nn
-
-import (
-	"github.com/ollama/ollama/ml"
-	"github.com/ollama/ollama/ml/nn/rope"
-)
-
-// fastRoPE is an interface for tensors that support fast rotary positional embedding.
-type fastRoPE interface {
-	RoPE(ctx ml.Context, positions ml.Tensor, dim int, base, scale float32, options ...func(*rope.Options)) ml.Tensor
-}
-
-// RoPE applies rotary positional embedding to tensor `t`.
-func RoPE(ctx ml.Context, t, positions ml.Tensor, dim int, base, scale float32, options ...func(*rope.Options)) ml.Tensor {
-	if t, ok := t.(fastRoPE); ok {
-		return t.RoPE(ctx, positions, dim, base, scale, options...)
-	}
-
-	panic("RoPE not implemented for this tensor type")
-}

ml/nn/rope/options.go

@@ -1,92 +0,0 @@
-// Package rope provides options for RoPE
-package rope
-
-import "github.com/ollama/ollama/ml"
-
-// Options contains optional parameters for RoPE function
-type Options struct {
-	Type    int
-	Factors ml.Tensor
-
-	// YaRN options
-	YaRN struct {
-		OriginalContextLength int
-		ExtrapolationFactor,
-		AttentionFactor,
-		BetaFast,
-		BetaSlow float32
-	}
-
-	// MRoPE options
-	MRoPE struct {
-		Sections []int
-	}
-}
-
-// WithTypeNeoX sets RoPE type to NeoX
-func WithTypeNeoX() func(*Options) {
-	return func(opts *Options) {
-		opts.Type = 2
-	}
-}
-
-// WithFactors sets custom rope factors
-func WithFactors(factors ml.Tensor) func(*Options) {
-	return func(opts *Options) {
-		if factors != nil {
-			opts.Factors = factors
-		}
-	}
-}
-
-// WithOriginalContextLength sets a custom context length
-func WithOriginalContextLength(n int) func(*Options) {
-	return func(opts *Options) {
-		opts.YaRN.OriginalContextLength = n
-	}
-}
-
-func WithExtrapolationFactor(extrapolationFactor float32) func(*Options) {
-	return func(opts *Options) {
-		opts.YaRN.ExtrapolationFactor = extrapolationFactor
-	}
-}
-
-func WithAttentionFactor(attentionFactor float32) func(*Options) {
-	return func(opts *Options) {
-		opts.YaRN.AttentionFactor = attentionFactor
-	}
-}
-
-func WithBetaFast(betaFast float32) func(*Options) {
-	return func(opts *Options) {
-		opts.YaRN.BetaFast = betaFast
-	}
-}
-
-func WithBetaSlow(betaSlow float32) func(*Options) {
-	return func(opts *Options) {
-		opts.YaRN.BetaSlow = betaSlow
-	}
-}
-
-func WithMRoPE(sections []int) func(*Options) {
-	return func(opts *Options) {
-		opts.Type |= 1 << 3
-		opts.MRoPE.Sections = sections
-	}
-}
-
-func WithVision(sections []int) func(*Options) {
-	return func(opts *Options) {
-		opts.Type |= 1<<3 | 1<<4
-		opts.MRoPE.Sections = sections
-	}
-}
-
-func WithInterleaveMRoPE(sections []int) func(*Options) {
-	return func(opts *Options) {
-		opts.Type |= 1<<3 | 1<<5
-		opts.MRoPE.Sections = sections
-	}
-}

model/input/input.go

@@ -1,72 +0,0 @@
-package input
-
-import "github.com/ollama/ollama/ml"
-
-// Multimodal is a multimodal embedding or a component of one.
-// For example, it could be a row of an image that can be processed
-// independently.
-type Multimodal struct {
-	// Tensor is the embedding data. Implementations may chose what to
-	// store here or it may be nil if not needed. However, any ml.Tensor
-	// objects must be stored here and not in Data.
-	Tensor ml.Tensor
-
-	// Data is implementation-specific opaque data, such as metadata on how
-	// to layout Tensor. It may be nil if not needed. It may also store larger
-	// objects such as complete images if they are to be processed later.
-	Data any
-}
-
-// Input represents one token in the input stream
-type Input struct {
-	// Token is a single element of text.
-	Token int32
-
-	// Multimodal is represents a non-text element such as an
-	// image (or part of one if the image can be processed in pieces).
-	// It may be used either together with Token or on its own.
-	Multimodal []Multimodal
-
-	// MultimodalHash is a unique representation of the data
-	// stored in Multimodal, used for caching and comparing
-	// equality.
-	MultimodalHash uint64
-
-	// SameBatch forces the following number of tokens to be processed
-	// in a single batch, breaking and extending batches as needed.
-	// Useful for things like images that must be processed in one
-	// shot.
-	SameBatch int
-}
-
-// MultimodalIndex is a multimodal element (such as an image)
-// together with an index into the slice of Inputs with the
-// corresponding token. Note that the index is not the same
-// as the position - to find that use the index with the
-// Positions slice.
-type MultimodalIndex struct {
-	Index      int
-	Multimodal []Multimodal
-}
-
-// Batch contains the inputs for a model forward pass
-type Batch struct {
-	// Inputs is the input tokens, including placeholders for multimodal inputs.
-	Inputs ml.Tensor
-
-	// Outputs are the set of indicies into Inputs for which output data should
-	// be returned.
-	Outputs ml.Tensor
-
-	// Positions is the position for each Input, relative to its sequence. Equal
-	// in length to Inputs.
-	Positions []int32
-
-	// Sequences is the sequence for each Input. Equal in length to Inputs.
-	Sequences []int
-
-	// Multimodal is a set of multimodal embeddings previously created by
-	// EncodeMultimodal, along with an index into Inputs. Unused for text-only
-	// models or for batches without multimodal elements.
-	Multimodal []MultimodalIndex
-}

model/model.go

@@ -1,353 +0,0 @@
-package model
-
-import (
-	"errors"
-	"fmt"
-	_ "image/jpeg"
-	_ "image/png"
-	"log/slog"
-	"os"
-	"reflect"
-	"strconv"
-	"strings"
-
-	_ "golang.org/x/image/bmp"
-	_ "golang.org/x/image/tiff"
-	_ "golang.org/x/image/webp"
-
-	"github.com/ollama/ollama/fs"
-	fsggml "github.com/ollama/ollama/fs/ggml"
-	"github.com/ollama/ollama/kvcache"
-	"github.com/ollama/ollama/logutil"
-	"github.com/ollama/ollama/ml"
-	"github.com/ollama/ollama/ml/nn/pooling"
-	"github.com/ollama/ollama/model/input"
-	"github.com/ollama/ollama/tokenizer"
-)
-
-var (
-	ErrNoVisionModel        = errors.New("this model is missing data required for image input")
-	ErrUnsupportedModel     = errors.New("model not supported")
-	ErrUnsupportedTokenizer = errors.New("tokenizer not supported")
-)
-
-// Model implements a specific model architecture, defining the forward pass and any model-specific configuration
-type Model interface {
-	Forward(ml.Context, input.Batch) (ml.Tensor, error)
-
-	Backend() ml.Backend
-	Config() config
-}
-
-// Validator is an optional interface that models can implement to perform
-// validation after tensors have been loaded. If validation fails, model
-// loading will fail with the returned error.
-type Validator interface {
-	Validate() error
-}
-
-// PostLoader is an optional interface that models can implement to run
-// initialization steps after backend weights have been loaded.
-type PostLoader interface {
-	PostLoad() error
-}
-
-// MultimodalProcessor must be implemented by multimodal models.
-type MultimodalProcessor interface {
-	// EncodeMultimodal processes a single input (such as an image) and
-	// generates an output (typically an embedding) that can be used by the model.
-	//
-	// The return value is one or more tensors, each with optional model-specific
-	// opaque metadata. Typically, the tensors might be views into an embedding
-	// with each view representing a chunk of data that can be processed independently
-	// in different batches.
-	//
-	// The result may be cached by the runner.
-	EncodeMultimodal(ml.Context, []byte) ([]input.Multimodal, error)
-
-	// PostTokenize is called after tokenization to allow the model to edit the
-	// input stream to correctly arrange multimodal elements.
-	//
-	// The input is a slice of tokens with the results of EncodeMultimodal interleaved
-	// in the order that the user provided them. Each element of the slice will be
-	// either a single token or single multimodal object.
-	//
-	// The model must ensure that inputs are stored according to how they will be
-	// processed and stored in the cache. For example, Llava-style models should insert
-	// placeholder tokens equal to the feature size of the corresponding image with
-	// the image itself attached to and split across these tokens. When Forward is called
-	// a partial subset of these tokens may be submitted according to the batch size.
-	//
-	// This function is also responsible for updating MultimodalHash for any Multimodal
-	// that is modified to ensure that there is a unique hash value that accurately
-	// represents the contents.
-	PostTokenize([]*input.Input) ([]*input.Input, error)
-}
-
-// Base implements the common fields and methods for all models
-type Base struct {
-	b ml.Backend
-	config
-}
-
-type config struct {
-	Cache kvcache.Cache
-}
-
-// Backend returns the underlying backend that will run the model
-func (m *Base) Backend() ml.Backend {
-	return m.b
-}
-
-func (m *Base) Config() config {
-	return m.config
-}
-
-var models = make(map[string]func(fs.Config) (Model, error))
-
-// Register registers a model constructor for the given architecture
-func Register(name string, f func(fs.Config) (Model, error)) {
-	if _, ok := models[name]; ok {
-		panic("model: model already registered")
-	}
-
-	models[name] = f
-}
-
-// New initializes a new model instance with the provided configuration based on the metadata in the model file
-func New(modelPath string, params ml.BackendParams) (Model, error) {
-	b, err := ml.NewBackend(modelPath, params)
-	if err != nil {
-		return nil, err
-	}
-
-	m, err := modelForArch(b.Config())
-	if err != nil {
-		return nil, err
-	}
-
-	base := Base{b: b, config: m.Config()}
-	v := reflect.ValueOf(m)
-	v.Elem().Set(populateFields(base, v.Elem()))
-
-	if validator, ok := m.(Validator); ok {
-		if err := validator.Validate(); err != nil {
-			return nil, err
-		}
-	}
-
-	return m, nil
-}
-
-func NewTextProcessor(s string) (tokenizer.Tokenizer, error) {
-	r, err := os.Open(s)
-	if err != nil {
-		return nil, err
-	}
-	defer r.Close()
-
-	meta, err := fsggml.Decode(r, -1)
-	if err != nil {
-		return nil, err
-	}
-
-	m, err := modelForArch(meta.KV())
-	if err != nil {
-		return nil, err
-	}
-
-	tp, ok := m.(tokenizer.Tokenizer)
-	if !ok {
-		return nil, ErrUnsupportedTokenizer
-	}
-	return tp, nil
-}
-
-func modelForArch(c fs.Config) (Model, error) {
-	arch := c.Architecture()
-	if pooling.Type(c.Uint("pooling_type")) != pooling.TypeNone {
-		arch = arch + "_embed"
-	}
-
-	f, ok := models[arch]
-	if !ok {
-		return nil, ErrUnsupportedModel
-	}
-
-	return f(c)
-}
-
-func populateFields(base Base, v reflect.Value, tags ...Tag) reflect.Value {
-	t := v.Type()
-
-	if t.Kind() == reflect.Struct {
-		allNil := true
-		for i := range t.NumField() {
-			tt := t.Field(i).Type
-			vv := v.Field(i)
-			if !vv.CanSet() {
-				continue
-			}
-
-			// make a copy
-			tagsCopy := tags
-			if tag := t.Field(i).Tag.Get("gguf"); tag != "" {
-				tagsCopy = append(tagsCopy, parseTag(tag))
-			}
-
-			if tt == reflect.TypeOf((*Base)(nil)).Elem() {
-				vv.Set(reflect.ValueOf(base))
-			} else if tt == reflect.TypeOf((*ml.Tensor)(nil)).Elem() {
-				var fn func([]Tag, string, string) [][]string
-				fn = func(tags []Tag, prefix, suffix string) (fullNames [][]string) {
-					if len(tags) > 0 {
-						var names []string
-						if tags[0].name != "" {
-							for _, n := range append([]string{tags[0].name}, tags[0].alternatives...) {
-								names = append(names, prefix+n+suffix)
-							}
-						}
-						childNames := fn(tags[1:], tags[0].prefix, tags[0].suffix)
-						if len(names) == 0 {
-							// current tag has no name, use child names only
-							fullNames = append(fullNames, childNames...)
-						} else if len(childNames) == 0 {
-							// current tag has names but no children, create branches for each name
-							for _, name := range names {
-								fullNames = append(fullNames, []string{name})
-							}
-						} else {
-							// merge each name with each child
-							for _, name := range names {
-								for _, childName := range childNames {
-									fullNames = append(fullNames, append([]string{name}, childName...))
-								}
-							}
-						}
-					}
-
-					return fullNames
-				}
-
-				names := fn(tagsCopy, "", "")
-				for _, name := range names {
-					if tensor := base.Backend().Get(strings.Join(name, ".")); tensor != nil {
-						logutil.Trace("found tensor", "", tensor)
-						vv.Set(reflect.ValueOf(tensor))
-						break
-					}
-				}
-			} else if tt.Kind() == reflect.Pointer || tt.Kind() == reflect.Interface {
-				setPointer(base, vv, tagsCopy)
-			} else if tt.Kind() == reflect.Slice || tt.Kind() == reflect.Array {
-				for i := range vv.Len() {
-					vvv := vv.Index(i)
-					if vvv.Kind() == reflect.Pointer || vvv.Kind() == reflect.Interface {
-						setPointer(base, vvv, append(tagsCopy, Tag{name: strconv.Itoa(i)}))
-					} else {
-						vvv.Set(populateFields(base, vvv, append(tagsCopy, Tag{name: strconv.Itoa(i)})...))
-					}
-				}
-			}
-
-			if !canNil(tt) || !vv.IsNil() {
-				allNil = false
-			}
-		}
-
-		if allNil {
-			return reflect.Zero(t)
-		}
-	}
-
-	return v
-}
-
-func setPointer(base Base, v reflect.Value, tags []Tag) {
-	vv := v
-	if v.Kind() == reflect.Interface {
-		if v.IsNil() {
-			return
-		}
-
-		vv = vv.Elem()
-	}
-
-	vv = reflect.Indirect(vv)
-	if v.IsNil() {
-		vv = reflect.New(v.Type().Elem()).Elem()
-	}
-
-	if f := populateFields(base, vv, tags...); f.CanAddr() {
-		v.Set(f.Addr())
-	}
-}
-
-type Tag struct {
-	name,
-	// prefix and suffix are applied to child tags
-	prefix,
-	suffix string
-	alternatives []string
-}
-
-func parseTag(s string) (tag Tag) {
-	parts := strings.Split(s, ",")
-	if len(parts) > 0 {
-		tag.name = parts[0]
-
-		for _, part := range parts[1:] {
-			if value, ok := strings.CutPrefix(part, "alt:"); ok && tag.name == "" {
-				// elevate alternative to primary if no primary given
-				tag.name = value
-				slog.Warn("gguf tag has alt: but no primary name", "tag", s)
-			} else if ok {
-				tag.alternatives = append(tag.alternatives, value)
-			}
-			if value, ok := strings.CutPrefix(part, "pre:"); ok {
-				tag.prefix = value
-			}
-			if value, ok := strings.CutPrefix(part, "suf:"); ok {
-				tag.suffix = value
-			}
-		}
-	}
-
-	return
-}
-
-func canNil(t reflect.Type) bool {
-	return t.Kind() == reflect.Chan ||
-		t.Kind() == reflect.Func ||
-		t.Kind() == reflect.Interface ||
-		t.Kind() == reflect.Map ||
-		t.Kind() == reflect.Pointer ||
-		t.Kind() == reflect.Slice
-}
-
-func Forward(ctx ml.Context, m Model, batch input.Batch) (ml.Tensor, error) {
-	if len(batch.Positions) != len(batch.Sequences) {
-		return nil, fmt.Errorf("length of positions (%v) must match length of seqs (%v)", len(batch.Positions), len(batch.Sequences))
-	}
-
-	if len(batch.Positions) < 1 {
-		return nil, errors.New("batch size cannot be less than 1")
-	}
-
-	cache := m.Config().Cache
-	if cache != nil {
-		err := cache.StartForward(ctx, batch, false)
-		if err != nil {
-			return nil, err
-		}
-	}
-
-	t, err := m.Forward(ctx, batch)
-	if err != nil {
-		return nil, err
-	}
-
-	ctx.Forward(t)
-
-	return t, nil
-}

model/models/gemma4/model.go

@@ -1,265 +0,0 @@
-package gemma4
-
-import (
-	"bytes"
-	"fmt"
-	"image"
-	"log/slog"
-	"slices"
-	"time"
-
-	"github.com/ollama/ollama/fs"
-	"github.com/ollama/ollama/kvcache"
-	"github.com/ollama/ollama/ml"
-	"github.com/ollama/ollama/ml/nn"
-	"github.com/ollama/ollama/ml/nn/rope"
-	"github.com/ollama/ollama/model"
-	"github.com/ollama/ollama/model/input"
-	"github.com/ollama/ollama/tokenizer"
-)
-
-type Model struct {
-	model.Base
-	tokenizer.Tokenizer
-
-	*VisionModel `gguf:"v"`
-	*TextModel
-	*AudioModel `gguf:"a"`
-
-	*MultiModalProjector      `gguf:"mm"`
-	*AudioMultimodalProjector `gguf:"mm.a"`
-
-	ImageProcessor
-
-	imageTokenID    int32
-	imageEndTokenID int32
-	audioTokenID    int32
-	audioEndTokenID int32
-
-	audioOpts *AudioModelOptions
-}
-
-var _ model.MultimodalProcessor = (*Model)(nil)
-
-type MultiModalProjector struct {
-	Projection *ClippableLinear `gguf:"input_projection"`
-}
-
-func (p *MultiModalProjector) Forward(ctx ml.Context, visionOutputs ml.Tensor, eps float32) ml.Tensor {
-	visionOutputs = p.Projection.Forward(ctx, visionOutputs)
-	// Post-projection RMSNorm without learned weight
-	visionOutputs = visionOutputs.RMSNorm(ctx, nil, eps)
-	return visionOutputs
-}
-
-func New(c fs.Config) (model.Model, error) {
-	vocabulary := tokenizer.Vocabulary{
-		Values: c.Strings("tokenizer.ggml.tokens"),
-		Scores: c.Floats("tokenizer.ggml.scores"),
-		Types:  c.Ints("tokenizer.ggml.token_type"),
-		Merges: c.Strings("tokenizer.ggml.merges"),
-		AddBOS: c.Bool("tokenizer.ggml.add_bos_token", false),
-		BOS:    []int32{int32(c.Uint("tokenizer.ggml.bos_token_id"))},
-		AddEOS: c.Bool("tokenizer.ggml.add_eos_token", false),
-		EOS: append(
-			[]int32{
-				int32(c.Uint("tokenizer.ggml.eos_token_id")),
-			},
-			c.Ints("tokenizer.ggml.eos_token_ids")...,
-		),
-	}
-
-	vocabulary.EOS = append(vocabulary.EOS, int32(c.Uint("tokenizer.ggml.eot_token_id", 106)))
-
-	// Gemma 4 uses BPE with SentencePiece-style ▁ space markers (not GPT-2 byte-level encoding).
-	// The tokenizer.json has merges and a Replace normalizer (space → ▁), with no pre-tokenizer.
-	t := tokenizer.NewBytePairEncodingWithOptions(&vocabulary, []string{},
-		tokenizer.WithSentencePieceNormalizer())
-
-	// Look up special token IDs for vision and audio
-	imageTokenID := int32(-1)
-	imageEndTokenID := int32(-1)
-	audioTokenID := int32(-1)
-	audioEndTokenID := int32(-1)
-	for i, tok := range vocabulary.Values {
-		switch tok {
-		case "<|image>":
-			imageTokenID = int32(i)
-		case "<image|>":
-			imageEndTokenID = int32(i)
-		case "<|audio>":
-			audioTokenID = int32(i)
-		case "<audio|>":
-			audioEndTokenID = int32(i)
-		}
-	}
-
-	slog.Info("gemma4: token IDs", "image", imageTokenID, "image_end", imageEndTokenID, "audio", audioTokenID, "audio_end", audioEndTokenID)
-
-	m := Model{
-		Tokenizer:                t,
-		TextModel:                newTextModel(c),
-		VisionModel:              newVisionModel(c),
-		AudioModel:               newAudioModel(c),
-		MultiModalProjector:      &MultiModalProjector{},
-		AudioMultimodalProjector: &AudioMultimodalProjector{},
-		ImageProcessor:           newImageProcessor(c),
-		imageTokenID:             imageTokenID,
-		imageEndTokenID:          imageEndTokenID,
-		audioTokenID:             audioTokenID,
-		audioEndTokenID:          audioEndTokenID,
-		audioOpts:                newAudioModelOptions(c),
-	}
-
-	slidingWindowLen := int32(c.Uint("attention.sliding_window"))
-	m.Cache = kvcache.NewWrapperCache(
-		kvcache.NewSWAMemCache(slidingWindowLen, 4096, m.Shift),
-		kvcache.NewCausalCache(m.Shift),
-	)
-
-	return &m, nil
-}
-
-func (m *Model) EncodeMultimodal(ctx ml.Context, multimodalData []byte) ([]input.Multimodal, error) {
-	// Audio input: detect WAV format and route to audio encoder.
-	if isAudioData(multimodalData) {
-		return m.encodeAudioMultimodal(ctx, multimodalData)
-	}
-
-	if len(m.VisionModel.Layers) == 0 {
-		return nil, model.ErrNoVisionModel
-	}
-
-	t0 := time.Now()
-	img, _, err := image.Decode(bytes.NewReader(multimodalData))
-	if err != nil {
-		return nil, err
-	}
-	slog.Info("vision: decode", "elapsed", time.Since(t0), "bounds", img.Bounds())
-
-	t1 := time.Now()
-	f32s, imgW, imgH, err := m.ImageProcessor.ProcessImage(img)
-	if err != nil {
-		return nil, err
-	}
-	slog.Info("vision: preprocess", "elapsed", time.Since(t1), "size", [2]int{imgW, imgH})
-
-	pixelValues := ctx.Input().FromFloats(f32s, imgW, imgH, m.ImageProcessor.numChannels)
-	slog.Info("vision: pixelValues", "shape", pixelValues.Shape(), "dim0", pixelValues.Dim(0), "dim1", pixelValues.Dim(1), "dim2", pixelValues.Dim(2))
-
-	numPatchesX := imgW / m.ImageProcessor.patchSize
-	numPatchesY := imgH / m.ImageProcessor.patchSize
-	slog.Info("vision: patches", "patchesX", numPatchesX, "patchesY", numPatchesY, "total", numPatchesX*numPatchesY, "patchSize", m.ImageProcessor.patchSize)
-
-	visionOutputs := m.VisionModel.Forward(ctx, pixelValues, numPatchesX, numPatchesY)
-	visionOutputs = visionPoolAndProject(ctx, visionOutputs, numPatchesX, numPatchesY, m.VisionModel.VisionModelOptions, m.MultiModalProjector, m.VisionModel.StdBias, m.VisionModel.StdScale)
-	slog.Info("vision: encoded", "elapsed", time.Since(t0), "shape", visionOutputs.Shape())
-
-	return []input.Multimodal{{Tensor: visionOutputs}}, nil
-}
-
-func (m *Model) PostLoad() error {
-	m.VisionModel.InitClamp(m.MultiModalProjector)
-	return nil
-}
-
-func (m *Model) encodeAudioMultimodal(ctx ml.Context, data []byte) ([]input.Multimodal, error) {
-	if m.AudioModel == nil || m.audioOpts == nil {
-		return nil, model.ErrNoVisionModel
-	}
-
-	t0 := time.Now()
-	samples, err := decodeWAV(data)
-	if err != nil {
-		return nil, err
-	}
-	slog.Info("audio: decode", "elapsed", time.Since(t0), "samples", len(samples), "duration_s", float64(len(samples))/audioSampleRate)
-
-	// Pad waveform to next multiple of 128.
-	if rem := len(samples) % 128; rem != 0 {
-		samples = append(samples, make([]float32, 128-rem)...)
-	}
-
-	// Compute mel spectrogram.
-	melData, numFrames := computeMelSpectrogram(samples)
-	if numFrames == 0 {
-		return nil, fmt.Errorf("audio too short to encode")
-	}
-	slog.Info("audio: mel", "frames", numFrames, "elapsed", time.Since(t0))
-
-	// Create input tensor [melBins, numFrames] (GGML ne order). FromFloats creates F32.
-	melTensor := ctx.Input().FromFloats(melData, melBins, numFrames)
-
-	// Run audio encoder.
-	audioOutputs := m.AudioModel.ForwardAudio(ctx, melTensor, m.AudioMultimodalProjector, m.audioOpts)
-	slog.Info("audio: encoded", "elapsed", time.Since(t0), "shape", audioOutputs.Shape())
-
-	return []input.Multimodal{{Tensor: audioOutputs, Data: audioTag{}}}, nil
-}
-
-// audioTag marks multimodal data as audio (vs vision) for PostTokenize.
-type audioTag struct{}
-
-func (m *Model) PostTokenize(inputs []*input.Input) ([]*input.Input, error) {
-	var result []*input.Input
-
-	for _, inp := range inputs {
-		if len(inp.Multimodal) == 0 {
-			result = append(result, inp)
-			continue
-		}
-
-		inputMultimodal := inp.Multimodal[0].Tensor
-		numTokens := inputMultimodal.Dim(1)
-
-		// Determine if this is audio or vision based on the tag.
-		_, isAudio := inp.Multimodal[0].Data.(audioTag)
-
-		var beginToken, endToken int32
-		if isAudio {
-			beginToken = m.audioTokenID
-			endToken = m.audioEndTokenID
-		} else {
-			beginToken = m.imageTokenID
-			endToken = m.imageEndTokenID
-		}
-
-		if beginToken >= 0 {
-			result = append(result, &input.Input{Token: beginToken, SameBatch: numTokens + 2})
-		}
-
-		result = append(result,
-			&input.Input{Multimodal: []input.Multimodal{{Tensor: inputMultimodal}}, MultimodalHash: inp.MultimodalHash},
-		)
-		result = append(result, slices.Repeat([]*input.Input{{Token: 0}}, numTokens-1)...)
-
-		if endToken >= 0 {
-			result = append(result, &input.Input{Token: endToken})
-		}
-	}
-
-	return result, nil
-}
-
-func (m *Model) Forward(ctx ml.Context, batch input.Batch) (ml.Tensor, error) {
-	hiddenState := m.TextModel.Forward(ctx, batch, m.Cache)
-
-	hiddenState = m.TextModel.Output.Forward(ctx, hiddenState)
-
-	if m.TextModel.TextOptions.finalLogitSoftcap > 0.0 {
-		hiddenState = hiddenState.Scale(ctx, 1.0/float64(m.TextModel.TextOptions.finalLogitSoftcap))
-		hiddenState = hiddenState.Tanh(ctx)
-		hiddenState = hiddenState.Scale(ctx, float64(m.TextModel.TextOptions.finalLogitSoftcap))
-	}
-
-	return hiddenState, nil
-}
-
-func (m *Model) Shift(ctx ml.Context, layer int, key, shift ml.Tensor) (ml.Tensor, error) {
-	ropeBase, ropeDims := m.TextModel.ropeForLayer(layer)
-	return nn.RoPE(ctx, key, shift, ropeDims, ropeBase, 1.0, rope.WithTypeNeoX()), nil
-}
-
-func init() {
-	model.Register("gemma4", New)
-}

model/models/gemma4/model_audio.go

@@ -1,611 +0,0 @@
-package gemma4
-
-import (
-	"math"
-
-	"github.com/ollama/ollama/fs"
-	"github.com/ollama/ollama/ml"
-	"github.com/ollama/ollama/ml/nn"
-)
-
-// AudioModel holds the audio encoder and configuration.
-type AudioModel struct {
-	// SSCP: Sub-Sample Convolution Projection.
-	SSCPConv0 *AudioConvBlock `gguf:"conv1d.0"`
-	SSCPConv1 *AudioConvBlock `gguf:"conv1d.1"`
-
-	// SSCP output projection (linear).
-	SSCPInputProj *nn.Linear `gguf:"pre_encode.out"`
-
-	// Conformer blocks.
-	Layers []AudioConformerBlock `gguf:"blk"`
-
-	// Output projection to embedder dimension.
-	OutputProj *AudioOutputProj `gguf:"output_proj"`
-
-	AudioModelOptions
-}
-
-type AudioOutputProj struct {
-	Weight ml.Tensor `gguf:"weight"`
-	Bias   ml.Tensor `gguf:"bias"`
-}
-
-// AudioModelOptions holds audio model hyperparameters.
-type AudioModelOptions struct {
-	hiddenSize     int
-	numHeads       int
-	headDim        int
-	ffnSize        int
-	numLayers      int
-	melBins        int
-	chunkSize      int
-	maxPast        int
-	maxFuture      int
-	contextSize    int
-	logitCap       float32
-	residualWeight float32
-	gradClip       float32
-	convKernelSize int
-	eps            float32
-}
-
-// AudioConvBlock is a single 2D convolution block for the SSCP.
-type AudioConvBlock struct {
-	Weight ml.Tensor     `gguf:"weight"`
-	Norm   *nn.LayerNorm `gguf:"norm"`
-}
-
-// AudioConformerBlock is a single conformer layer.
-// All tensors are flat at the block level (a.blk.N.<name>) using underscore naming.
-type AudioConformerBlock struct {
-	// Block-level norm
-	Norm *nn.RMSNorm `gguf:"layer_pre_norm"`
-
-	// FFW start
-	FFWNorm     *nn.RMSNorm           `gguf:"ffn_norm"`
-	FFWUp       *AudioClippableLinear `gguf:"ffn_up"`
-	FFWDown     *AudioClippableLinear `gguf:"ffn_down"`
-	FFWPostNorm *nn.RMSNorm           `gguf:"ffn_post_norm"`
-
-	// FFW end
-	FFWNorm1     *nn.RMSNorm           `gguf:"ffn_norm_1"`
-	FFWUp1       *AudioClippableLinear `gguf:"ffn_up_1"`
-	FFWDown1     *AudioClippableLinear `gguf:"ffn_down_1"`
-	FFWPostNorm1 *nn.RMSNorm           `gguf:"ffn_post_norm_1"`
-
-	// Attention
-	AttnQ        *AudioClippableLinear `gguf:"attn_q"`
-	AttnK        *AudioClippableLinear `gguf:"attn_k"`
-	AttnV        *AudioClippableLinear `gguf:"attn_v"`
-	AttnOut      *AudioClippableLinear `gguf:"attn_out"`
-	AttnPreNorm  *nn.RMSNorm           `gguf:"ln1"`
-	AttnPostNorm *nn.RMSNorm           `gguf:"ln2"`
-	LinearPos    ml.Tensor             `gguf:"linear_pos.weight"`
-	PerDimScale  ml.Tensor             `gguf:"per_dim_scale.weight"`
-
-	// LightConv1d
-	ConvPW1  *AudioClippableLinear `gguf:"conv_pw1"`
-	ConvPW2  *AudioClippableLinear `gguf:"conv_pw2"`
-	ConvDW   ml.Tensor             `gguf:"conv_dw.weight"`
-	ConvNorm *nn.RMSNorm           `gguf:"conv_norm"`
-	NormConv *nn.RMSNorm           `gguf:"norm_conv"`
-}
-
-// AudioClippableLinear is a linear layer with optional input/output clamping.
-type AudioClippableLinear struct {
-	Weight    ml.Tensor `gguf:"weight"`
-	Bias      ml.Tensor `gguf:"bias"`
-	InputMin  ml.Tensor `gguf:"input_min"`
-	InputMax  ml.Tensor `gguf:"input_max"`
-	OutputMin ml.Tensor `gguf:"output_min"`
-	OutputMax ml.Tensor `gguf:"output_max"`
-
-	// Cached scalar clamp values (populated on first forward).
-	inMin, inMax, outMin, outMax float32
-	clampsLoaded                 bool
-}
-
-func (l *AudioClippableLinear) loadClamps() {
-	if l.clampsLoaded {
-		return
-	}
-	l.clampsLoaded = true
-	if l.InputMin != nil {
-		vals := l.InputMin.BackendGet()
-		if len(vals) > 0 {
-			l.inMin = vals[0]
-		}
-	}
-	if l.InputMax != nil {
-		vals := l.InputMax.BackendGet()
-		if len(vals) > 0 {
-			l.inMax = vals[0]
-		}
-	}
-	if l.OutputMin != nil {
-		vals := l.OutputMin.BackendGet()
-		if len(vals) > 0 {
-			l.outMin = vals[0]
-		}
-	}
-	if l.OutputMax != nil {
-		vals := l.OutputMax.BackendGet()
-		if len(vals) > 0 {
-			l.outMax = vals[0]
-		}
-	}
-}
-
-func (l *AudioClippableLinear) Forward(ctx ml.Context, x ml.Tensor) ml.Tensor {
-	l.loadClamps()
-	if l.inMax != 0 {
-		x = x.Clamp(ctx, l.inMin, l.inMax)
-	}
-	out := l.Weight.Mulmat(ctx, x)
-	if l.Bias != nil {
-		out = out.Add(ctx, l.Bias)
-	}
-	if l.outMax != 0 {
-		out = out.Clamp(ctx, l.outMin, l.outMax)
-	}
-	return out
-}
-
-// AudioMultimodalProjector is the audio-to-text embedding projector.
-type AudioMultimodalProjector struct {
-	Projection *AudioClippableLinear `gguf:"input_projection"`
-	FC         *AudioFC              `gguf:"fc"`
-}
-
-type AudioFC struct {
-	Weight ml.Tensor `gguf:"weight"`
-	Bias   ml.Tensor `gguf:"bias"`
-}
-
-func (p *AudioMultimodalProjector) Forward(ctx ml.Context, x ml.Tensor, eps float32) ml.Tensor {
-	// FC: output projection from conformer to embedder dimension.
-	x = p.FC.Weight.Mulmat(ctx, x)
-	if p.FC.Bias != nil {
-		x = x.Add(ctx, p.FC.Bias)
-	}
-	// Pre-projection RMSNorm (without learned weight) — matches Python's embedding_pre_projection_norm.
-	x = x.RMSNorm(ctx, nil, eps)
-	// Embedding projection to text hidden size.
-	x = p.Projection.Forward(ctx, x)
-	return x
-}
-
-// ForwardAudio encodes mel spectrogram features into soft tokens.
-// melFeatures: float32 tensor with ne[0]=melBins, ne[1]=numFrames.
-// Returns: [hiddenSize, numTokens] tensor.
-func (m *AudioModel) ForwardAudio(ctx ml.Context, melFeatures ml.Tensor, proj *AudioMultimodalProjector, opts *AudioModelOptions) ml.Tensor {
-	// SSCP Conv2D input: ne[0]=F (freq/width), ne[1]=T (time/height), ne[2]=C_in, ne[3]=B
-	// melFeatures is [melBins, numFrames], add channel and batch dims.
-	x := melFeatures.Reshape(ctx, melFeatures.Dim(0), melFeatures.Dim(1), 1, 1)
-
-	// SSCP Conv block 0: [F, T, 1, 1] → [F', T', C0, 1]
-	x = forwardConvBlock(ctx, m.SSCPConv0, x, opts)
-
-	// SSCP Conv block 1: [F', T', C0, 1] → [F'', T'', C1, 1]
-	x = forwardConvBlock(ctx, m.SSCPConv1, x, opts)
-
-	// After conv blocks, layout is [F'', T'', C_out, B].
-	// Permute to [C_out*F'', T'', B] for linear projection (channels+freq in ne[0]).
-	fOut := x.Dim(0)
-	tOut := x.Dim(1)
-	cOut := x.Dim(2)
-	// Permute [F'', T'', C, B] → [C, F'', T'', B]
-	// (1,2,0,3): old[0]→pos1, old[1]→pos2, old[2]→pos0
-	x = x.Permute(ctx, 1, 2, 0, 3).Contiguous(ctx)
-	x = x.Reshape(ctx, cOut*fOut, tOut)
-
-	// Linear projection to hidden size.
-	x = m.SSCPInputProj.Forward(ctx, x)
-
-	// Build causal-valid mask for conformer attention.
-	causalMask := buildCausalValidMaskF32(opts.chunkSize, opts.maxPast, opts.maxFuture)
-
-	// Run conformer blocks.
-	for i := range m.Layers {
-		x = m.Layers[i].Forward(ctx, x, causalMask, opts, i)
-	}
-
-	// Output projection.
-	if m.OutputProj != nil {
-		x = m.OutputProj.Weight.Mulmat(ctx, x)
-		if m.OutputProj.Bias != nil {
-			x = x.Add(ctx, m.OutputProj.Bias)
-		}
-	}
-
-	// Audio embedder: project to text embedding space.
-	if proj != nil {
-		x = proj.Forward(ctx, x, opts.eps)
-	}
-
-	return x
-}
-
-// forwardConvBlock runs a single SSCP Conv2D block.
-// Conv2D receiver is the kernel, argument is the input data.
-// Input: [F, T, C_in, B]. Output: [F', T', C_out, B].
-func forwardConvBlock(ctx ml.Context, block *AudioConvBlock, x ml.Tensor, opts *AudioModelOptions) ml.Tensor {
-	// Conv2D: kernel.Conv2D(ctx, input, s0, s1, p0, p1, d0, d1)
-	// Kernel is 3x3, stride 2x2, padding 1x1 (matching SSCP config).
-	// Output layout: [F', T', C_out, B]
-	// Make weight contiguous — the shape reversal in the converter creates
-	// a tensor where the physical data order doesn't match ne[]/stride[].
-	weight := block.Weight.Contiguous(ctx)
-	x = weight.Conv2D(ctx, x, 2, 2, 1, 1, 1, 1)
-
-	// LayerNorm needs channels in ne[0]. Permute [F', T', C_out, B] → [C_out, F', T', B],
-	// norm, then permute back.
-	// GGML permute: axis i says where old axis i goes.
-	// (1,2,0,3): old[0]→pos1, old[1]→pos2, old[2]→pos0 → [C_out, F', T', B]
-	x = x.Permute(ctx, 1, 2, 0, 3).Contiguous(ctx)
-	x = block.Norm.Forward(ctx, x, opts.eps)
-	// (2,0,1,3): old[0]→pos2, old[1]→pos0, old[2]→pos1 → [F', T', C_out, B]
-	x = x.Permute(ctx, 2, 0, 1, 3).Contiguous(ctx)
-
-	x = x.RELU(ctx)
-	return x
-}
-
-// Forward runs a single conformer block.
-func (cb *AudioConformerBlock) Forward(ctx ml.Context, x ml.Tensor, causalMask []float32, opts *AudioModelOptions, blockIdx int) ml.Tensor {
-	// FFW start (half-residual).
-	x = cb.forwardFFW(ctx, cb.FFWNorm, cb.FFWUp, cb.FFWDown, cb.FFWPostNorm, x, opts)
-
-	// Self-attention.
-	x = cb.forwardAttention(ctx, x, causalMask, opts, blockIdx)
-
-	// Lightweight Conv1d.
-	x = cb.forwardLightConv(ctx, x, opts, blockIdx)
-
-	// FFW end (half-residual).
-	x = cb.forwardFFW(ctx, cb.FFWNorm1, cb.FFWUp1, cb.FFWDown1, cb.FFWPostNorm1, x, opts)
-
-	// Gradient clipping + final norm.
-	x = x.Clamp(ctx, -opts.gradClip, opts.gradClip)
-	x = cb.Norm.Forward(ctx, x, opts.eps)
-
-	return x
-}
-
-// forwardFFW runs a feedforward module with half-residual connection.
-func (cb *AudioConformerBlock) forwardFFW(ctx ml.Context, preNorm *nn.RMSNorm, up, down *AudioClippableLinear, postNorm *nn.RMSNorm, x ml.Tensor, opts *AudioModelOptions) ml.Tensor {
-	residual := x
-	x = x.Clamp(ctx, -opts.gradClip, opts.gradClip)
-	x = preNorm.Forward(ctx, x, opts.eps)
-	x = up.Forward(ctx, x)
-	x = x.SILU(ctx)
-	x = down.Forward(ctx, x)
-	x = x.Clamp(ctx, -opts.gradClip, opts.gradClip)
-	x = postNorm.Forward(ctx, x, opts.eps)
-	x = x.Scale(ctx, float64(opts.residualWeight))
-	return residual.Add(ctx, x)
-}
-
-// forwardAttention runs the conformer block-local attention with relative position embeddings.
-func (cb *AudioConformerBlock) forwardAttention(ctx ml.Context, x ml.Tensor, causalMask []float32, opts *AudioModelOptions, blockIdx int) ml.Tensor {
-	residual := x
-	x = x.Clamp(ctx, -opts.gradClip, opts.gradClip)
-	x = cb.AttnPreNorm.Forward(ctx, x, opts.eps)
-
-	hiddenSize := x.Dim(0)
-	seqLen := x.Dim(1)
-
-	// QKV projections: [hiddenSize, seqLen] → [headDim, numHeads, seqLen]
-	q := cb.AttnQ.Forward(ctx, x).Reshape(ctx, opts.headDim, opts.numHeads, seqLen)
-	k := cb.AttnK.Forward(ctx, x).Reshape(ctx, opts.headDim, opts.numHeads, seqLen)
-	v := cb.AttnV.Forward(ctx, x).Reshape(ctx, opts.headDim, opts.numHeads, seqLen)
-
-	// Per-dim scaling for queries: (headDim^-0.5 / log(2)) * softplus(per_dim_scale)
-	// per_dim_scale is already softplus'd from the converter.
-	qScale := float64(math.Pow(float64(opts.headDim), -0.5)) / math.Log(2)
-	q = q.Scale(ctx, qScale)
-	if cb.PerDimScale != nil {
-		q = q.Mul(ctx, cb.PerDimScale)
-	}
-
-	// Key scaling: softplus(1) / log(2) — matches the query base scaling convention.
-	kScale := math.Log(1+math.E) / math.Log(2)
-	k = k.Scale(ctx, kScale)
-
-	// Build sinusoidal position embeddings for the block-local context.
-	maxSpan := opts.maxPast + opts.maxFuture + 1 // 13 unique relative positions
-	posEmb := cb.buildPositionEmbeddings(ctx, maxSpan, opts)
-	// posEmb: [headDim, numHeads, maxSpan]
-
-	// Block-local attention: process chunks of size chunkSize.
-	chunkSize := opts.chunkSize
-	numChunks := (seqLen + chunkSize - 1) / chunkSize
-	contextSize := opts.contextSize
-
-	// Pad q/k/v to multiple of chunkSize on the time dimension (dim 2).
-	padT := numChunks*chunkSize - seqLen
-	if padT > 0 {
-		q = q.Pad(ctx, 0, 0, padT, 0)
-		k = k.Pad(ctx, 0, 0, padT, 0)
-		v = v.Pad(ctx, 0, 0, padT, 0)
-	}
-	paddedLen := numChunks * chunkSize
-
-	// Pad k/v for context extraction: add maxPast on left, (maxFuture+chunkSize-1) on right.
-	// Use Pad (right) + PadExt (left) workaround since PadExt+Slice has issues.
-	// Actually use Concat with zero tensors for reliable left-padding.
-	padLeft := opts.maxPast
-	padRight := opts.maxFuture + chunkSize - 1
-	zeroLeft := ctx.Input().FromFloats(make([]float32, opts.headDim*opts.numHeads*padLeft), opts.headDim, opts.numHeads, padLeft)
-	zeroRight := ctx.Input().FromFloats(make([]float32, opts.headDim*opts.numHeads*padRight), opts.headDim, opts.numHeads, padRight)
-	kPadded := zeroLeft.Concat(ctx, k, 2).Concat(ctx, zeroRight, 2)
-	vPadded := zeroLeft.Concat(ctx, v, 2).Concat(ctx, zeroRight, 2)
-
-	// Reshape q into chunks: [headDim, numHeads, numChunks, chunkSize]
-	qChunked := q.Reshape(ctx, opts.headDim, opts.numHeads, numChunks, chunkSize)
-
-	// Process each chunk and collect results.
-	chunkOutputs := make([]ml.Tensor, numChunks)
-	for u := range numChunks {
-		// Extract query block: [headDim, numHeads, 1, chunkSize] → [headDim, numHeads, chunkSize]
-		qBlock := qChunked.Slice(ctx, 2, u, u+1, 1).Reshape(ctx, opts.headDim, opts.numHeads, chunkSize)
-
-		// Extract key/value context: [headDim, numHeads, contextSize]
-		cStart := u * chunkSize // offset in kPadded (padLeft already accounts for left context)
-		kCtx := kPadded.Slice(ctx, 2, cStart, cStart+contextSize, 1).Contiguous(ctx)
-		vCtx := vPadded.Slice(ctx, 2, cStart, cStart+contextSize, 1).Contiguous(ctx)
-
-		// Content-content logits: qBlock^T @ kCtx → [chunkSize, contextSize] per head.
-		// Mulmat(a, b) = a^T @ b. We want Q^T K, so: kCtx.Mulmat(qBlock) but that gives
-		// [numHeads, chunkSize, contextSize] with wrong batching.
-		// Instead: permute to [headDim, chunkSize, numHeads] and [headDim, contextSize, numHeads]
-		// then Mulmat batches over numHeads.
-		// GGML permute(0,2,1,3): old[0]→0, old[1]→2, old[2]→1
-		qP := qBlock.Permute(ctx, 0, 2, 1, 3) // [headDim, chunkSize, numHeads]
-		kP := kCtx.Permute(ctx, 0, 2, 1, 3)   // [headDim, contextSize, numHeads]
-
-		termAC := kP.MulmatFullPrec(ctx, qP) // [contextSize, chunkSize, numHeads]
-
-		// Content-position logits: qBlock^T @ posEmb → [chunkSize, maxSpan] per head.
-		pP := posEmb.Permute(ctx, 0, 2, 1, 3)   // [headDim, maxSpan, numHeads]
-		termBDRaw := pP.MulmatFullPrec(ctx, qP) // [maxSpan, chunkSize, numHeads]
-
-		// Relative shift: [maxSpan, chunkSize, numHeads] → [contextSize, chunkSize, numHeads]
-		termBD := cb.relativeShiftGGML(ctx, termBDRaw, maxSpan, chunkSize, contextSize, opts.numHeads)
-
-		// Combined logits.
-		logits := termAC.Add(ctx, termBD)
-
-		// Logit softcap: tanh(logits / cap) * cap
-		logits = logits.Scale(ctx, 1.0/float64(opts.logitCap))
-		logits = logits.Tanh(ctx)
-		logits = logits.Scale(ctx, float64(opts.logitCap))
-
-		// Apply combined causal + validity mask.
-		// causalMask [chunkSize * contextSize]: 1=causal-allowed, 0=masked.
-		// Validity: context positions before the actual sequence start are invalid.
-		// For chunk u, context position c corresponds to actual time: u*chunkSize + c - padLeft.
-		// Valid if 0 <= actual_time < seqLen.
-		// Mask tensor layout: [contextSize, chunkSize, 1] with ne[0]=contextSize contiguous.
-		// Element at (context=j, chunk=i) is at flat index: i*contextSize + j.
-		maskData := make([]float32, contextSize*chunkSize)
-		for i := range chunkSize {
-			for j := range contextSize {
-				actualTime := u*chunkSize + j - padLeft
-				causalOK := causalMask[i*contextSize+j] > 0
-				validOK := actualTime >= 0 && actualTime < seqLen
-				if causalOK && validOK {
-					maskData[i*contextSize+j] = 0
-				} else {
-					maskData[i*contextSize+j] = -1e9
-				}
-			}
-		}
-		mask := ctx.Input().FromFloats(maskData, contextSize, chunkSize, 1) // 3D for broadcasting over numHeads
-		logits = logits.Add(ctx, mask)
-
-		// Softmax over context dimension (dim 0 = contextSize).
-		logits = logits.Softmax(ctx) // softmax over ne[0]=contextSize
-
-		// Weighted sum: logits^T @ vCtx.
-		// logits: [contextSize, chunkSize, numHeads], vCtx: [headDim, numHeads, contextSize]
-		// vCtx permuted: [headDim, contextSize, numHeads]
-		vP := vCtx.Permute(ctx, 0, 2, 1, 3) // [headDim, contextSize, numHeads]
-		// Weighted sum: for each head, value[headDim, contextSize] @ weights[contextSize, chunkSize]
-		// = [headDim, chunkSize].
-		// Mulmat(a, b) = a^T @ b. Need a=[contextSize, headDim, numHeads], b=[contextSize, chunkSize, numHeads].
-		vPT := vP.Permute(ctx, 1, 0, 2, 3).Contiguous(ctx) // [contextSize, headDim, numHeads]
-		chunkOut := vPT.Mulmat(ctx, logits)                // [headDim, chunkSize, numHeads]
-
-		// Permute back to [headDim, numHeads, chunkSize]
-		chunkOut = chunkOut.Permute(ctx, 0, 2, 1, 3).Contiguous(ctx)
-		chunkOutputs[u] = chunkOut
-	}
-
-	// Concatenate chunk outputs along time dimension.
-	var attnOut ml.Tensor
-	if numChunks == 1 {
-		attnOut = chunkOutputs[0]
-	} else {
-		attnOut = chunkOutputs[0]
-		for _, co := range chunkOutputs[1:] {
-			attnOut = attnOut.Concat(ctx, co, 2)
-		}
-	}
-
-	// Trim to original sequence length if we padded.
-	if paddedLen > seqLen {
-		attnOut = attnOut.Slice(ctx, 2, 0, seqLen, 1).Contiguous(ctx)
-	}
-
-	// Reshape to [hiddenSize, seqLen] and project.
-	attnOut = attnOut.Reshape(ctx, hiddenSize, seqLen)
-	x = cb.AttnOut.Forward(ctx, attnOut)
-	x = x.Clamp(ctx, -opts.gradClip, opts.gradClip)
-	x = cb.AttnPostNorm.Forward(ctx, x, opts.eps)
-
-	return residual.Add(ctx, x)
-}
-
-// buildPositionEmbeddings builds sinusoidal position embeddings and projects through linear_pos.
-// Returns [headDim, numHeads, maxSpan] tensor.
-func (cb *AudioConformerBlock) buildPositionEmbeddings(ctx ml.Context, maxSpan int, opts *AudioModelOptions) ml.Tensor {
-	halfDim := opts.hiddenSize / 2
-	hiddenSize := opts.hiddenSize
-
-	// inv_timescales: exp(-i * log(10000) / max(D/2-1, 1))
-	logInc := math.Log(10000.0) / math.Max(float64(halfDim-1), 1)
-
-	// Sinusoidal embeddings for relative positions [maxPast, maxPast-1, ..., -maxFuture].
-	posData := make([]float32, hiddenSize*maxSpan)
-	for p := range maxSpan {
-		relPos := float64(opts.maxPast - p)
-		for d := range halfDim {
-			angle := relPos * math.Exp(float64(-d)*logInc)
-			posData[p*hiddenSize+d] = float32(math.Sin(angle))
-			posData[p*hiddenSize+halfDim+d] = float32(math.Cos(angle))
-		}
-	}
-
-	// Create [hiddenSize, maxSpan] input tensor.
-	posEmb := ctx.Input().FromFloats(posData, hiddenSize, maxSpan)
-
-	// Project through linear_pos: [hiddenSize, maxSpan] → Mulmat → [numHeads*headDim, maxSpan]
-	projPos := cb.LinearPos.Mulmat(ctx, posEmb)
-
-	// Reshape to [headDim, numHeads, maxSpan].
-	return projPos.Reshape(ctx, opts.headDim, opts.numHeads, maxSpan)
-}
-
-// relativeShiftGGML performs the relative shift to extract correct position logits.
-// Input: [maxSpan, chunkSize, numHeads]. Output: [contextSize, chunkSize, numHeads].
-func (cb *AudioConformerBlock) relativeShiftGGML(ctx ml.Context, x ml.Tensor, maxSpan, chunkSize, contextSize, numHeads int) ml.Tensor {
-	// The shift trick: pad ne[0] to contextSize+1, reshape to flatten first two dims,
-	// skip first (contextSize+1-maxSpan) elements, take contextSize*chunkSize elements, reshape back.
-	padAmt := contextSize + 1 - maxSpan
-	if padAmt > 0 {
-		x = x.Pad(ctx, padAmt, 0, 0, 0) // [maxSpan+padAmt, chunkSize, numHeads] = [contextSize+1, chunkSize, numHeads]
-	}
-	// Reshape to [(contextSize+1)*chunkSize, numHeads]
-	x = x.Reshape(ctx, (contextSize+1)*chunkSize, numHeads)
-	// Take the first contextSize*chunkSize elements (the standard relative shift trick).
-	x = x.Slice(ctx, 0, 0, contextSize*chunkSize, 1).Contiguous(ctx)
-	// Reshape to [contextSize, chunkSize, numHeads]
-	return x.Reshape(ctx, contextSize, chunkSize, numHeads)
-}
-
-// forwardLightConv runs the lightweight depthwise convolution module.
-func (cb *AudioConformerBlock) forwardLightConv(ctx ml.Context, x ml.Tensor, opts *AudioModelOptions, blockIdx int) ml.Tensor {
-	residual := x
-
-	x = cb.ConvNorm.Forward(ctx, x, opts.eps)
-	x = cb.ConvPW1.Forward(ctx, x) // [2*D, T, B]
-
-	// GLU: split in half along dim 0, sigmoid gate, multiply.
-	d := x.Dim(0) / 2
-	data := x.Slice(ctx, 0, 0, d, 1).Contiguous(ctx)
-	gate := x.Slice(ctx, 0, d, d*2, 1).Contiguous(ctx).Sigmoid(ctx)
-	x = data.Mul(ctx, gate) // [D, T, B]
-
-	// Depthwise Conv1d: manual implementation using model weight tensor slices.
-	// Kernel cb.ConvDW shape: [K=5, D=1024] (ne[0]=K, ne[1]=D) after shape reversal.
-	// Actually in GGML, ne[0]=K=5 contiguous, ne[1]=D=1024.
-	// We need per-tap weights [D] and shifted input copies.
-	kernelSize := cb.ConvDW.Dim(0) // K=5
-	seqLen := x.Dim(1)
-
-	// Transpose kernel to [D, K] for per-tap slicing.
-	// GGML permute(1,0,2,3): old[0]→pos1, old[1]→pos0 → swap ne[0] and ne[1]
-	kernelT := cb.ConvDW.Permute(ctx, 1, 0, 2, 3).Contiguous(ctx) // [D, K]
-
-	var convOut ml.Tensor
-	for k := range kernelSize {
-		shift := kernelSize - 1 - k
-		var shifted ml.Tensor
-		if shift == 0 {
-			shifted = x
-		} else {
-			trimmed := x.Slice(ctx, 1, 0, seqLen-shift, 1).Contiguous(ctx)
-			shifted = trimmed.PadExt(ctx, 0, 0, shift, 0, 0, 0, 0, 0)
-		}
-
-		wk := kernelT.Slice(ctx, 1, k, k+1, 1).Contiguous(ctx) // [D, 1]
-		term := shifted.Mul(ctx, wk)
-		if convOut == nil {
-			convOut = term
-		} else {
-			convOut = convOut.Add(ctx, term)
-		}
-	}
-	x = convOut
-
-	x = x.Clamp(ctx, -opts.gradClip, opts.gradClip)
-	x = cb.NormConv.Forward(ctx, x, opts.eps)
-	x = x.SILU(ctx)
-	x = cb.ConvPW2.Forward(ctx, x)
-
-	return x.Add(ctx, residual)
-}
-
-func newAudioModel(c fs.Config) *AudioModel {
-	numLayers := int(c.Uint("audio.block_count", 0))
-	if numLayers == 0 {
-		return nil
-	}
-	return &AudioModel{
-		Layers: make([]AudioConformerBlock, numLayers),
-	}
-}
-
-func newAudioModelOptions(c fs.Config) *AudioModelOptions {
-	hiddenSize := int(c.Uint("audio.embedding_length", 0))
-	if hiddenSize == 0 {
-		return nil
-	}
-	numHeads := int(c.Uint("audio.attention.head_count", 8))
-	headDim := hiddenSize / numHeads
-	chunkSize := 12 // default conformer chunk size
-	maxPast := 12   // conf_attention_context_left - 1
-	maxFuture := 0  // conf_attention_context_right
-	convKernel := int(c.Uint("audio.conv_kernel_size", 5))
-
-	eps := c.Float("audio.attention.layer_norm_epsilon", 1e-6)
-
-	return &AudioModelOptions{
-		hiddenSize:     hiddenSize,
-		numHeads:       numHeads,
-		headDim:        headDim,
-		ffnSize:        int(c.Uint("audio.feed_forward_length", uint32(hiddenSize*4))),
-		numLayers:      int(c.Uint("audio.block_count", 12)),
-		melBins:        int(c.Uint("audio.num_mel_bins", 128)),
-		chunkSize:      chunkSize,
-		maxPast:        maxPast,
-		maxFuture:      maxFuture,
-		contextSize:    chunkSize + maxPast + maxFuture,
-		logitCap:       50.0,
-		residualWeight: 0.5,
-		gradClip:       1e10,
-		convKernelSize: convKernel,
-		eps:            float32(eps),
-	}
-}
-
-// buildCausalValidMaskF32 creates the causal-valid mask for block-local attention.
-// Returns flat [chunkSize * contextSize] float32 data (1.0 = allowed, 0.0 = masked).
-func buildCausalValidMaskF32(chunkSize, maxPast, maxFuture int) []float32 {
-	contextSize := chunkSize + maxPast + maxFuture
-	upperDiag := maxPast + maxFuture
-
-	result := make([]float32, chunkSize*contextSize)
-	for r := range chunkSize {
-		for c := range contextSize {
-			lower := (r <= c)           // tril(contextSize, chunkSize) transposed
-			upper := (c <= r+upperDiag) // tril(chunkSize, contextSize, diag=upperDiag)
-			if lower && upper {
-				result[r*contextSize+c] = 1.0
-			}
-		}
-	}
-	return result
-}

model/models/gemma4/model_text.go

@@ -1,475 +0,0 @@
-package gemma4
-
-import (
-	"math"
-
-	"github.com/ollama/ollama/fs"
-	"github.com/ollama/ollama/kvcache"
-	"github.com/ollama/ollama/ml"
-	"github.com/ollama/ollama/ml/nn"
-	"github.com/ollama/ollama/ml/nn/rope"
-	"github.com/ollama/ollama/model/input"
-)
-
-const (
-	cacheTypeSWA = iota
-	cacheTypeCausal
-)
-
-type TextOptions struct {
-	hiddenSize              int
-	numHeads, numKVHeads    int
-	numGlobalKVHeads        int
-	headDim, globalHeadDim  int
-	hiddenLayers            int
-	hiddenSizePerLayerInput int
-
-	eps               float32
-	ropeBase          float32
-	ropeLocalBase     float32
-	partialRotaryDims int // RoPE dims for full-attention (global) layers
-
-	slidingWindowPattern []bool
-	// kvDonorMap maps shared layer index -> donor layer index.
-	// Donor is the last non-shared layer of the same type (sliding/full).
-	kvDonorMap map[int]int
-
-	finalLogitSoftcap float32
-
-	numExperts     int
-	numExpertsUsed int
-}
-
-func (o *TextOptions) isLocal(layer int) bool {
-	if layer < len(o.slidingWindowPattern) {
-		return o.slidingWindowPattern[layer]
-	}
-	return false
-}
-
-func (o *TextOptions) ropeForLayer(layer int) (base float32, dims int) {
-	if o.isLocal(layer) {
-		return o.ropeLocalBase, o.headDim
-	}
-	return o.ropeBase, o.partialRotaryDims
-}
-
-func (o *TextOptions) kvHeadsForLayer(layer int) int {
-	if o.isLocal(layer) {
-		return o.numKVHeads
-	}
-	if o.numGlobalKVHeads > 0 {
-		return o.numGlobalKVHeads
-	}
-	return o.numKVHeads
-}
-
-func (o *TextOptions) headDimForLayer(layer int) int {
-	if o.isLocal(layer) {
-		return o.headDim
-	}
-	return o.globalHeadDim
-}
-
-type TextModel struct {
-	TokenEmbedding *nn.Embedding `gguf:"token_embd"`
-	*PerLayerProjector
-	Layers     []TextLayer `gguf:"blk"`
-	OutputNorm *nn.RMSNorm `gguf:"output_norm"`
-	Output     *nn.Linear  `gguf:"output,alt:token_embd"`
-	TextOptions
-}
-
-func newTextModel(c fs.Config) *TextModel {
-	numLayers := int(c.Uint("block_count"))
-
-	// Head dimensions: key_length is global head dim, key_length_swa is local (SWA) head dim.
-	globalHeadDim := int(c.Uint("attention.key_length", 512))
-	headDim := int(c.Uint("attention.key_length_swa", 256))
-
-	// RoPE dimensions for global (full attention) layers with proportional RoPE.
-	// The freq_factors tensor handles partial rotation (1.0 for rotated pairs,
-	// 1e30 for non-rotated), so ropeDims equals the full global head dim.
-	partialRotaryDims := int(c.Uint("rope.dimension_count", 0))
-	if partialRotaryDims == 0 {
-		partialFactor := c.Float("rope.partial_rotary_factor", 1.0)
-		partialRotaryDims = int(float32(globalHeadDim) * partialFactor)
-	}
-
-	ropeBase := c.Float("rope.freq_base", 1000000.0)
-	ropeLocalBase := c.Float("rope.freq_base_swa", 0)
-	if ropeLocalBase == 0 {
-		ropeLocalBase = c.Float("rope.local.freq_base", 10000.0)
-	}
-
-	numGlobalKVHeads := int(c.Uint("attention.global_head_count_kv", 0))
-	slidingPattern := c.Bools("attention.sliding_window_pattern")
-
-	// KV heads: try per-layer array first (MoE models), then fall back to scalar
-	numKVHeads := 0
-	kvHeadsArray := c.Ints("attention.head_count_kv")
-	if len(kvHeadsArray) > 0 {
-		numKVHeads = int(kvHeadsArray[0])
-		if numGlobalKVHeads == 0 && len(slidingPattern) > 0 {
-			for i, isLocal := range slidingPattern {
-				if !isLocal && i < len(kvHeadsArray) {
-					numGlobalKVHeads = int(kvHeadsArray[i])
-					break
-				}
-			}
-		}
-	}
-	if numKVHeads == 0 {
-		numKVHeads = int(c.Uint("attention.head_count_kv", 0))
-	}
-
-	// Compute KV sharing donor map (same logic as MLX)
-	sharedLayers := int(c.Uint("attention.shared_kv_layers", 0))
-	kvDonorMap := make(map[int]int)
-	if sharedLayers > 0 && len(slidingPattern) > 0 {
-		firstShared := numLayers - sharedLayers
-		for i := firstShared; i < numLayers; i++ {
-			isLocal := slidingPattern[i]
-			// Find last non-shared layer of same type
-			for j := firstShared - 1; j >= 0; j-- {
-				if slidingPattern[j] == isLocal {
-					kvDonorMap[i] = j
-					break
-				}
-			}
-		}
-	}
-
-	return &TextModel{
-		Layers: make([]TextLayer, numLayers),
-		TextOptions: TextOptions{
-			hiddenSize:              int(c.Uint("embedding_length")),
-			numHeads:                int(c.Uint("attention.head_count")),
-			numKVHeads:              numKVHeads,
-			numGlobalKVHeads:        numGlobalKVHeads,
-			headDim:                 headDim,
-			globalHeadDim:           globalHeadDim,
-			hiddenLayers:            numLayers,
-			hiddenSizePerLayerInput: int(c.Uint("embedding_length_per_layer_input", 0)),
-			eps:                     c.Float("attention.layer_norm_rms_epsilon", 1e-06),
-			ropeBase:                ropeBase,
-			ropeLocalBase:           ropeLocalBase,
-			partialRotaryDims:       partialRotaryDims,
-			slidingWindowPattern:    slidingPattern,
-			kvDonorMap:              kvDonorMap,
-			finalLogitSoftcap:       c.Float("final_logit_softcapping", 0.0),
-			numExperts:              int(c.Uint("expert_count", 0)),
-			numExpertsUsed:          int(c.Uint("expert_used_count", 0)),
-		},
-	}
-}
-
-func (m *TextModel) Forward(ctx ml.Context, batch input.Batch, cache kvcache.Cache) ml.Tensor {
-	positions := ctx.Input().FromInts(batch.Positions, len(batch.Positions))
-
-	hiddenState := m.TokenEmbedding.Forward(ctx, batch.Inputs)
-	hiddenState = hiddenState.Scale(ctx, math.Sqrt(float64(m.hiddenSize)))
-
-	// Inject vision embeddings into the hidden state
-	var except []int
-	for _, image := range batch.Multimodal {
-		visionOutputs := image.Multimodal[0].Tensor
-		ctx.Forward(visionOutputs.Copy(ctx, hiddenState.View(ctx, image.Index*hiddenState.Stride(1), visionOutputs.Dim(0)*visionOutputs.Dim(1))))
-
-		for i := range visionOutputs.Dim(1) {
-			except = append(except, image.Index+i)
-		}
-	}
-
-	// PLE
-	var perLayerInputs ml.Tensor
-	if m.PerLayerProjector != nil {
-		perLayerInputs = m.PerLayerProjector.Forward(ctx, batch, hiddenState, &m.TextOptions)
-	}
-
-	for i := range len(m.Layers) {
-		layer := m.Layers[i]
-		if cache != nil {
-			cache.SetLayer(i)
-			cacheType := cacheTypeSWA
-			if !m.isLocal(i) {
-				cacheType = cacheTypeCausal
-			}
-			wc := cache.(*kvcache.WrapperCache)
-			wc.SetLayerType(cacheType)
-
-			if causal, ok := wc.UnderlyingCache().(*kvcache.Causal); ok {
-				causal.SetCausal(ctx, kvcache.CausalOptions{Except: except})
-			}
-		}
-
-		var lastLayerOutputs ml.Tensor
-		if i == len(m.Layers)-1 {
-			lastLayerOutputs = batch.Outputs
-		}
-
-		var perLayerInput ml.Tensor
-		if perLayerInputs != nil {
-			perLayerInput = perLayerInputs.View(ctx, i*perLayerInputs.Stride(1), perLayerInputs.Dim(0), perLayerInputs.Stride(2), perLayerInputs.Dim(2))
-		}
-
-		// KV sharing: layers >= firstShared reuse K/V from donor layers
-		isShared := false
-		if donorLayer, ok := m.kvDonorMap[i]; ok {
-			// Set cache layer to donor so Get() reads donor's K/V
-			cache.SetLayer(donorLayer)
-			isShared = true
-		}
-		hiddenState = layer.Forward(ctx, i, hiddenState, positions, perLayerInput, lastLayerOutputs, cache, isShared, &m.TextOptions)
-	}
-
-	return m.OutputNorm.Forward(ctx, hiddenState, m.eps)
-}
-
-// PerLayerProjector implements PLE.
-type PerLayerProjector struct {
-	TokenEmbedding *nn.Embedding `gguf:"per_layer_token_embd"`
-	Projector      *nn.Linear    `gguf:"per_layer_model_proj"`
-	Norm           *nn.RMSNorm   `gguf:"per_layer_proj_norm"`
-}
-
-func (p *PerLayerProjector) Forward(ctx ml.Context, batch input.Batch, inputs ml.Tensor, opts *TextOptions) ml.Tensor {
-	inputsPerLayer := p.TokenEmbedding.Forward(ctx, batch.Inputs)
-	inputsPerLayer = inputsPerLayer.Scale(ctx, math.Sqrt(float64(opts.hiddenSizePerLayerInput)))
-	// Reshape to [pleDim, numLayers, numTokens] — matching projection shape
-	inputsPerLayer = inputsPerLayer.Reshape(ctx, opts.hiddenSizePerLayerInput, opts.hiddenLayers, inputs.Dim(1))
-
-	perLayerProjection := p.Projector.Forward(ctx, inputs)
-	perLayerProjection = perLayerProjection.Scale(ctx, 1.0/math.Sqrt(float64(opts.hiddenSize)))
-	perLayerProjection = perLayerProjection.Reshape(ctx, opts.hiddenSizePerLayerInput, opts.hiddenLayers, inputs.Dim(1))
-	perLayerProjection = p.Norm.Forward(ctx, perLayerProjection, opts.eps)
-
-	if inputsPerLayer != nil {
-		perLayerProjection = perLayerProjection.Add(ctx, inputsPerLayer)
-		perLayerProjection = perLayerProjection.Scale(ctx, 1/math.Sqrt(2))
-	}
-
-	return perLayerProjection
-}
-
-type TextSelfAttention struct {
-	Query       *nn.Linear  `gguf:"attn_q"`
-	QueryNorm   *nn.RMSNorm `gguf:"attn_q_norm"`
-	Key         *nn.Linear  `gguf:"attn_k"`
-	KeyNorm     *nn.RMSNorm `gguf:"attn_k_norm"`
-	Value       *nn.Linear  `gguf:"attn_v"`
-	Output      *nn.Linear  `gguf:"attn_output"`
-	RopeFactors ml.Tensor   `gguf:"rope_freqs.weight"` // proportional RoPE freq_factors
-}
-
-func (sa *TextSelfAttention) Forward(ctx ml.Context, layer int, hiddenState, positions ml.Tensor, cache kvcache.Cache, sharedKV bool, opts *TextOptions) ml.Tensor {
-	batchSize := hiddenState.Dim(1)
-	hd := opts.headDimForLayer(layer)
-	kvHeads := opts.kvHeadsForLayer(layer)
-	ropeBase, ropeDims := opts.ropeForLayer(layer)
-
-	q := sa.Query.Forward(ctx, hiddenState)
-	q = q.Reshape(ctx, hd, opts.numHeads, batchSize)
-	q = sa.QueryNorm.Forward(ctx, q, opts.eps)
-
-	var k, v ml.Tensor
-	if !sharedKV {
-		k = sa.Key.Forward(ctx, hiddenState)
-		k = k.Reshape(ctx, hd, kvHeads, batchSize)
-
-		if sa.Value != nil {
-			v = sa.Value.Forward(ctx, hiddenState)
-			v = v.Reshape(ctx, hd, kvHeads, batchSize)
-		} else {
-			// K=V: use raw K projection (before K norm) as V
-			v = k
-		}
-
-		k = sa.KeyNorm.Forward(ctx, k, opts.eps)
-		v = v.RMSNorm(ctx, nil, opts.eps) // V norm: unweighted RMSNorm
-	}
-
-	// RoPE with proportional freq_factors on global layers
-	ropeOpts := []func(*rope.Options){rope.WithTypeNeoX()}
-	if sa.RopeFactors != nil && !opts.isLocal(layer) {
-		ropeOpts = append(ropeOpts, rope.WithFactors(sa.RopeFactors))
-	}
-	q = nn.RoPE(ctx, q, positions, ropeDims, ropeBase, 1.0, ropeOpts...)
-	if k != nil {
-		k = nn.RoPE(ctx, k, positions, ropeDims, ropeBase, 1.0, ropeOpts...)
-	}
-
-	attention := nn.Attention(ctx, q, k, v, 1.0, cache)
-
-	attention = attention.Reshape(ctx, hd*opts.numHeads, batchSize)
-	return sa.Output.Forward(ctx, attention)
-}
-
-type TextMLP struct {
-	Gate *nn.Linear `gguf:"ffn_gate"`
-	Up   *nn.Linear `gguf:"ffn_up"`
-	Down *nn.Linear `gguf:"ffn_down"`
-}
-
-func (mlp *TextMLP) Forward(ctx ml.Context, hiddenState ml.Tensor) ml.Tensor {
-	hiddenState = mlp.Gate.Forward(ctx, hiddenState).GELU(ctx, mlp.Up.Forward(ctx, hiddenState))
-	return mlp.Down.Forward(ctx, hiddenState)
-}
-
-// TextRouter implements the Gemma 4 MoE router.
-type TextRouter struct {
-	Proj  *nn.Linear `gguf:"ffn_gate_inp"`
-	Scale ml.Tensor  `gguf:"ffn_gate_inp.scale"`
-}
-
-func (r *TextRouter) Forward(ctx ml.Context, hiddenState ml.Tensor, opts *TextOptions) (routingWeights, selectedExperts ml.Tensor) {
-	// RMSNorm without learned weight
-	x := hiddenState.RMSNorm(ctx, nil, opts.eps)
-	// Scale by 1/sqrt(hidden_size)
-	x = x.Scale(ctx, 1.0/math.Sqrt(float64(opts.hiddenSize)))
-	// Multiply by learned scale parameter
-	x = x.Mul(ctx, r.Scale)
-	// Project to expert logits
-	expertScores := r.Proj.Forward(ctx, x)
-	// Softmax over experts
-	routingWeights = expertScores.Softmax(ctx)
-	// TopK expert selection
-	selectedExperts = routingWeights.TopK(ctx, opts.numExpertsUsed)
-	return routingWeights, selectedExperts
-}
-
-// TextMoEBlock implements the Gemma 4 sparse MoE.
-type TextMoEBlock struct {
-	GateUp    *nn.LinearBatch `gguf:"ffn_gate_up_exps"`
-	Gate      *nn.LinearBatch `gguf:"ffn_gate_exps"`
-	Up        *nn.LinearBatch `gguf:"ffn_up_exps"`
-	Down      *nn.LinearBatch `gguf:"ffn_down_exps"`
-	DownScale ml.Tensor       `gguf:"ffn_down_exps.scale,alt:ffn_gate_inp.per_expert_scale"`
-}
-
-func (moe *TextMoEBlock) Forward(ctx ml.Context, hiddenState, routingWeights, selectedExperts ml.Tensor, opts *TextOptions) ml.Tensor {
-	// Select routing weights for chosen experts and renormalize
-	routingWeights = routingWeights.Reshape(ctx, 1, opts.numExperts, hiddenState.Dim(1)).Rows(ctx, selectedExperts)
-	routingWeights = routingWeights.Reshape(ctx, opts.numExpertsUsed, hiddenState.Dim(1))
-	routingWeights = routingWeights.Div(ctx, routingWeights.SumRows(ctx))
-	routingWeights = routingWeights.Reshape(ctx, 1, opts.numExpertsUsed, hiddenState.Dim(1))
-
-	hiddenState = hiddenState.Reshape(ctx, hiddenState.Dim(0), 1, hiddenState.Dim(1))
-
-	// Expert computation using LinearBatch (MulmatID selecting experts by index)
-	var gateOut, upOut ml.Tensor
-	if moe.GateUp != nil && moe.GateUp.Weight != nil {
-		gateUp := moe.GateUp.Forward(ctx, hiddenState, selectedExperts)
-		nFF := gateUp.Dim(0) / 2
-		gateOut = gateUp.Slice(ctx, 0, 0, nFF, 1)
-		upOut = gateUp.Slice(ctx, 0, nFF, gateUp.Dim(0), 1)
-	} else {
-		gateOut = moe.Gate.Forward(ctx, hiddenState, selectedExperts)
-		upOut = moe.Up.Forward(ctx, hiddenState, selectedExperts)
-	}
-	hiddenState = gateOut.GELU(ctx, upOut)
-	experts := moe.Down.Forward(ctx, hiddenState, selectedExperts)
-
-	// Apply per-expert down projection scale when present.
-	if moe.DownScale != nil {
-		expertScales := moe.DownScale.Reshape(ctx, opts.numExperts, 1)
-		expertScales = expertScales.Repeat(ctx, 1, hiddenState.Dim(2))
-		expertScales = expertScales.Reshape(ctx, 1, opts.numExperts, hiddenState.Dim(2)).Rows(ctx, selectedExperts)
-		expertScales = expertScales.Reshape(ctx, opts.numExpertsUsed, hiddenState.Dim(2))
-		expertScales = expertScales.Reshape(ctx, 1, opts.numExpertsUsed, hiddenState.Dim(2))
-		experts = experts.Mul(ctx, expertScales)
-	}
-
-	// Apply routing weights
-	experts = experts.Mul(ctx, routingWeights)
-
-	// Sum across experts
-	nextStates := experts.View(ctx, 0, experts.Dim(0), experts.Stride(2), experts.Dim(2))
-	for i := 1; i < opts.numExpertsUsed; i++ {
-		nextStates = nextStates.Add(ctx, experts.View(ctx, i*experts.Stride(1), experts.Dim(0), experts.Stride(2), experts.Dim(2)))
-	}
-
-	return nextStates
-}
-
-type TextLayer struct {
-	AttentionNorm     *nn.RMSNorm `gguf:"attn_norm"`
-	SelfAttention     *TextSelfAttention
-	PostAttentionNorm *nn.RMSNorm `gguf:"post_attention_norm,alt:attn_post_norm"`
-	MLPNorm           *nn.RMSNorm `gguf:"ffn_norm,alt:ffn_pre_norm"`
-	MLP               *TextMLP
-	PostMLPNorm       *nn.RMSNorm `gguf:"post_ffw_norm,alt:ffn_post_norm"`
-
-	// MoE (present only for models with enable_moe_block=true)
-	Router       *TextRouter
-	MoE          *TextMoEBlock
-	MoENorm      *nn.RMSNorm `gguf:"pre_ffw_norm_2,alt:ffn_pre_norm_2"`
-	PostMoENorm  *nn.RMSNorm `gguf:"post_ffw_norm_2,alt:ffn_post_norm_2"`
-	PostMLPNorm1 *nn.RMSNorm `gguf:"post_ffw_norm_1,alt:ffn_post_norm_1"` // used instead of PostMLPNorm when MoE is present
-
-	PerLayerInputGate  *nn.Linear  `gguf:"inp_gate"`
-	PerLayerProjection *nn.Linear  `gguf:"proj"`
-	PostPerLayerNorm   *nn.RMSNorm `gguf:"post_norm"`
-	LayerScalar        ml.Tensor   `gguf:"layer_scalar,alt:layer_output_scale.weight"`
-}
-
-func (l *TextLayer) Forward(ctx ml.Context, layer int, hiddenState, positions, perLayerInput, outputs ml.Tensor, cache kvcache.Cache, sharedKV bool, opts *TextOptions) ml.Tensor {
-	residual := hiddenState
-
-	hiddenState = l.AttentionNorm.Forward(ctx, hiddenState, opts.eps)
-	hiddenState = l.SelfAttention.Forward(ctx, layer, hiddenState, positions, cache, sharedKV, opts)
-	hiddenState = l.PostAttentionNorm.Forward(ctx, hiddenState, opts.eps)
-
-	if outputs != nil {
-		hiddenState = hiddenState.Rows(ctx, outputs)
-		residual = residual.Rows(ctx, outputs)
-		if perLayerInput != nil {
-			perLayerInput = perLayerInput.Rows(ctx, outputs)
-		}
-	}
-
-	hiddenState = hiddenState.Add(ctx, residual)
-	residual = hiddenState
-
-	// MLP (+ optional MoE in parallel)
-	hasSplitExperts := l.MoE != nil && l.MoE.Gate != nil && l.MoE.Up != nil && l.MoE.Gate.Weight != nil && l.MoE.Up.Weight != nil
-	hasFusedExperts := l.MoE != nil && l.MoE.GateUp != nil && l.MoE.GateUp.Weight != nil
-	if l.Router != nil && l.MoE != nil && l.MoE.Down != nil && l.MoE.Down.Weight != nil && (hasSplitExperts || hasFusedExperts) {
-		// MoE layers: run MLP and MoE in parallel, sum results
-		mlpState := l.MLPNorm.Forward(ctx, hiddenState, opts.eps)
-		mlpState = l.MLP.Forward(ctx, mlpState)
-		mlpState = l.PostMLPNorm1.Forward(ctx, mlpState, opts.eps)
-
-		routingWeights, selectedExperts := l.Router.Forward(ctx, hiddenState, opts)
-		moeState := l.MoENorm.Forward(ctx, hiddenState, opts.eps)
-		moeState = l.MoE.Forward(ctx, moeState, routingWeights, selectedExperts, opts)
-		moeState = l.PostMoENorm.Forward(ctx, moeState, opts.eps)
-
-		// Combine MLP + MoE, apply outer post-FFN norm, then add residual
-		combined := mlpState.Add(ctx, moeState)
-		combined = l.PostMLPNorm.Forward(ctx, combined, opts.eps)
-		hiddenState = combined.Add(ctx, residual)
-	} else {
-		// Dense layers: MLP only
-		hiddenState = l.MLPNorm.Forward(ctx, hiddenState, opts.eps)
-		hiddenState = l.MLP.Forward(ctx, hiddenState)
-		hiddenState = l.PostMLPNorm.Forward(ctx, hiddenState, opts.eps)
-		hiddenState = hiddenState.Add(ctx, residual)
-	}
-
-	// PLE injection (after MLP residual)
-	if perLayerInput != nil && l.PerLayerInputGate != nil {
-		pleState := l.PerLayerInputGate.Forward(ctx, hiddenState)
-		pleState = pleState.GELU(ctx, perLayerInput)
-		pleState = l.PerLayerProjection.Forward(ctx, pleState)
-		pleState = l.PostPerLayerNorm.Forward(ctx, pleState, opts.eps)
-		hiddenState = hiddenState.Add(ctx, pleState)
-	}
-
-	// Layer scalar applied at end of layer (full-attention layers only)
-	if l.LayerScalar != nil {
-		hiddenState = hiddenState.Mul(ctx, l.LayerScalar)
-	}
-
-	return hiddenState
-}

model/models/gemma4/model_vision.go

@@ -1,384 +0,0 @@
-package gemma4
-
-import (
-	"math"
-
-	"github.com/ollama/ollama/fs"
-	"github.com/ollama/ollama/ml"
-	"github.com/ollama/ollama/ml/nn"
-	"github.com/ollama/ollama/ml/nn/rope"
-)
-
-const batchSize = 1
-
-// ClippableLinear is a linear layer with optional input/output clamping.
-// Required by Gemma4 vision encoder for numerical stability with F16 weights.
-type ClippableLinear struct {
-	Weight ml.Tensor `gguf:"weight"`
-
-	InputMin  ml.Tensor `gguf:"input_min"`
-	InputMax  ml.Tensor `gguf:"input_max"`
-	OutputMin ml.Tensor `gguf:"output_min"`
-	OutputMax ml.Tensor `gguf:"output_max"`
-
-	inMin, inMax, outMin, outMax float32
-	hasClamp                     bool
-	clampsLoaded                 bool
-}
-
-func scalarValue(t ml.Tensor) (float32, bool) {
-	if t == nil {
-		return 0, false
-	}
-
-	data := t.BackendGet()
-	if len(data) == 0 {
-		return 0, false
-	}
-
-	return data[0], true
-}
-
-func (l *ClippableLinear) loadClampFromScalars() {
-	if l.clampsLoaded {
-		return
-	}
-	l.clampsLoaded = true
-
-	const (
-		defaultMin = -math.MaxFloat32
-		defaultMax = math.MaxFloat32
-	)
-
-	inMin, hasInMin := scalarValue(l.InputMin)
-	inMax, hasInMax := scalarValue(l.InputMax)
-	outMin, hasOutMin := scalarValue(l.OutputMin)
-	outMax, hasOutMax := scalarValue(l.OutputMax)
-
-	if !(hasInMin || hasInMax || hasOutMin || hasOutMax) {
-		return
-	}
-
-	l.hasClamp = true
-	l.inMin = defaultMin
-	l.inMax = defaultMax
-	l.outMin = defaultMin
-	l.outMax = defaultMax
-
-	if hasInMin {
-		l.inMin = inMin
-	}
-	if hasInMax {
-		l.inMax = inMax
-	}
-	if hasOutMin {
-		l.outMin = outMin
-	}
-	if hasOutMax {
-		l.outMax = outMax
-	}
-}
-
-func (l *ClippableLinear) Forward(ctx ml.Context, x ml.Tensor) ml.Tensor {
-	if l.hasClamp {
-		x = x.Clamp(ctx, l.inMin, l.inMax)
-	}
-	out := l.Weight.Mulmat(ctx, x)
-	if l.hasClamp {
-		out = out.Clamp(ctx, l.outMin, l.outMax)
-	}
-	return out
-}
-
-// InitClamp distributes packed clamp values from v.clamp_data to ClippableLinear structs.
-// If scalar clamp tensors (input_min/max, output_min/max) are present, they are used too.
-// Layout: numLayers × 7 linears (q,k,v,out,gate,up,down) × 4 floats (inMin,inMax,outMin,outMax)
-// then 4 floats for the projector.
-func (m *VisionModel) InitClamp(proj *MultiModalProjector) {
-	if m.clampInitDone {
-		return
-	}
-	m.clampInitDone = true
-
-	linears := func(l *VisionEncoderLayer) []*ClippableLinear {
-		return []*ClippableLinear{
-			l.SelfAttention.Query, l.SelfAttention.Key, l.SelfAttention.Value,
-			l.SelfAttention.Output, l.MLP.Gate, l.MLP.Up, l.MLP.Down,
-		}
-	}
-
-	for i := range m.Layers {
-		for _, cl := range linears(&m.Layers[i]) {
-			if cl != nil {
-				cl.loadClampFromScalars()
-			}
-		}
-	}
-	if proj != nil && proj.Projection != nil {
-		proj.Projection.loadClampFromScalars()
-	}
-
-	// Load packed clamp data when present (legacy Ollama format).
-	if m.ClampData == nil {
-		return
-	}
-
-	// Read all clamp values from packed F32 tensor
-	data := m.ClampData.BackendGet()
-	if len(data) == 0 {
-		return
-	}
-
-	// Distribute to layer linears: 7 per layer × 4 values each
-	for i := range m.Layers {
-		for li, cl := range linears(&m.Layers[i]) {
-			if cl == nil {
-				continue
-			}
-			idx := (i*7 + li) * 4
-			if idx+3 < len(data) {
-				cl.inMin = data[idx]
-				cl.inMax = data[idx+1]
-				cl.outMin = data[idx+2]
-				cl.outMax = data[idx+3]
-				cl.hasClamp = true
-			}
-		}
-	}
-
-	// Projector clamp values (last 4 floats)
-	if proj != nil && proj.Projection != nil {
-		projIdx := len(m.Layers) * 7 * 4
-		if projIdx+3 < len(data) {
-			proj.Projection.inMin = data[projIdx]
-			proj.Projection.inMax = data[projIdx+1]
-			proj.Projection.outMin = data[projIdx+2]
-			proj.Projection.outMax = data[projIdx+3]
-			proj.Projection.hasClamp = true
-		}
-	}
-}
-
-type VisionSelfAttention struct {
-	Query     *ClippableLinear `gguf:"attn_q"`
-	Key       *ClippableLinear `gguf:"attn_k"`
-	Value     *ClippableLinear `gguf:"attn_v"`
-	QueryNorm *nn.RMSNorm      `gguf:"attn_q_norm"`
-	KeyNorm   *nn.RMSNorm      `gguf:"attn_k_norm"`
-	Output    *ClippableLinear `gguf:"attn_out"`
-}
-
-func (sa *VisionSelfAttention) Forward(ctx ml.Context, hiddenState, posX, posY, attnMask ml.Tensor, opts *VisionModelOptions) ml.Tensor {
-	numPatches := hiddenState.Dim(1)
-	headDim := opts.hiddenSize / opts.numHeads
-
-	query := sa.Query.Forward(ctx, hiddenState)
-	key := sa.Key.Forward(ctx, hiddenState)
-	value := sa.Value.Forward(ctx, hiddenState)
-
-	query = query.Reshape(ctx, headDim, opts.numHeads, numPatches, batchSize)
-	key = key.Reshape(ctx, headDim, opts.numHeads, numPatches, batchSize)
-	value = value.Reshape(ctx, headDim, opts.numHeads, numPatches, batchSize)
-
-	// Q/K norms (Gemma-style: x * (1 + weight) / rms(x))
-	query = sa.QueryNorm.Forward(ctx, query, opts.eps)
-	key = sa.KeyNorm.Forward(ctx, key, opts.eps)
-
-	// V norm (RMSNorm without learned weights)
-	value = value.RMSNorm(ctx, nil, opts.eps)
-
-	// 2D RoPE: split head dim in half, apply NeoX RoPE with x positions to first half,
-	// y positions to second half, then concatenate.
-	halfDim := headDim / 2
-	ropeOpts := rope.WithTypeNeoX()
-
-	qFirst := query.View(ctx, 0, halfDim, query.Stride(1), opts.numHeads, query.Stride(2), numPatches)
-	qFirst = nn.RoPE(ctx, qFirst, posX, halfDim, opts.ropeTheta, 1.0, ropeOpts)
-
-	kFirst := key.View(ctx, 0, halfDim, key.Stride(1), opts.numHeads, key.Stride(2), numPatches)
-	kFirst = nn.RoPE(ctx, kFirst, posX, halfDim, opts.ropeTheta, 1.0, ropeOpts)
-
-	halfOffset := halfDim * query.Stride(0)
-	qSecond := query.View(ctx, halfOffset, halfDim, query.Stride(1), opts.numHeads, query.Stride(2), numPatches)
-	qSecond = nn.RoPE(ctx, qSecond, posY, halfDim, opts.ropeTheta, 1.0, ropeOpts)
-
-	halfOffsetK := halfDim * key.Stride(0)
-	kSecond := key.View(ctx, halfOffsetK, halfDim, key.Stride(1), opts.numHeads, key.Stride(2), numPatches)
-	kSecond = nn.RoPE(ctx, kSecond, posY, halfDim, opts.ropeTheta, 1.0, ropeOpts)
-
-	query = qFirst.Concat(ctx, qSecond, 0)
-	key = kFirst.Concat(ctx, kSecond, 0)
-
-	// Use flash attention for numerical stability (handles large attention scores
-	// from unclamped RMSNorm weights, e.g. 26B has addOne weights up to 19.5)
-	attention := nn.Attention(ctx, query, key, value, 1.0, nil)
-	attention = attention.Reshape(ctx, opts.hiddenSize, attention.Dim(2), batchSize)
-
-	return sa.Output.Forward(ctx, attention)
-}
-
-type VisionMLP struct {
-	Gate *ClippableLinear `gguf:"ffn_gate"`
-	Up   *ClippableLinear `gguf:"ffn_up"`
-	Down *ClippableLinear `gguf:"ffn_down"`
-}
-
-func (mlp *VisionMLP) Forward(ctx ml.Context, hiddenState ml.Tensor) ml.Tensor {
-	gate := mlp.Gate.Forward(ctx, hiddenState)
-	up := mlp.Up.Forward(ctx, hiddenState)
-	hiddenState = gate.QuickGELU(ctx, up)
-	return mlp.Down.Forward(ctx, hiddenState)
-}
-
-type VisionEncoderLayer struct {
-	AttentionNorm     *nn.RMSNorm `gguf:"ln1"`
-	SelfAttention     *VisionSelfAttention
-	PostAttentionNorm *nn.RMSNorm `gguf:"attn_post_norm"`
-
-	FFNNorm     *nn.RMSNorm `gguf:"ln2"`
-	MLP         *VisionMLP
-	PostFFNNorm *nn.RMSNorm `gguf:"ffn_post_norm"`
-
-	LayerOutputScale ml.Tensor `gguf:"out_scale.weight"`
-}
-
-func (e *VisionEncoderLayer) Forward(ctx ml.Context, hiddenState, posX, posY, attnMask ml.Tensor, opts *VisionModelOptions) ml.Tensor {
-	residual := hiddenState
-
-	// Pre-attention norm -> self attention -> post-attention norm
-	hiddenState = e.AttentionNorm.Forward(ctx, hiddenState, opts.eps)
-	hiddenState = e.SelfAttention.Forward(ctx, hiddenState, posX, posY, attnMask, opts)
-	hiddenState = e.PostAttentionNorm.Forward(ctx, hiddenState, opts.eps)
-
-	// Residual connection
-	hiddenState = hiddenState.Add(ctx, residual)
-	residual = hiddenState
-
-	// Pre-FFN norm -> FFN -> post-FFN norm
-	hiddenState = e.FFNNorm.Forward(ctx, hiddenState, opts.eps)
-	hiddenState = e.MLP.Forward(ctx, hiddenState)
-	hiddenState = e.PostFFNNorm.Forward(ctx, hiddenState, opts.eps)
-
-	// Residual connection
-	hiddenState = hiddenState.Add(ctx, residual)
-
-	// Per-layer output scale
-	if e.LayerOutputScale != nil {
-		hiddenState = hiddenState.Mul(ctx, e.LayerOutputScale)
-	}
-
-	return hiddenState
-}
-
-type VisionModelOptions struct {
-	hiddenSize int
-	numHeads   int
-	patchSize  int
-	nMerge     int
-	eps        float32
-	ropeTheta  float32
-}
-
-type VisionModel struct {
-	PatchEmbedding    *nn.Conv2D `gguf:"patch_embd"`
-	PositionEmbedding ml.Tensor  `gguf:"position_embd.weight"`
-	ClampData         ml.Tensor  `gguf:"clamp_data"`
-	StdBias           ml.Tensor  `gguf:"std_bias"`
-	StdScale          ml.Tensor  `gguf:"std_scale"`
-
-	Layers []VisionEncoderLayer `gguf:"blk"`
-
-	*VisionModelOptions
-	clampInitDone bool
-}
-
-func (m *VisionModel) Forward(ctx ml.Context, pixelValues ml.Tensor, numPatchesX, numPatchesY int) ml.Tensor {
-	numPatches := numPatchesX * numPatchesY
-
-	// Patch embedding via Conv2D
-	hiddenState := m.PatchEmbedding.Forward(ctx, pixelValues, m.patchSize, m.patchSize, 0, 0, 1, 1)
-	hiddenState = hiddenState.Reshape(ctx, numPatches, m.hiddenSize)
-	hiddenState = hiddenState.Permute(ctx, 1, 0, 2, 3).Contiguous(ctx)
-
-	// Conv2D with F16 weights produces F16 output via im2col; cast to F32 for encoder precision
-	hiddenState = hiddenState.Cast(ctx, ml.DTypeF32)
-
-	// 2D positional embeddings from 3D tensor [nEmbd, maxPos, 2]
-	posSize := m.PositionEmbedding.Dim(1)
-	nb1 := m.PositionEmbedding.Stride(1)
-	tblX := m.PositionEmbedding.View(ctx, 0, m.hiddenSize, nb1, posSize)
-	tblY := m.PositionEmbedding.View(ctx, posSize*nb1, m.hiddenSize, nb1, posSize)
-
-	// Position indices for patches
-	posXData := make([]int32, numPatches)
-	posYData := make([]int32, numPatches)
-	for i := range numPatches {
-		posXData[i] = int32(i % numPatchesX)
-		posYData[i] = int32(i / numPatchesX)
-	}
-
-	posXEmb := ctx.Input().FromInts(posXData, numPatches)
-	posYEmb := ctx.Input().FromInts(posYData, numPatches)
-
-	hiddenState = hiddenState.Add(ctx, tblX.Rows(ctx, posXEmb))
-	hiddenState = hiddenState.Add(ctx, tblY.Rows(ctx, posYEmb))
-
-	// No attention mask — all positions are real patches
-	var attnMask ml.Tensor
-
-	// RoPE positions
-	posXRope := ctx.Input().FromInts(posXData, numPatches)
-	posYRope := ctx.Input().FromInts(posYData, numPatches)
-
-	// Vision transformer layers
-	for i := range m.Layers {
-		hiddenState = m.Layers[i].Forward(ctx, hiddenState, posXRope, posYRope, attnMask, m.VisionModelOptions)
-	}
-
-	return hiddenState
-}
-
-func newVisionModel(c fs.Config) *VisionModel {
-	return &VisionModel{
-		Layers: make([]VisionEncoderLayer, c.Uint("vision.block_count")),
-		VisionModelOptions: &VisionModelOptions{
-			hiddenSize: int(c.Uint("vision.embedding_length")),
-			numHeads:   int(c.Uint("vision.attention.head_count")),
-			patchSize:  int(c.Uint("vision.patch_size", 16)),
-			nMerge:     int(c.Uint("vision.projector.scale_factor", 3)),
-			eps:        c.Float("vision.attention.layer_norm_epsilon", 1e-6),
-			ropeTheta:  100.0,
-		},
-	}
-}
-
-func visionPoolAndProject(ctx ml.Context, hiddenState ml.Tensor, numPatchesX, numPatchesY int, opts *VisionModelOptions, proj *MultiModalProjector, stdBias, stdScale ml.Tensor) ml.Tensor {
-	hiddenSize := opts.hiddenSize
-
-	// Reshape from [hiddenSize, numPatches] to spatial layout for pooling
-	hiddenState = hiddenState.Permute(ctx, 1, 0, 2, 3).Contiguous(ctx)
-	hiddenState = hiddenState.Reshape(ctx, numPatchesX, numPatchesY, hiddenSize)
-
-	// AvgPool2D with kernel=stride=nMerge
-	hiddenState = hiddenState.AvgPool2D(ctx, opts.nMerge, opts.nMerge, 0)
-
-	// Reshape back to [hiddenSize, numMergedPatches]
-	mergedX := numPatchesX / opts.nMerge
-	mergedY := numPatchesY / opts.nMerge
-	hiddenState = hiddenState.Reshape(ctx, mergedX*mergedY, hiddenSize)
-	hiddenState = hiddenState.Permute(ctx, 1, 0, 2, 3).Contiguous(ctx)
-
-	hiddenState = hiddenState.Cast(ctx, ml.DTypeF32)
-	hiddenState = hiddenState.Scale(ctx, math.Sqrt(float64(hiddenSize)))
-
-	// Optional vision standardization before projection.
-	if stdBias != nil && stdScale != nil {
-		hiddenState = hiddenState.Sub(ctx, stdBias)
-		hiddenState = hiddenState.Mul(ctx, stdScale)
-	}
-
-	// Project to text embedding dimension
-	hiddenState = proj.Forward(ctx, hiddenState, opts.eps)
-
-	return hiddenState
-}

model/models/gemma4/process_audio.go

@@ -1,280 +0,0 @@
-package gemma4
-
-import (
-	"encoding/binary"
-	"fmt"
-	"math"
-	"math/cmplx"
-)
-
-// Audio preprocessing constants.
-const (
-	audioSampleRate    = 16000
-	melBins            = 128
-	frameLengthMs      = 20.0
-	hopLengthMs        = 10.0
-	minFrequency       = 0.0
-	maxFrequency       = 8000.0
-	melFloor           = 1e-3
-	maxAudioSoftTokens = 750
-)
-
-// Computed from the above constants.
-var (
-	frameLength = int(math.Round(audioSampleRate * frameLengthMs / 1000.0)) // 320
-	hopLength   = int(math.Round(audioSampleRate * hopLengthMs / 1000.0))   // 160
-)
-
-// decodeWAV extracts mono float32 PCM samples from a WAV file, resampled to 16kHz.
-func decodeWAV(data []byte) ([]float32, error) {
-	if len(data) < 12 {
-		return nil, fmt.Errorf("WAV file too short")
-	}
-	if string(data[0:4]) != "RIFF" || string(data[8:12]) != "WAVE" {
-		return nil, fmt.Errorf("not a WAV file")
-	}
-
-	var audioFormat uint16
-	var numChannels, sampleRate, bitsPerSample int
-	var audioData []byte
-	foundFmt := false
-
-	offset := 12
-	for offset+8 <= len(data) {
-		chunkID := string(data[offset : offset+4])
-		chunkSize := int(binary.LittleEndian.Uint32(data[offset+4 : offset+8]))
-		chunkData := data[offset+8 : min(offset+8+chunkSize, len(data))]
-
-		switch chunkID {
-		case "fmt ":
-			if len(chunkData) < 16 {
-				return nil, fmt.Errorf("fmt chunk too short")
-			}
-			audioFormat = binary.LittleEndian.Uint16(chunkData[0:2])
-			numChannels = int(binary.LittleEndian.Uint16(chunkData[2:4]))
-			sampleRate = int(binary.LittleEndian.Uint32(chunkData[4:8]))
-			bitsPerSample = int(binary.LittleEndian.Uint16(chunkData[14:16]))
-			if audioFormat == 0xFFFE && len(chunkData) >= 26 {
-				audioFormat = binary.LittleEndian.Uint16(chunkData[24:26])
-			}
-			foundFmt = true
-		case "data":
-			audioData = chunkData
-		}
-
-		offset += 8 + chunkSize
-		if chunkSize%2 != 0 {
-			offset++
-		}
-	}
-
-	if !foundFmt {
-		return nil, fmt.Errorf("no fmt chunk found in WAV file")
-	}
-	if audioFormat != 1 && audioFormat != 3 {
-		return nil, fmt.Errorf("unsupported WAV format: %d (need PCM=1 or float=3)", audioFormat)
-	}
-	if audioData == nil {
-		return nil, fmt.Errorf("no data chunk found in WAV file")
-	}
-
-	samples := decodeWAVSamples(audioData, audioFormat, bitsPerSample, numChannels)
-	if sampleRate != audioSampleRate {
-		samples = resampleLinear(samples, sampleRate, audioSampleRate)
-	}
-	return samples, nil
-}
-
-func decodeWAVSamples(data []byte, format uint16, bits, channels int) []float32 {
-	bytesPerSample := bits / 8
-	totalSamples := len(data) / (bytesPerSample * channels)
-	mono := make([]float32, totalSamples)
-
-	for i := range totalSamples {
-		var sum float64
-		for ch := range channels {
-			off := (i*channels + ch) * bytesPerSample
-			if off+bytesPerSample > len(data) {
-				break
-			}
-			switch {
-			case format == 1 && bits == 16:
-				v := int16(binary.LittleEndian.Uint16(data[off : off+2]))
-				sum += float64(v) / 32768.0
-			case format == 1 && bits == 32:
-				v := int32(binary.LittleEndian.Uint32(data[off : off+4]))
-				sum += float64(v) / 2147483648.0
-			case format == 1 && bits == 24:
-				v := int32(data[off]) | int32(data[off+1])<<8 | int32(data[off+2])<<16
-				if v&0x800000 != 0 {
-					v |= ^0xFFFFFF
-				}
-				sum += float64(v) / 8388608.0
-			case format == 3 && bits == 32:
-				v := math.Float32frombits(binary.LittleEndian.Uint32(data[off : off+4]))
-				sum += float64(v)
-			case format == 1 && bits == 8:
-				sum += (float64(data[off]) - 128.0) / 128.0
-			}
-		}
-		mono[i] = float32(sum / float64(channels))
-	}
-	return mono
-}
-
-func resampleLinear(samples []float32, fromRate, toRate int) []float32 {
-	n := int(float64(len(samples)) / float64(fromRate) * float64(toRate))
-	out := make([]float32, n)
-	for i := range n {
-		pos := float64(i) * float64(len(samples)-1) / float64(n-1)
-		idx := int(pos)
-		frac := float32(pos - float64(idx))
-		if idx+1 < len(samples) {
-			out[i] = samples[idx]*(1-frac) + samples[idx+1]*frac
-		} else {
-			out[i] = samples[idx]
-		}
-	}
-	return out
-}
-
-// computeMelSpectrogram computes the log mel spectrogram from PCM samples.
-// Returns shape [numFrames, melBins] as float32 slice, and numFrames.
-func computeMelSpectrogram(samples []float32) ([]float32, int) {
-	fftLen := 1
-	for fftLen < frameLength {
-		fftLen <<= 1
-	}
-	fftLen *= 2 // fft_overdrive=True
-
-	// Hanning-nonzero window.
-	window := make([]float64, frameLength)
-	arg := math.Pi * 2.0 / float64(frameLength)
-	for i := range frameLength {
-		window[i] = 0.5 - 0.5*math.Cos(arg*(float64(i)+0.5))
-	}
-
-	numFreqBins := fftLen/2 + 1
-	melFilters := buildMelFilterBank(numFreqBins, melBins, minFrequency, maxFrequency, audioSampleRate)
-
-	frameSizeForUnfold := frameLength + 1
-	numFrames := (len(samples) - frameSizeForUnfold) / hopLength
-	if numFrames <= 0 {
-		return nil, 0
-	}
-
-	result := make([]float32, numFrames*melBins)
-	fftInput := make([]complex128, fftLen)
-
-	for f := range numFrames {
-		start := f * hopLength
-		for i := range frameLength {
-			fftInput[i] = complex(float64(samples[start+i])*window[i], 0)
-		}
-		for i := frameLength; i < fftLen; i++ {
-			fftInput[i] = 0
-		}
-
-		fft(fftInput)
-
-		for m := range melBins {
-			var melVal float64
-			for k := range numFreqBins {
-				mag := cmplx.Abs(fftInput[k])
-				melVal += mag * float64(melFilters[k*melBins+m])
-			}
-			if melVal < melFloor {
-				melVal = melFloor
-			}
-			result[f*melBins+m] = float32(math.Log(melVal))
-		}
-	}
-
-	return result, numFrames
-}
-
-func buildMelFilterBank(numFreqBins, numMels int, fMin, fMax float64, sr int) []float32 {
-	hzToMel := func(f float64) float64 {
-		return 2595.0 * math.Log10(1.0+f/700.0)
-	}
-	melToHz := func(m float64) float64 {
-		return 700.0 * (math.Pow(10.0, m/2595.0) - 1.0)
-	}
-
-	melMin := hzToMel(fMin)
-	melMax := hzToMel(fMax)
-
-	melPts := make([]float64, numMels+2)
-	for i := range melPts {
-		melPts[i] = melMin + float64(i)*(melMax-melMin)/float64(numMels+1)
-	}
-	filterFreqs := make([]float64, numMels+2)
-	for i, m := range melPts {
-		filterFreqs[i] = melToHz(m)
-	}
-
-	fftFreqs := make([]float64, numFreqBins)
-	for i := range fftFreqs {
-		fftFreqs[i] = float64(i) * float64(sr) / float64(2*(numFreqBins-1))
-	}
-
-	filters := make([]float32, numFreqBins*numMels)
-	for m := range numMels {
-		fLeft := filterFreqs[m]
-		fCenter := filterFreqs[m+1]
-		fRight := filterFreqs[m+2]
-		for k := range numFreqBins {
-			f := fftFreqs[k]
-			var v float64
-			if f >= fLeft && f <= fCenter && fCenter > fLeft {
-				v = (f - fLeft) / (fCenter - fLeft)
-			} else if f > fCenter && f <= fRight && fRight > fCenter {
-				v = (fRight - f) / (fRight - fCenter)
-			}
-			if v > 0 {
-				filters[k*numMels+m] = float32(v)
-			}
-		}
-	}
-	return filters
-}
-
-// fft performs an in-place Cooley-Tukey radix-2 FFT.
-func fft(x []complex128) {
-	n := len(x)
-	if n <= 1 {
-		return
-	}
-
-	j := 0
-	for i := 1; i < n; i++ {
-		bit := n >> 1
-		for j&bit != 0 {
-			j ^= bit
-			bit >>= 1
-		}
-		j ^= bit
-		if i < j {
-			x[i], x[j] = x[j], x[i]
-		}
-	}
-
-	for size := 2; size <= n; size <<= 1 {
-		halfSize := size / 2
-		w := complex(math.Cos(2*math.Pi/float64(size)), -math.Sin(2*math.Pi/float64(size)))
-		for start := 0; start < n; start += size {
-			wn := complex(1, 0)
-			for k := range halfSize {
-				t := wn * x[start+k+halfSize]
-				x[start+k+halfSize] = x[start+k] - t
-				x[start+k] = x[start+k] + t
-				wn *= w
-			}
-		}
-	}
-}
-
-// isAudioData checks if the data starts with WAV magic bytes.
-func isAudioData(data []byte) bool {
-	return len(data) >= 12 && string(data[0:4]) == "RIFF" && string(data[8:12]) == "WAVE"
-}

model/models/gemma4/process_image.go

@@ -1,103 +0,0 @@
-package gemma4
-
-import (
-	"image"
-	"math"
-
-	"golang.org/x/image/draw"
-
-	"github.com/ollama/ollama/fs"
-)
-
-type ImageProcessor struct {
-	patchSize   int
-	numChannels int
-	nMerge      int
-	minPixels   int
-	maxPixels   int
-}
-
-func newImageProcessor(c fs.Config) ImageProcessor {
-	patchSize := int(c.Uint("vision.patch_size", 16))
-	nMerge := int(c.Uint("vision.projector.scale_factor", 3))
-	numChannels := int(c.Uint("vision.num_channels", 3))
-
-	// Token limits from reference: min=40, max=280 output tokens after pooling.
-	// Convert to pixel counts: tokens * nMerge^2 * patchSize^2
-	minTokens := 40
-	maxTokens := 280
-	patchArea := patchSize * patchSize * nMerge * nMerge
-	minPixels := minTokens * patchArea
-	maxPixels := maxTokens * patchArea
-
-	return ImageProcessor{
-		patchSize:   patchSize,
-		numChannels: numChannels,
-		nMerge:      nMerge,
-		minPixels:   minPixels,
-		maxPixels:   maxPixels,
-	}
-}
-
-// ProcessImage resizes an image preserving aspect ratio, aligning dimensions
-// to (patchSize * nMerge) boundaries, and normalizes pixels to [-1, 1].
-// Returns the float32 pixel data and the actual output dimensions.
-func (p *ImageProcessor) ProcessImage(img image.Image) ([]float32, int, int, error) {
-	// Compute target size preserving aspect ratio
-	alignSize := p.patchSize * p.nMerge
-	targetW, targetH := p.smartResize(img.Bounds().Dx(), img.Bounds().Dy(), alignSize)
-
-	// Resize directly without alpha compositing, matching MLX reference.
-	dst := image.NewRGBA(image.Rect(0, 0, targetW, targetH))
-	draw.BiLinear.Scale(dst, dst.Bounds(), img, img.Bounds(), draw.Over, nil)
-
-	// Normalize to [-1, 1] using mean=0.5, std=0.5: (pixel/255 - 0.5) / 0.5 = 2*pixel/255 - 1
-	data := p.pack(dst)
-	return data, targetW, targetH, nil
-}
-
-// smartResize computes target dimensions that preserve aspect ratio and
-// align to alignSize boundaries. It scales the image to fill the maximum
-// patch budget (maxPixels), matching the MLX reference.
-func (p *ImageProcessor) smartResize(origW, origH, alignSize int) (int, int) {
-	totalPx := origW * origH
-
-	var targetW, targetH int
-	if p.maxPixels > 0 && totalPx > 0 {
-		factor := math.Sqrt(float64(p.maxPixels) / float64(totalPx))
-		targetH = max(alignSize, int(math.Floor(factor*float64(origH)/float64(alignSize)))*alignSize)
-		targetW = max(alignSize, int(math.Floor(factor*float64(origW)/float64(alignSize)))*alignSize)
-	} else {
-		targetH = max(alignSize, (origH/alignSize)*alignSize)
-		targetW = max(alignSize, (origW/alignSize)*alignSize)
-	}
-
-	return targetW, targetH
-}
-
-// pack extracts RGB values from an image and normalizes to [-1, 1].
-// Returns channel-first layout: [R..., G..., B...].
-func (p *ImageProcessor) pack(img image.Image) []float32 {
-	bounds := img.Bounds()
-	w := bounds.Dx()
-	h := bounds.Dy()
-	size := w * h
-
-	pixelVals := make([]float32, 3*size)
-	rOff, gOff, bOff := 0, size, 2*size
-
-	for y := bounds.Min.Y; y < bounds.Max.Y; y++ {
-		for x := bounds.Min.X; x < bounds.Max.X; x++ {
-			c := img.At(x, y)
-			r, g, b, _ := c.RGBA()
-			idx := (y-bounds.Min.Y)*w + (x - bounds.Min.X)
-
-			// Normalize [0, 255] -> [-1, 1]: 2 * (val/255) - 1
-			pixelVals[rOff+idx] = float32(r>>8)/255.0*2.0 - 1.0
-			pixelVals[gOff+idx] = float32(g>>8)/255.0*2.0 - 1.0
-			pixelVals[bOff+idx] = float32(b>>8)/255.0*2.0 - 1.0
-		}
-	}
-
-	return pixelVals
-}

model/models/laguna/model.go

@@ -1,444 +0,0 @@
-package laguna
-
-import (
-	"fmt"
-	"math"
-
-	"github.com/ollama/ollama/fs"
-	"github.com/ollama/ollama/kvcache"
-	"github.com/ollama/ollama/ml"
-	"github.com/ollama/ollama/ml/nn"
-	"github.com/ollama/ollama/ml/nn/rope"
-	"github.com/ollama/ollama/model"
-	"github.com/ollama/ollama/model/input"
-	"github.com/ollama/ollama/tokenizer"
-)
-
-const (
-	cacheTypeSWA = iota
-	cacheTypeCausal
-)
-
-type Options struct {
-	hiddenSize int
-	headDim    int
-
-	numHeads   []int
-	numKVHeads int
-
-	eps float32
-
-	slidingWindow        int
-	slidingWindowPattern []bool
-
-	fullRopeDim                   int
-	fullRopeBase, fullRopeScale   float32
-	fullRopeOriginalContextLength int
-	fullRopeAttentionFactor       float32
-	fullRopeBetaFast              float32
-	fullRopeBetaSlow              float32
-
-	swaRopeDim                 int
-	swaRopeBase, swaRopeScale  float32
-	numExperts, numExpertsUsed int
-	normTopKProb               bool
-	routedScalingFactor        float32
-	decoderSparseStep          int
-	denseLayers                map[int]bool
-}
-
-func (o *Options) numHeadsForLayer(layer int) int {
-	if layer < len(o.numHeads) && o.numHeads[layer] > 0 {
-		return o.numHeads[layer]
-	}
-	if len(o.numHeads) > 0 && o.numHeads[0] > 0 {
-		return o.numHeads[0]
-	}
-	return 1
-}
-
-func (o *Options) layerIsSliding(layer int) bool {
-	return layer < len(o.slidingWindowPattern) && o.slidingWindowPattern[layer]
-}
-
-func (o *Options) layerUsesMoE(layer int) bool {
-	if o.numExperts == 0 || o.denseLayers[layer] {
-		return false
-	}
-	step := o.decoderSparseStep
-	if step <= 0 {
-		step = 1
-	}
-	return (layer+1)%step == 0
-}
-
-func (o *Options) applyRotaryPositionEmbeddings(ctx ml.Context, layer int, states, positions ml.Tensor) ml.Tensor {
-	opts := []func(*rope.Options){rope.WithTypeNeoX()}
-	if o.layerIsSliding(layer) {
-		return nn.RoPE(ctx, states, positions, o.swaRopeDim, o.swaRopeBase, 1./o.swaRopeScale, opts...)
-	}
-
-	opts = append(opts,
-		rope.WithOriginalContextLength(o.fullRopeOriginalContextLength),
-		rope.WithExtrapolationFactor(1),
-		rope.WithAttentionFactor(o.fullRopeAttentionFactor),
-		rope.WithBetaFast(o.fullRopeBetaFast),
-		rope.WithBetaSlow(o.fullRopeBetaSlow),
-	)
-	return nn.RoPE(ctx, states, positions, o.fullRopeDim, o.fullRopeBase, 1./o.fullRopeScale, opts...)
-}
-
-type Attention struct {
-	Query     *nn.Linear  `gguf:"attn_q"`
-	QueryNorm *nn.RMSNorm `gguf:"attn_q_norm"`
-	Key       *nn.Linear  `gguf:"attn_k"`
-	KeyNorm   *nn.RMSNorm `gguf:"attn_k_norm"`
-	Value     *nn.Linear  `gguf:"attn_v"`
-	Gate      *nn.Linear  `gguf:"attn_g"`
-	Output    *nn.Linear  `gguf:"attn_output"`
-}
-
-func (sa *Attention) Forward(ctx ml.Context, layer int, hiddenStates, positions ml.Tensor, cache kvcache.Cache, opts *Options) ml.Tensor {
-	batchSize := hiddenStates.Dim(1)
-	numHeads := opts.numHeadsForLayer(layer)
-
-	query := sa.Query.Forward(ctx, hiddenStates)
-	key := sa.Key.Forward(ctx, hiddenStates)
-	value := sa.Value.Forward(ctx, hiddenStates)
-	gate := sa.Gate.Forward(ctx, hiddenStates)
-
-	query = query.Reshape(ctx, opts.headDim, numHeads, batchSize)
-	key = key.Reshape(ctx, opts.headDim, opts.numKVHeads, batchSize)
-	value = value.Reshape(ctx, opts.headDim, opts.numKVHeads, batchSize)
-
-	query = sa.QueryNorm.Forward(ctx, query, opts.eps)
-	key = sa.KeyNorm.Forward(ctx, key, opts.eps)
-
-	query = opts.applyRotaryPositionEmbeddings(ctx, layer, query, positions)
-	key = opts.applyRotaryPositionEmbeddings(ctx, layer, key, positions)
-
-	attention := nn.Attention(ctx, query, key, value, 1./math.Sqrt(float64(opts.headDim)), cache)
-	// Laguna applies the per-head gate softplus in float32, then casts back.
-	gate = gate.Cast(ctx, ml.DTypeF32).Softplus(ctx).Cast(ctx, attention.DType())
-	attention = attention.Mul(ctx, gate.Reshape(ctx, 1, numHeads, batchSize))
-	attention = attention.Reshape(ctx, opts.headDim*numHeads, batchSize)
-	return sa.Output.Forward(ctx, attention)
-}
-
-type MLP interface {
-	Forward(ml.Context, ml.Tensor, *Options) ml.Tensor
-}
-
-type dense struct {
-	Gate *nn.Linear `gguf:"ffn_gate"`
-	Up   *nn.Linear `gguf:"ffn_up"`
-	Down *nn.Linear `gguf:"ffn_down"`
-}
-
-func (mlp *dense) Forward(ctx ml.Context, hiddenStates ml.Tensor, _ *Options) ml.Tensor {
-	hiddenStates = mlp.Gate.Forward(ctx, hiddenStates).SILU(ctx, mlp.Up.Forward(ctx, hiddenStates))
-	return mlp.Down.Forward(ctx, hiddenStates)
-}
-
-type sparse struct {
-	Router       *nn.Linear      `gguf:"ffn_gate_inp"`
-	Gate         *nn.LinearBatch `gguf:"ffn_gate_exps"`
-	Up           *nn.LinearBatch `gguf:"ffn_up_exps"`
-	Down         *nn.LinearBatch `gguf:"ffn_down_exps"`
-	SharedExpert *dense          `gguf:",suf:_shexp"`
-	ExpProbsBias ml.Tensor       `gguf:"exp_probs_b.bias,alt:exp_probs_b"`
-}
-
-func (moe *sparse) topKIndices(ctx ml.Context, scores ml.Tensor, opts *Options) ml.Tensor {
-	if moe.ExpProbsBias != nil {
-		scores = scores.Add(ctx, moe.ExpProbsBias)
-	}
-	return scores.TopK(ctx, opts.numExpertsUsed)
-}
-
-func (moe *sparse) Forward(ctx ml.Context, hiddenStates ml.Tensor, opts *Options) ml.Tensor {
-	residual := hiddenStates
-
-	scores := moe.Router.Forward(ctx, hiddenStates).Cast(ctx, ml.DTypeF32).Sigmoid(ctx)
-	selectedExperts := moe.topKIndices(ctx, scores, opts)
-	routingWeights := scores.Reshape(ctx, 1, opts.numExperts, hiddenStates.Dim(1)).Rows(ctx, selectedExperts)
-	if opts.normTopKProb {
-		routingWeights = routingWeights.Reshape(ctx, opts.numExpertsUsed, hiddenStates.Dim(1))
-		routingWeights = routingWeights.Div(ctx, routingWeights.SumRows(ctx))
-		routingWeights = routingWeights.Reshape(ctx, 1, opts.numExpertsUsed, hiddenStates.Dim(1))
-	}
-	routingWeights = routingWeights.Scale(ctx, float64(opts.routedScalingFactor))
-
-	hiddenStates = hiddenStates.Reshape(ctx, hiddenStates.Dim(0), 1, hiddenStates.Dim(1))
-	upStates := moe.Up.Forward(ctx, hiddenStates, selectedExperts)
-	hiddenStates = moe.Gate.Forward(ctx, hiddenStates, selectedExperts).SILU(ctx, upStates)
-
-	experts := moe.Down.Forward(ctx, hiddenStates, selectedExperts)
-	experts = experts.Mul(ctx, routingWeights)
-
-	nextStates := experts.View(ctx, 0, experts.Dim(0), experts.Stride(2), experts.Dim(2))
-	for i := 1; i < opts.numExpertsUsed; i++ {
-		nextStates = nextStates.Add(ctx, experts.View(ctx, i*experts.Stride(1), experts.Dim(0), experts.Stride(2), experts.Dim(2)))
-	}
-
-	return nextStates.Add(ctx, moe.SharedExpert.Forward(ctx, residual, opts))
-}
-
-type Layer struct {
-	AttentionNorm *nn.RMSNorm `gguf:"attn_norm"`
-	*Attention
-
-	MLPNorm *nn.RMSNorm `gguf:"ffn_norm"`
-	MLP
-}
-
-func (l *Layer) Forward(ctx ml.Context, layer int, hiddenStates, positions, outputs ml.Tensor, cache kvcache.Cache, opts *Options) ml.Tensor {
-	residual := hiddenStates
-	hiddenStates = l.AttentionNorm.Forward(ctx, hiddenStates, opts.eps)
-	hiddenStates = l.Attention.Forward(ctx, layer, hiddenStates, positions, cache, opts)
-
-	if outputs != nil {
-		hiddenStates = hiddenStates.Rows(ctx, outputs)
-		residual = residual.Rows(ctx, outputs)
-	}
-
-	hiddenStates = hiddenStates.Add(ctx, residual)
-	residual = hiddenStates
-
-	hiddenStates = l.MLPNorm.Forward(ctx, hiddenStates, opts.eps)
-	hiddenStates = l.MLP.Forward(ctx, hiddenStates, opts)
-	return hiddenStates.Add(ctx, residual)
-}
-
-type Model struct {
-	model.Base
-	tokenizer.Tokenizer
-
-	TokenEmbedding *nn.Embedding `gguf:"token_embd"`
-	Layers         []Layer       `gguf:"blk"`
-	OutputNorm     *nn.RMSNorm   `gguf:"output_norm"`
-	Output         *nn.Linear    `gguf:"output,alt:token_embd"`
-
-	*Options
-}
-
-func New(c fs.Config) (model.Model, error) {
-	if c.Bool("attention.sink_enabled") {
-		return nil, fmt.Errorf("laguna: SWA attention sinks are not supported")
-	}
-	if c.Uint("attention.gating_type") != 1 {
-		return nil, fmt.Errorf("laguna: unsupported attention gating type %d", c.Uint("attention.gating_type"))
-	}
-	if !c.Bool("attention.qk_norm") {
-		return nil, fmt.Errorf("laguna: Q/K RMSNorm is required")
-	}
-	if gating := c.Uint("expert_gating_func"); gating != 2 {
-		return nil, fmt.Errorf("laguna: unsupported expert gating function %d", gating)
-	}
-
-	numLayers := int(c.Uint("block_count"))
-	opts := newOptions(c, numLayers)
-	layers := make([]Layer, numLayers)
-	for i := range layers {
-		if opts.layerUsesMoE(i) {
-			layers[i].MLP = &sparse{}
-		} else {
-			layers[i].MLP = &dense{}
-		}
-	}
-
-	var pre []string
-	switch c.String("tokenizer.ggml.pre") {
-	case "laguna":
-		pre = []string{
-			`(?:\r?\n)+(?!\r?\n)`,
-			`(?i:'s|'t|'re|'ve|'m|'ll|'d)|[^\r\n\p{L}\p{N}]?\p{L}+|\p{N}| ?[^\s\p{L}\p{N}]+[\r\n]*|\s*[\r\n]+|\s+(?!\S)|\s+`,
-		}
-	default:
-		return nil, model.ErrUnsupportedTokenizer
-	}
-
-	m := Model{
-		Tokenizer: tokenizer.NewBytePairEncoding(
-			&tokenizer.Vocabulary{
-				Values: c.Strings("tokenizer.ggml.tokens"),
-				Types:  c.Ints("tokenizer.ggml.token_type"),
-				Merges: c.Strings("tokenizer.ggml.merges"),
-				AddBOS: c.Bool("tokenizer.ggml.add_bos_token", true),
-				BOS:    []int32{int32(c.Uint("tokenizer.ggml.bos_token_id"))},
-				AddEOS: c.Bool("tokenizer.ggml.add_eos_token", false),
-				EOS: append(
-					[]int32{int32(c.Uint("tokenizer.ggml.eos_token_id"))},
-					c.Ints("tokenizer.ggml.eos_token_ids")...,
-				),
-			},
-			pre...,
-		),
-		Layers:  layers,
-		Options: opts,
-	}
-
-	m.Cache = kvcache.NewWrapperCache(
-		kvcache.NewSWACache(int32(opts.slidingWindow), m.Shift),
-		kvcache.NewCausalCache(m.Shift),
-	)
-	return &m, nil
-}
-
-func newOptions(c fs.Config, numLayers int) *Options {
-	denseLayers := make(map[int]bool)
-	for _, layer := range configUints(c, "dense_layers") {
-		denseLayers[int(layer)] = true
-	}
-	for i := range c.Uint("leading_dense_block_count") {
-		denseLayers[int(i)] = true
-	}
-
-	fullRopeScale := c.Float("rope.scaling.factor", 1)
-	if fullRopeScale == 0 {
-		fullRopeScale = 1
-	}
-	swaRopeScale := c.Float("rope.swa.scaling.factor", 1)
-	if swaRopeScale == 0 {
-		swaRopeScale = 1
-	}
-	fullRopeType := c.String("rope.scaling.type")
-	fullRopeAttentionFactor := lagunaAttentionFactor(fullRopeType, fullRopeScale, c.Float("rope.scaling.attn_factor"))
-
-	return &Options{
-		hiddenSize:                    int(c.Uint("embedding_length")),
-		headDim:                       int(c.Uint("attention.key_length")),
-		numHeads:                      expandIntArray(configUints(c, "attention.head_count"), numLayers, c.Uint("attention.head_count", 1)),
-		numKVHeads:                    int(c.Uint("attention.head_count_kv")),
-		eps:                           c.Float("attention.layer_norm_rms_epsilon", 1e-6),
-		slidingWindow:                 int(c.Uint("attention.sliding_window", 512)),
-		slidingWindowPattern:          slidingWindowPattern(c, numLayers),
-		fullRopeDim:                   int(c.Uint("rope.dimension_count", c.Uint("attention.key_length"))),
-		fullRopeBase:                  c.Float("rope.freq_base", 500000),
-		fullRopeScale:                 fullRopeScale,
-		fullRopeOriginalContextLength: int(c.Uint("rope.scaling.original_context_length", 4096)),
-		fullRopeAttentionFactor:       fullRopeAttentionFactor,
-		fullRopeBetaFast:              c.Float("rope.scaling.beta_fast", 64),
-		fullRopeBetaSlow:              c.Float("rope.scaling.beta_slow", 1),
-		swaRopeDim:                    int(c.Uint("rope.swa.dimension_count", c.Uint("attention.key_length"))),
-		swaRopeBase:                   c.Float("rope.swa.freq_base", 10000),
-		swaRopeScale:                  swaRopeScale,
-		numExperts:                    int(c.Uint("expert_count")),
-		numExpertsUsed:                int(c.Uint("expert_used_count")),
-		normTopKProb:                  c.Bool("expert_weights_norm", true),
-		routedScalingFactor:           c.Float("expert_weights_scale", 1),
-		decoderSparseStep:             int(c.Uint("decoder_sparse_step", 1)),
-		denseLayers:                   denseLayers,
-	}
-}
-
-func lagunaAttentionFactor(ropeType string, scaleFactor, attentionFactor float32) float32 {
-	if attentionFactor != 0 {
-		return attentionFactor
-	}
-	if ropeType == "yarn" && scaleFactor > 1 {
-		return float32(0.1*math.Log(float64(scaleFactor)) + 1)
-	}
-	return 1
-}
-
-func slidingWindowPattern(c fs.Config, numLayers int) []bool {
-	pattern := c.Bools("attention.sliding_window_pattern")
-	if len(pattern) == numLayers {
-		return pattern
-	}
-
-	layerTypes := configUints(c, "attention.layer_types")
-	if len(layerTypes) == numLayers {
-		pattern = make([]bool, numLayers)
-		for i, layerType := range layerTypes {
-			pattern[i] = layerType == 1
-		}
-		return pattern
-	}
-
-	return make([]bool, numLayers)
-}
-
-func configUints(c fs.Config, key string) []uint32 {
-	keyExists := c.Value(c.Architecture()+"."+key) != nil || c.Value(key) != nil
-	if cc, ok := c.(interface {
-		Uints(string, ...[]uint32) []uint32
-	}); ok {
-		if values := cc.Uints(key); len(values) > 0 && (keyExists || !(len(values) == 1 && values[0] == 0)) {
-			return values
-		}
-	}
-
-	ints := c.Ints(key)
-	if len(ints) > 0 && (keyExists || !(len(ints) == 1 && ints[0] == 0)) {
-		values := make([]uint32, len(ints))
-		for i, v := range ints {
-			values[i] = uint32(v)
-		}
-		return values
-	}
-
-	if scalar := c.Uint(key); scalar != 0 {
-		return []uint32{scalar}
-	}
-	return nil
-}
-
-func expandIntArray(values []uint32, n int, fallback uint32) []int {
-	if len(values) == 0 {
-		values = []uint32{fallback}
-	}
-	defaultValue := values[0]
-	if len(values) == 1 {
-		defaultValue = values[0]
-	}
-
-	out := make([]int, n)
-	for i := range out {
-		if i < len(values) {
-			out[i] = int(values[i])
-		} else {
-			out[i] = int(defaultValue)
-		}
-	}
-	return out
-}
-
-func (m *Model) Shift(ctx ml.Context, layer int, key, shift ml.Tensor) (ml.Tensor, error) {
-	return m.Options.applyRotaryPositionEmbeddings(ctx, layer, key, shift), nil
-}
-
-func (m *Model) Forward(ctx ml.Context, batch input.Batch) (ml.Tensor, error) {
-	positions := ctx.Input().FromInts(batch.Positions, len(batch.Positions))
-	hiddenStates := m.TokenEmbedding.Forward(ctx, batch.Inputs)
-
-	for i, layer := range m.Layers {
-		if m.Cache != nil {
-			m.Cache.SetLayer(i)
-			if wrapper, ok := m.Cache.(*kvcache.WrapperCache); ok {
-				cacheType := cacheTypeCausal
-				if m.Options.layerIsSliding(i) {
-					cacheType = cacheTypeSWA
-				}
-				wrapper.SetLayerType(cacheType)
-			}
-		}
-
-		var outputs ml.Tensor
-		if i == len(m.Layers)-1 {
-			outputs = batch.Outputs
-		}
-
-		hiddenStates = layer.Forward(ctx, i, hiddenStates, positions, outputs, m.Cache, m.Options)
-	}
-
-	hiddenStates = m.OutputNorm.Forward(ctx, hiddenStates, m.eps)
-	return m.Output.Forward(ctx, hiddenStates), nil
-}
-
-func init() {
-	model.Register("laguna", New)
-}
-
-var _ model.Model = (*Model)(nil)

model/models/laguna/model_test.go

@@ -1,237 +0,0 @@
-package laguna
-
-import (
-	"iter"
-	"math"
-	"testing"
-)
-
-type testConfig map[string]any
-
-func (c testConfig) Architecture() string { return "laguna" }
-
-func (c testConfig) key(key string) string {
-	switch {
-	case len(key) >= len("tokenizer.") && key[:len("tokenizer.")] == "tokenizer.":
-		return key
-	case len(key) >= len("general.") && key[:len("general.")] == "general.":
-		return key
-	default:
-		return "laguna." + key
-	}
-}
-
-func (c testConfig) String(key string, defaultValue ...string) string {
-	if v, ok := c[c.key(key)].(string); ok {
-		return v
-	}
-	if len(defaultValue) > 0 {
-		return defaultValue[0]
-	}
-	return ""
-}
-
-func (c testConfig) Uint(key string, defaultValue ...uint32) uint32 {
-	switch v := c[c.key(key)].(type) {
-	case uint32:
-		return v
-	case int:
-		return uint32(v)
-	}
-	if len(defaultValue) > 0 {
-		return defaultValue[0]
-	}
-	return 0
-}
-
-func (c testConfig) Float(key string, defaultValue ...float32) float32 {
-	if v, ok := c[c.key(key)].(float32); ok {
-		return v
-	}
-	if len(defaultValue) > 0 {
-		return defaultValue[0]
-	}
-	return 0
-}
-
-func (c testConfig) Bool(key string, defaultValue ...bool) bool {
-	if v, ok := c[c.key(key)].(bool); ok {
-		return v
-	}
-	if len(defaultValue) > 0 {
-		return defaultValue[0]
-	}
-	return false
-}
-
-func (c testConfig) Strings(key string, defaultValue ...[]string) []string {
-	if v, ok := c[c.key(key)].([]string); ok {
-		return v
-	}
-	if len(defaultValue) > 0 {
-		return defaultValue[0]
-	}
-	return nil
-}
-
-func (c testConfig) Ints(key string, defaultValue ...[]int32) []int32 {
-	if v, ok := c[c.key(key)].([]int32); ok {
-		return v
-	}
-	if len(defaultValue) > 0 {
-		return defaultValue[0]
-	}
-	return nil
-}
-
-func (c testConfig) Uints(key string, defaultValue ...[]uint32) []uint32 {
-	if v, ok := c[c.key(key)].([]uint32); ok {
-		return v
-	}
-	if len(defaultValue) > 0 {
-		return defaultValue[0]
-	}
-	return nil
-}
-
-func (c testConfig) Floats(key string, defaultValue ...[]float32) []float32 {
-	if v, ok := c[c.key(key)].([]float32); ok {
-		return v
-	}
-	if len(defaultValue) > 0 {
-		return defaultValue[0]
-	}
-	return nil
-}
-
-func (c testConfig) Bools(key string, defaultValue ...[]bool) []bool {
-	if v, ok := c[c.key(key)].([]bool); ok {
-		return v
-	}
-	if len(defaultValue) > 0 {
-		return defaultValue[0]
-	}
-	return nil
-}
-
-func (c testConfig) Len() int { return len(c) }
-
-func (c testConfig) Keys() iter.Seq[string] {
-	return func(yield func(string) bool) {
-		for key := range c {
-			if !yield(key) {
-				return
-			}
-		}
-	}
-}
-
-func (c testConfig) Value(key string) any { return c[key] }
-
-func TestNewOptionsLayerConfig(t *testing.T) {
-	cfg := testConfig{
-		"laguna.block_count":                          uint32(4),
-		"laguna.embedding_length":                     uint32(128),
-		"laguna.attention.key_length":                 uint32(16),
-		"laguna.attention.head_count":                 []uint32{8, 16, 16, 16},
-		"laguna.attention.head_count_kv":              uint32(4),
-		"laguna.attention.layer_norm_rms_epsilon":     float32(1e-6),
-		"laguna.attention.sliding_window":             uint32(512),
-		"laguna.attention.sliding_window_pattern":     []bool{false, true, true, true},
-		"laguna.rope.dimension_count":                 uint32(8),
-		"laguna.rope.freq_base":                       float32(500000),
-		"laguna.rope.scaling.factor":                  float32(32),
-		"laguna.rope.scaling.original_context_length": uint32(4096),
-		"laguna.rope.swa.dimension_count":             uint32(16),
-		"laguna.rope.swa.freq_base":                   float32(10000),
-		"laguna.expert_count":                         uint32(32),
-		"laguna.expert_used_count":                    uint32(4),
-		"laguna.expert_weights_norm":                  true,
-		"laguna.expert_weights_scale":                 float32(2.5),
-		"laguna.decoder_sparse_step":                  uint32(1),
-		"laguna.dense_layers":                         []uint32{0},
-	}
-
-	opts := newOptions(cfg, 4)
-	if got := opts.numHeadsForLayer(0); got != 8 {
-		t.Fatalf("layer 0 heads = %d, want 8", got)
-	}
-	if got := opts.numHeadsForLayer(1); got != 16 {
-		t.Fatalf("layer 1 heads = %d, want 16", got)
-	}
-	if opts.layerIsSliding(0) {
-		t.Fatal("layer 0 should be full attention")
-	}
-	if !opts.layerIsSliding(1) {
-		t.Fatal("layer 1 should be sliding attention")
-	}
-	if opts.layerUsesMoE(0) {
-		t.Fatal("layer 0 should be dense")
-	}
-	if !opts.layerUsesMoE(1) {
-		t.Fatal("layer 1 should use MoE")
-	}
-	if opts.fullRopeDim != 8 || opts.swaRopeDim != 16 {
-		t.Fatalf("rope dims = full %d swa %d, want 8/16", opts.fullRopeDim, opts.swaRopeDim)
-	}
-}
-
-func TestNewOptionsYarnAttentionFactorFallback(t *testing.T) {
-	cfg := testConfig{
-		"laguna.block_count":             uint32(1),
-		"laguna.embedding_length":        uint32(128),
-		"laguna.attention.key_length":    uint32(16),
-		"laguna.attention.head_count":    uint32(8),
-		"laguna.attention.head_count_kv": uint32(4),
-		"laguna.rope.scaling.type":       "yarn",
-		"laguna.rope.scaling.factor":     float32(32),
-	}
-
-	opts := newOptions(cfg, 1)
-	want := float32(0.1*math.Log(32) + 1)
-	if got := opts.fullRopeAttentionFactor; math.Abs(float64(got-want)) > 1e-6 {
-		t.Fatalf("fullRopeAttentionFactor = %v, want %v", got, want)
-	}
-}
-
-func TestNewRejectsUnsupportedLagunaVariants(t *testing.T) {
-	tests := []struct {
-		name string
-		cfg  testConfig
-	}{
-		{
-			name: "attention sinks",
-			cfg: testConfig{
-				"laguna.attention.sink_enabled": true,
-			},
-		},
-		{
-			name: "non per-head gate",
-			cfg: testConfig{
-				"laguna.attention.gating_type": uint32(0),
-			},
-		},
-		{
-			name: "missing qk norm",
-			cfg: testConfig{
-				"laguna.attention.gating_type": uint32(1),
-			},
-		},
-		{
-			name: "non sigmoid experts",
-			cfg: testConfig{
-				"laguna.attention.gating_type": uint32(1),
-				"laguna.attention.qk_norm":     true,
-				"laguna.expert_gating_func":    uint32(1),
-			},
-		},
-	}
-
-	for _, tt := range tests {
-		t.Run(tt.name, func(t *testing.T) {
-			if _, err := New(tt.cfg); err == nil {
-				t.Fatal("expected unsupported variant error")
-			}
-		})
-	}
-}

model/parsers/parsers.go

@@ -54,6 +54,8 @@ func ParserForName(name string) Parser {
 		p = &Qwen3Parser{hasThinkingSupport: true, defaultThinking: true}
 	case "qwen3.5":
 		p = &Qwen35Parser{}
+	case "ornith":
+		p = &Qwen35Parser{}
 	case "qwen3-coder":
 		p = &Qwen3CoderParser{}
 	case "qwen3-vl-instruct":

model/parsers/parsers_test.go

@@ -65,6 +65,7 @@ func TestBuiltInParsersStillWork(t *testing.T) {
 		{"lfm2"},
 		{"lfm2-thinking"},
 		{"qwen3.5"},
+		{"ornith"},
 		{"harmony"},
 	}
 

model/renderers/ornith.go

@@ -0,0 +1,16 @@
+package renderers
+
+type OrnithRenderer struct {
+	Qwen35Renderer
+}
+
+func newOrnithRenderer() Renderer {
+	return &OrnithRenderer{
+		Qwen35Renderer: Qwen35Renderer{
+			isThinking:                      true,
+			alwaysRenderAssistantThinkBlock: true,
+			emitEmptyThinkOnNoThink:         true,
+			useImgTags:                      RenderImgTags,
+		},
+	}
+}

model/renderers/ornith_test.go

@@ -0,0 +1,74 @@
+package renderers
+
+import (
+	"testing"
+
+	"github.com/ollama/ollama/api"
+)
+
+func TestOrnithRendererMatchesAssistantHistoryThinkBlocks(t *testing.T) {
+	msgs := []api.Message{
+		{Role: "user", Content: "Say hello."},
+		{Role: "assistant", Content: "Hello."},
+		{Role: "user", Content: "Now say bye."},
+	}
+
+	got, err := RenderWithRenderer("ornith", msgs, nil, nil)
+	if err != nil {
+		t.Fatalf("render failed: %v", err)
+	}
+
+	want := `<|im_start|>user
+Say hello.<|im_end|>
+<|im_start|>assistant
+<think>
+
+</think>
+
+Hello.<|im_end|>
+<|im_start|>user
+Now say bye.<|im_end|>
+<|im_start|>assistant
+<think>
+`
+	if got != want {
+		t.Fatalf("unexpected Ornith render\n--- got ---\n%q\n--- want ---\n%q", got, want)
+	}
+}
+
+func TestOrnithRendererKeepsAssistantThinkBlocksWhenThinkingDisabled(t *testing.T) {
+	msgs := []api.Message{
+		{Role: "user", Content: "Say hello."},
+		{
+			Role:     "assistant",
+			Thinking: "Keep it short.",
+			Content:  "Hello.",
+		},
+		{Role: "user", Content: "Now say bye."},
+	}
+
+	got, err := RenderWithRenderer("ornith", msgs, nil, &api.ThinkValue{Value: false})
+	if err != nil {
+		t.Fatalf("render failed: %v", err)
+	}
+
+	want := `<|im_start|>user
+Say hello.<|im_end|>
+<|im_start|>assistant
+<think>
+Keep it short.
+</think>
+
+Hello.<|im_end|>
+<|im_start|>user
+Now say bye.<|im_end|>
+<|im_start|>assistant
+<think>
+
+</think>
+
+`
+	if got != want {
+		t.Fatalf("unexpected Ornith render with thinking disabled\n--- got ---\n%q\n--- want ---\n%q", got, want)
+	}
+}

model/renderers/qwen35.go

@@ -39,8 +39,9 @@ Reminder:
 type Qwen35Renderer struct {
 	isThinking bool
 
-	emitEmptyThinkOnNoThink bool
-	useImgTags              bool
+	alwaysRenderAssistantThinkBlock bool
+	emitEmptyThinkOnNoThink         bool
+	useImgTags                      bool
 }
 
 func (r *Qwen35Renderer) LeadingBOS() string {
@@ -140,9 +141,10 @@ func (r *Qwen35Renderer) Render(messages []api.Message, tools []api.Tool, think
 		if message.Role == "user" || (message.Role == "system" && i != 0) {
 			sb.WriteString(imStartTag + message.Role + "\n" + content + imEndTag + "\n")
 		} else if message.Role == "assistant" {
-			contentReasoning, content := splitQwen35ReasoningContent(content, message.Thinking, isThinking)
+			renderAssistantThinkBlock := r.alwaysRenderAssistantThinkBlock || (isThinking && i > lastQueryIndex)
+			contentReasoning, content := splitQwen35ReasoningContent(content, message.Thinking, renderAssistantThinkBlock)
 
-			if isThinking && i > lastQueryIndex {
+			if renderAssistantThinkBlock {
 				sb.WriteString(imStartTag + message.Role + "\n<think>\n" + contentReasoning + "\n</think>\n\n" + content)
 			} else {
 				sb.WriteString(imStartTag + message.Role + "\n" + content)

model/renderers/renderer.go

@@ -69,6 +69,8 @@ func rendererForName(name string) Renderer {
 	case "qwen3.5":
 		renderer := &Qwen35Renderer{isThinking: true, emitEmptyThinkOnNoThink: true, useImgTags: RenderImgTags}
 		return renderer
+	case "ornith":
+		return newOrnithRenderer()
 	case "cogito":
 		renderer := &CogitoRenderer{isThinking: true}
 		return renderer

scripts/build_darwin.sh

@@ -38,6 +38,11 @@ done
 shift $(( $OPTIND - 1 ))
 
 _build_darwin() {
+    BUILD_CPUS=$(getconf _NPROCESSORS_ONLN)
+    BUILD_JOBS=${OLLAMA_BUILD_PARALLEL:-$BUILD_CPUS}
+    BUILD_LOAD=${OLLAMA_BUILD_LOAD:-$BUILD_CPUS}
+    status "Build parallelism: $BUILD_JOBS, load limit: $BUILD_LOAD"
+
     SOURCE_BUILD=build/darwin-sources
     status "Preparing shared native sources"
     cmake -S . -B "$SOURCE_BUILD" -DOLLAMA_MLX_BACKENDS=metal_v3 -DOLLAMA_LLAMA_BACKENDS=
@@ -81,7 +86,7 @@ _build_darwin() {
             $MLX_EXTRA_ARGS
 
         GOOS=darwin GOARCH=$ARCH CGO_ENABLED=1 CGO_CFLAGS="$MLX_CGO_CFLAGS" CGO_LDFLAGS="$MLX_CGO_LDFLAGS" \
-            cmake --build "$BUILD_DIR" --target ollama-local --target ollama-mlx-backends --parallel
+            cmake --build "$BUILD_DIR" --target ollama-local --target ollama-mlx-backends --parallel "$BUILD_JOBS" -- -l "$BUILD_LOAD"
     done
 }
 

server/create.go

@@ -11,6 +11,7 @@ import (
 	"io/fs"
 	"log/slog"
 	"maps"
+	"math"
 	"net"
 	"net/http"
 	"net/url"
@@ -268,34 +269,65 @@ func (s *Server) CreateHandler(c *gin.Context) {
 				}
 			}
 
-			strFromInfo := func(k string) string {
+			strFromInfo := func(k string) (string, error) {
 				v, ok := r.Info[k]
 				if ok {
-					val := v.(string)
-					return val
+					val, ok := v.(string)
+					if !ok {
+						return "", fmt.Errorf("info field %q must be a string", k)
+					}
+					return val, nil
 				}
-				return ""
+				return "", nil
 			}
 
-			vFromInfo := func(k string) float64 {
+			intFromInfo := func(k string) (int, error) {
 				v, ok := r.Info[k]
 				if ok {
-					val := v.(float64)
-					return val
+					val, ok := v.(float64)
+					if !ok {
+						return 0, fmt.Errorf("info field %q must be a number", k)
+					}
+					if val < 0 || math.Trunc(val) != val || val > float64(maxCreateInfoInt()) {
+						return 0, fmt.Errorf("info field %q must be a non-negative integer", k)
+					}
+					return int(val), nil
 				}
-				return 0
+				return 0, nil
 			}
 
-			config.ModelFamily = strFromInfo("model_family")
+			if config.ModelFamily, err = strFromInfo("model_family"); err != nil {
+				ch <- gin.H{"error": err.Error(), "status": http.StatusBadRequest}
+				return
+			}
 			if config.ModelFamily != "" {
 				config.ModelFamilies = []string{config.ModelFamily}
 			}
 
-			config.BaseName = strFromInfo("base_name")
-			config.FileType = strFromInfo("quantization_level")
-			config.ModelType = strFromInfo("parameter_size")
-			config.ContextLen = int(vFromInfo("context_length"))
-			config.EmbedLen = int(vFromInfo("embedding_length"))
+			if config.BaseName, err = strFromInfo("base_name"); err != nil {
+				ch <- gin.H{"error": err.Error(), "status": http.StatusBadRequest}
+				return
+			}
+			if config.FileType, err = strFromInfo("quantization_level"); err != nil {
+				ch <- gin.H{"error": err.Error(), "status": http.StatusBadRequest}
+				return
+			}
+			if config.ModelType, err = strFromInfo("parameter_size"); err != nil {
+				ch <- gin.H{"error": err.Error(), "status": http.StatusBadRequest}
+				return
+			}
+			contextLen, err := intFromInfo("context_length")
+			if err != nil {
+				ch <- gin.H{"error": err.Error(), "status": http.StatusBadRequest}
+				return
+			}
+			config.ContextLen = contextLen
+			embedLen, err := intFromInfo("embedding_length")
+			if err != nil {
+				ch <- gin.H{"error": err.Error(), "status": http.StatusBadRequest}
+				return
+			}
+			config.EmbedLen = embedLen
 		}
 
 		if err := createModel(r, name, baseLayers, config, fn); err != nil {
@@ -440,6 +472,10 @@ func convertModelFromFilesWithMediaType(files map[string]string, baseLayers []*l
 	}
 }
 
+func maxCreateInfoInt() int {
+	return int(^uint(0) >> 1)
+}
+
 func detectModelTypeFromFiles(files map[string]string) string {
 	for fn := range files {
 		if strings.HasSuffix(fn, ".safetensors") {

server/internal/cache/blob/cache.go

@@ -1,544 +0,0 @@
-// Package blob implements a content-addressable disk cache for blobs and
-// manifests.
-package blob
-
-import (
-	"bytes"
-	"crypto/sha256"
-	"errors"
-	"fmt"
-	"hash"
-	"io"
-	"io/fs"
-	"iter"
-	"os"
-	"path/filepath"
-	"strings"
-	"time"
-
-	"github.com/ollama/ollama/server/internal/internal/names"
-)
-
-// Entry contains metadata about a blob in the cache.
-type Entry struct {
-	Digest Digest
-	Size   int64
-	Time   time.Time // when added to the cache
-}
-
-// DiskCache caches blobs and manifests on disk.
-//
-// The cache is rooted at a directory, which is created if it does not exist.
-//
-// Blobs are stored in the "blobs" subdirectory, and manifests are stored in the
-// "manifests" subdirectory. A example directory structure might look like:
-//
-//	<dir>/
-//	  blobs/
-//	    sha256-<digest> - <blob data>
-//	  manifests/
-//	    <host>/
-//	      <namespace>/
-//	        <name>/
-//	          <tag> - <manifest data>
-//
-// The cache is safe for concurrent use.
-//
-// Name casing is preserved in the cache, but is not significant when resolving
-// names. For example, "Foo" and "foo" are considered the same name.
-//
-// The cache is not safe for concurrent use. It guards concurrent writes, but
-// does not prevent duplicated effort. Because blobs are immutable, duplicate
-// writes should result in the same file being written to disk.
-type DiskCache struct {
-	// Dir specifies the top-level directory where blobs and manifest
-	// pointers are stored.
-	dir string
-	now func() time.Time
-
-	testHookBeforeFinalWrite func(f *os.File)
-}
-
-// PutBytes is a convenience function for c.Put(d, strings.NewReader(s), int64(len(s))).
-func PutBytes[S string | []byte](c *DiskCache, d Digest, data S) error {
-	return c.Put(d, bytes.NewReader([]byte(data)), int64(len(data)))
-}
-
-// Open opens a cache rooted at the given directory. If the directory does not
-// exist, it is created. If the directory is not a directory, an error is
-// returned.
-func Open(dir string) (*DiskCache, error) {
-	if dir == "" {
-		return nil, errors.New("blob: empty directory name")
-	}
-
-	info, err := os.Stat(dir)
-	if err == nil && !info.IsDir() {
-		return nil, fmt.Errorf("%q is not a directory", dir)
-	}
-	if err := os.MkdirAll(dir, 0o777); err != nil {
-		return nil, err
-	}
-
-	subdirs := []string{"blobs", "manifests"}
-	for _, subdir := range subdirs {
-		if err := os.MkdirAll(filepath.Join(dir, subdir), 0o777); err != nil {
-			return nil, err
-		}
-	}
-
-	// TODO(bmizerany): support shards
-	c := &DiskCache{
-		dir: dir,
-		now: time.Now,
-	}
-	return c, nil
-}
-
-func readAndSum(filename string, limit int64) (data []byte, _ Digest, err error) {
-	f, err := os.Open(filename)
-	if err != nil {
-		return nil, Digest{}, err
-	}
-	defer f.Close()
-
-	h := sha256.New()
-	r := io.TeeReader(f, h)
-	data, err = io.ReadAll(io.LimitReader(r, limit))
-	if err != nil {
-		return nil, Digest{}, err
-	}
-	var d Digest
-	h.Sum(d.sum[:0])
-	return data, d, nil
-}
-
-//lint:ignore U1000 used for debugging purposes as needed in tests
-var debug = false
-
-// debugger returns a function that can be used to add a step to the error message.
-// The error message will be a list of steps that were taken before the error occurred.
-// The steps are added in the order they are called.
-//
-// To set the error message, call the returned function with an empty string.
-//
-//lint:ignore U1000 used for debugging purposes as needed in tests
-func debugger(err *error) func(step string) {
-	if !debug {
-		return func(string) {}
-	}
-	var steps []string
-	return func(step string) {
-		if step == "" && *err != nil {
-			*err = fmt.Errorf("%q: %w", steps, *err)
-			return
-		}
-		steps = append(steps, step)
-		if len(steps) > 100 {
-			// shift hints in case of a bug that causes a lot of hints
-			copy(steps, steps[1:])
-			steps = steps[:100]
-		}
-	}
-}
-
-// Resolve resolves a name to a digest. The name is expected to
-// be in either of the following forms:
-//
-//	@<digest>
-//	<name>@<digest>
-//	<name>
-//
-// If a digest is provided, it is returned as is and nothing else happens.
-//
-// If a name is provided for a manifest that exists in the cache, the digest
-// of the manifest is returned. If there is no manifest in the cache, it
-// returns [fs.ErrNotExist].
-//
-// To cover the case where a manifest may change without the cache knowing
-// (e.g. it was reformatted or modified by hand), the manifest data read and
-// hashed is passed to a PutBytes call to ensure that the manifest is in the
-// blob store. This is done to ensure that future calls to [Get] succeed in
-// these cases.
-func (c *DiskCache) Resolve(name string) (Digest, error) {
-	name, digest := splitNameDigest(name)
-	if digest != "" {
-		return ParseDigest(digest)
-	}
-
-	// We want to address manifests files by digest using Get. That requires
-	// them to be blobs. This cannot be directly accomplished by looking in
-	// the blob store because manifests can change without Ollama knowing
-	// (e.g. a user modifies a manifests by hand then pushes it to update
-	// their model). We also need to support the blob caches inherited from
-	// older versions of Ollama, which do not store manifests in the blob
-	// store, so for these cases, we need to handle adding the manifests to
-	// the blob store, just in time.
-	//
-	// So now we read the manifests file, hash it, and copy it to the blob
-	// store if it's not already there.
-	//
-	// This should be cheap because manifests are small, and accessed
-	// infrequently.
-	file, err := c.manifestPath(name)
-	if err != nil {
-		return Digest{}, err
-	}
-
-	data, d, err := readAndSum(file, 1<<20)
-	if err != nil {
-		return Digest{}, err
-	}
-
-	// Ideally we'd read the "manifest" file as a manifest to the blob file,
-	// but we are not changing this yet, so copy the manifest to the blob
-	// store so it can be addressed by digest subsequent calls to Get.
-	if err := PutBytes(c, d, data); err != nil {
-		return Digest{}, err
-	}
-	return d, nil
-}
-
-// Put writes a new blob to the cache, identified by its digest. The operation
-// reads content from r, which must precisely match both the specified size and
-// digest.
-//
-// Concurrent write safety is achieved through file locking. The implementation
-// guarantees write integrity by enforcing size limits and content validation
-// before allowing the file to reach its final state.
-func (c *DiskCache) Put(d Digest, r io.Reader, size int64) error {
-	return c.copyNamedFile(c.GetFile(d), r, d, size)
-}
-
-// Import imports a blob from the provided reader into the cache. It reads the
-// entire content of the reader, calculates its digest, and stores it in the
-// cache.
-//
-// Import should be considered unsafe for use with untrusted data, such as data
-// read from a network. The caller is responsible for ensuring the integrity of
-// the data being imported.
-func (c *DiskCache) Import(r io.Reader, size int64) (Digest, error) {
-	// users that want to change the temp dir can set TEMPDIR.
-	f, err := os.CreateTemp("", "blob-")
-	if err != nil {
-		return Digest{}, err
-	}
-	defer os.Remove(f.Name())
-
-	// Copy the blob to a temporary file.
-	h := sha256.New()
-	r = io.TeeReader(r, h)
-	n, err := io.Copy(f, r)
-	if err != nil {
-		return Digest{}, err
-	}
-	if n != size {
-		return Digest{}, fmt.Errorf("blob: expected %d bytes, got %d", size, n)
-	}
-
-	// Check the digest.
-	var d Digest
-	h.Sum(d.sum[:0])
-	if err := f.Close(); err != nil {
-		return Digest{}, err
-	}
-	name := c.GetFile(d)
-	// Rename the temporary file to the final file.
-	if err := os.Rename(f.Name(), name); err != nil {
-		return Digest{}, err
-	}
-	os.Chtimes(name, c.now(), c.now()) // mainly for tests
-	return d, nil
-}
-
-// Get retrieves a blob from the cache using the provided digest. The operation
-// fails if the digest is malformed or if any errors occur during blob
-// retrieval.
-func (c *DiskCache) Get(d Digest) (Entry, error) {
-	name := c.GetFile(d)
-	info, err := os.Stat(name)
-	if err != nil {
-		return Entry{}, err
-	}
-	if info.Size() == 0 {
-		return Entry{}, fs.ErrNotExist
-	}
-	return Entry{
-		Digest: d,
-		Size:   info.Size(),
-		Time:   info.ModTime(),
-	}, nil
-}
-
-// Link creates a symbolic reference in the cache that maps the provided name
-// to a blob identified by its digest, making it retrievable by name using
-// [Resolve].
-//
-// It returns an error if either the name or digest is invalid, or if link
-// creation encounters any issues.
-func (c *DiskCache) Link(name string, d Digest) error {
-	manifest, err := c.manifestPath(name)
-	if err != nil {
-		return err
-	}
-	f, err := os.OpenFile(c.GetFile(d), os.O_RDONLY, 0)
-	if err != nil {
-		return err
-	}
-	defer f.Close()
-
-	// TODO(bmizerany): test this happens only if the blob was found to
-	// avoid leaving debris
-	if err := os.MkdirAll(filepath.Dir(manifest), 0o777); err != nil {
-		return err
-	}
-
-	info, err := f.Stat()
-	if err != nil {
-		return err
-	}
-
-	// Copy manifest to cache directory.
-	return c.copyNamedFile(manifest, f, d, info.Size())
-}
-
-// Unlink unlinks the manifest by name from the cache. If the name is not
-// found. If a manifest is removed ok will be true, otherwise false. If an
-// error occurs, it returns ok false, and the error.
-func (c *DiskCache) Unlink(name string) (ok bool, _ error) {
-	manifest, err := c.manifestPath(name)
-	if err != nil {
-		return false, err
-	}
-	err = os.Remove(manifest)
-	if errors.Is(err, fs.ErrNotExist) {
-		return false, nil
-	}
-	return true, err
-}
-
-// GetFile returns the absolute path to the file, in the cache, for the given
-// digest. It does not check if the file exists.
-//
-// The returned path should not be stored, used outside the lifetime of the
-// cache, or interpreted in any way.
-func (c *DiskCache) GetFile(d Digest) string {
-	filename := fmt.Sprintf("sha256-%x", d.sum)
-	return absJoin(c.dir, "blobs", filename)
-}
-
-// Links returns a sequence of link names. The sequence is in lexical order.
-// Names are converted from their relative path form to their name form but are
-// not guaranteed to be valid. Callers should validate the names before using.
-func (c *DiskCache) Links() iter.Seq2[string, error] {
-	return func(yield func(string, error) bool) {
-		for path, err := range c.links() {
-			if err != nil {
-				yield("", err)
-				return
-			}
-			if !yield(pathToName(path), nil) {
-				return
-			}
-		}
-	}
-}
-
-// pathToName converts a path to a name. It is the inverse of nameToPath. The
-// path is assumed to be in filepath.ToSlash format.
-func pathToName(s string) string {
-	s = strings.TrimPrefix(s, "manifests/")
-	rr := []rune(s)
-	for i := len(rr) - 1; i > 0; i-- {
-		if rr[i] == '/' {
-			rr[i] = ':'
-			return string(rr)
-		}
-	}
-	return s
-}
-
-// manifestPath finds the first manifest file on disk that matches the given
-// name using a case-insensitive comparison. If no manifest file is found, it
-// returns the path where the manifest file would be if it existed.
-//
-// If two manifest files exists on disk that match the given name using a
-// case-insensitive comparison, the one that sorts first, lexically, is
-// returned.
-func (c *DiskCache) manifestPath(name string) (string, error) {
-	np, err := nameToPath(name)
-	if err != nil {
-		return "", err
-	}
-
-	maybe := filepath.Join("manifests", np)
-	for l, err := range c.links() {
-		if err != nil {
-			return "", err
-		}
-		if strings.EqualFold(maybe, l) {
-			return filepath.Join(c.dir, l), nil
-		}
-	}
-	return filepath.Join(c.dir, maybe), nil
-}
-
-// links returns a sequence of links in the cache in lexical order.
-func (c *DiskCache) links() iter.Seq2[string, error] {
-	// TODO(bmizerany): reuse empty dirnames if exist
-	return func(yield func(string, error) bool) {
-		fsys := os.DirFS(c.dir)
-		manifests, err := fs.Glob(fsys, "manifests/*/*/*/*")
-		if err != nil {
-			yield("", err)
-			return
-		}
-		for _, manifest := range manifests {
-			if !yield(manifest, nil) {
-				return
-			}
-		}
-	}
-}
-
-type checkWriter struct {
-	size int64
-	d    Digest
-	f    *os.File
-	h    hash.Hash
-
-	w   io.Writer // underlying writer; set by creator
-	n   int64
-	err error
-
-	testHookBeforeFinalWrite func(*os.File)
-}
-
-func (w *checkWriter) seterr(err error) error {
-	if w.err == nil {
-		w.err = err
-	}
-	return err
-}
-
-// Write writes p to the underlying hash and writer. The last write to the
-// underlying writer is guaranteed to be the last byte of p as verified by the
-// hash.
-func (w *checkWriter) Write(p []byte) (int, error) {
-	if w.err != nil {
-		return 0, w.err
-	}
-
-	_, err := w.h.Write(p)
-	if err != nil {
-		return 0, w.seterr(err)
-	}
-	nextSize := w.n + int64(len(p))
-	if nextSize == w.size {
-		// last write. check hash.
-		sum := w.h.Sum(nil)
-		if !bytes.Equal(sum, w.d.sum[:]) {
-			return 0, w.seterr(fmt.Errorf("file content changed underfoot"))
-		}
-		if w.testHookBeforeFinalWrite != nil {
-			w.testHookBeforeFinalWrite(w.f)
-		}
-	}
-	if nextSize > w.size {
-		return 0, w.seterr(fmt.Errorf("content exceeds expected size: %d > %d", nextSize, w.size))
-	}
-	n, err := w.w.Write(p)
-	w.n += int64(n)
-	return n, w.seterr(err)
-}
-
-// copyNamedFile copies file into name, expecting it to have the given Digest
-// and size, if that file is not present already.
-func (c *DiskCache) copyNamedFile(name string, file io.Reader, out Digest, size int64) error {
-	info, err := os.Stat(name)
-	if err == nil && info.Size() == size {
-		// File already exists with correct size. This is good enough.
-		// We can skip expensive hash checks.
-		//
-		// TODO: Do the hash check, but give caller a way to skip it.
-		return nil
-	}
-
-	// Copy file to cache directory.
-	mode := os.O_RDWR | os.O_CREATE
-	if err == nil && info.Size() > size { // shouldn't happen but fix in case
-		mode |= os.O_TRUNC
-	}
-	f, err := os.OpenFile(name, mode, 0o666)
-	if err != nil {
-		return err
-	}
-	defer f.Close()
-	if size == 0 {
-		// File now exists with correct size.
-		// Only one possible zero-length file, so contents are OK too.
-		// Early return here makes sure there's a "last byte" for code below.
-		return nil
-	}
-
-	// From here on, if any of the I/O writing the file fails,
-	// we make a best-effort attempt to truncate the file f
-	// before returning, to avoid leaving bad bytes in the file.
-
-	// Copy file to f, but also into h to double-check hash.
-	cw := &checkWriter{
-		d:    out,
-		size: size,
-		h:    sha256.New(),
-		f:    f,
-		w:    f,
-
-		testHookBeforeFinalWrite: c.testHookBeforeFinalWrite,
-	}
-	n, err := io.Copy(cw, file)
-	if err != nil {
-		f.Truncate(0)
-		return err
-	}
-	if n < size {
-		f.Truncate(0)
-		return io.ErrUnexpectedEOF
-	}
-
-	if err := f.Close(); err != nil {
-		// Data might not have been written,
-		// but file may look like it is the right size.
-		// To be extra careful, remove cached file.
-		os.Remove(name)
-		return err
-	}
-	os.Chtimes(name, c.now(), c.now()) // mainly for tests
-
-	return nil
-}
-
-func splitNameDigest(s string) (name, digest string) {
-	i := strings.LastIndexByte(s, '@')
-	if i < 0 {
-		return s, ""
-	}
-	return s[:i], s[i+1:]
-}
-
-var errInvalidName = errors.New("invalid name")
-
-func nameToPath(name string) (_ string, err error) {
-	n := names.Parse(name)
-	if !n.IsFullyQualified() {
-		return "", errInvalidName
-	}
-	return filepath.Join(n.Host(), n.Namespace(), n.Model(), n.Tag()), nil
-}
-
-func absJoin(pp ...string) string {
-	abs, err := filepath.Abs(filepath.Join(pp...))
-	if err != nil {
-		panic(err) // this should never happen
-	}
-	return abs
-}

server/internal/cache/blob/cache_test.go

@@ -1,688 +0,0 @@
-package blob
-
-import (
-	"crypto/sha256"
-	"errors"
-	"fmt"
-	"io"
-	"io/fs"
-	"os"
-	"path/filepath"
-	"slices"
-	"strings"
-	"testing"
-	"time"
-
-	"github.com/ollama/ollama/server/internal/testutil"
-)
-
-func init() {
-	debug = true
-}
-
-var epoch = func() time.Time {
-	d := time.Date(2021, 1, 1, 0, 0, 0, 0, time.UTC)
-	if d.IsZero() {
-		panic("time zero")
-	}
-	return d
-}()
-
-func TestOpenErrors(t *testing.T) {
-	exe, err := os.Executable()
-	if err != nil {
-		panic(err)
-	}
-
-	cases := []struct {
-		dir string
-		err string
-	}{
-		{t.TempDir(), ""},
-		{"", "empty directory name"},
-		{exe, "not a directory"},
-	}
-
-	for _, tt := range cases {
-		t.Run(tt.dir, func(t *testing.T) {
-			_, err := Open(tt.dir)
-			if tt.err == "" {
-				if err != nil {
-					t.Fatal(err)
-				}
-				return
-			}
-			if err == nil {
-				t.Fatal("expected error")
-			}
-			if !strings.Contains(err.Error(), tt.err) {
-				t.Fatalf("err = %v, want %q", err, tt.err)
-			}
-		})
-	}
-}
-
-func TestGetFile(t *testing.T) {
-	t.Chdir(t.TempDir())
-
-	c, err := Open(".")
-	if err != nil {
-		t.Fatal(err)
-	}
-
-	d := mkdigest("1")
-	got := c.GetFile(d)
-	cleaned := filepath.Clean(got)
-	if cleaned != got {
-		t.Fatalf("got is unclean: %q", got)
-	}
-	if !filepath.IsAbs(got) {
-		t.Fatal("got is not absolute")
-	}
-	abs, _ := filepath.Abs(c.dir)
-	if !strings.HasPrefix(got, abs) {
-		t.Fatalf("got is not local to %q", c.dir)
-	}
-}
-
-func TestBasic(t *testing.T) {
-	c, err := Open(t.TempDir())
-	if err != nil {
-		t.Fatal(err)
-	}
-	now := epoch
-	c.now = func() time.Time { return now }
-
-	checkEntry := entryChecker(t, c)
-	checkFailed := func(err error) {
-		if err == nil {
-			t.Helper()
-			t.Fatal("expected error")
-		}
-	}
-
-	_, err = c.Resolve("invalid")
-	checkFailed(err)
-
-	_, err = c.Resolve("h/n/m:t")
-	checkFailed(err)
-
-	dx := mkdigest("x")
-
-	d, err := c.Resolve(fmt.Sprintf("h/n/m:t@%s", dx))
-	if err != nil {
-		t.Fatal(err)
-	}
-	if d != dx {
-		t.Fatalf("d = %v, want %v", d, dx)
-	}
-
-	_, err = c.Get(Digest{})
-	checkFailed(err)
-
-	// not committed yet
-	_, err = c.Get(dx)
-	checkFailed(err)
-
-	err = PutBytes(c, dx, "!")
-	checkFailed(err)
-
-	err = PutBytes(c, dx, "x")
-	if err != nil {
-		t.Fatal(err)
-	}
-	checkEntry(dx, 1, now)
-
-	t0 := now
-	now = now.Add(1*time.Hour + 1*time.Minute)
-	err = PutBytes(c, dx, "x")
-	if err != nil {
-		t.Fatal(err)
-	}
-
-	// check not updated
-	checkEntry(dx, 1, t0)
-}
-
-type sleepFunc func(d time.Duration) time.Time
-
-func openTester(t *testing.T) (*DiskCache, sleepFunc) {
-	t.Helper()
-	c, err := Open(t.TempDir())
-	if err != nil {
-		t.Fatal(err)
-	}
-	now := epoch
-	c.now = func() time.Time { return now }
-	return c, func(d time.Duration) time.Time {
-		now = now.Add(d)
-		return now
-	}
-}
-
-func TestManifestPath(t *testing.T) {
-	check := testutil.Checker(t)
-
-	c, sleep := openTester(t)
-
-	d1 := mkdigest("1")
-	err := PutBytes(c, d1, "1")
-	check(err)
-
-	err = c.Link("h/n/m:t", d1)
-	check(err)
-
-	t0 := sleep(0)
-	sleep(1 * time.Hour)
-	err = c.Link("h/n/m:t", d1) // nop expected
-	check(err)
-
-	file := must(c.manifestPath("h/n/m:t"))
-	info, err := os.Stat(file)
-	check(err)
-	testutil.CheckTime(t, info.ModTime(), t0)
-}
-
-func TestManifestExistsWithoutBlob(t *testing.T) {
-	t.Chdir(t.TempDir())
-
-	check := testutil.Checker(t)
-
-	c, err := Open(".")
-	check(err)
-
-	checkEntry := entryChecker(t, c)
-
-	man := must(c.manifestPath("h/n/m:t"))
-	os.MkdirAll(filepath.Dir(man), 0o777)
-	testutil.WriteFile(t, man, "1")
-
-	got, err := c.Resolve("h/n/m:t")
-	check(err)
-
-	want := mkdigest("1")
-	if got != want {
-		t.Fatalf("got = %v, want %v", got, want)
-	}
-
-	e, err := c.Get(got)
-	check(err)
-	checkEntry(got, 1, e.Time)
-}
-
-func TestPut(t *testing.T) {
-	c, sleep := openTester(t)
-
-	check := testutil.Checker(t)
-	checkEntry := entryChecker(t, c)
-
-	d := mkdigest("hello, world")
-	err := PutBytes(c, d, "hello")
-	if err == nil {
-		t.Fatal("expected error")
-	}
-
-	got, err := c.Get(d)
-	if !errors.Is(err, fs.ErrNotExist) {
-		t.Fatalf("expected error, got %v", got)
-	}
-
-	// Put a valid blob
-	err = PutBytes(c, d, "hello, world")
-	check(err)
-	checkEntry(d, 12, sleep(0))
-
-	// Put a blob with content that does not hash to the digest
-	err = PutBytes(c, d, "hello")
-	if err == nil {
-		t.Fatal("expected error")
-	}
-	checkNotExists(t, c, d)
-
-	// Put the valid blob back and check it
-	err = PutBytes(c, d, "hello, world")
-	check(err)
-	checkEntry(d, 12, sleep(0))
-
-	// Put a blob that errors during Read
-	err = c.Put(d, &errOnBangReader{s: "!"}, 1)
-	if err == nil {
-		t.Fatal("expected error")
-	}
-	checkNotExists(t, c, d)
-
-	// Put valid blob back and check it
-	err = PutBytes(c, d, "hello, world")
-	check(err)
-	checkEntry(d, 12, sleep(0))
-
-	// Put a blob with mismatched size
-	err = c.Put(d, strings.NewReader("hello, world"), 11)
-	if err == nil {
-		t.Fatal("expected error")
-	}
-	checkNotExists(t, c, d)
-
-	// Final byte does not match the digest (testing commit phase)
-	err = PutBytes(c, d, "hello, world$")
-	if err == nil {
-		t.Fatal("expected error")
-	}
-	checkNotExists(t, c, d)
-
-	reset := c.setTestHookBeforeFinalWrite(func(f *os.File) {
-		// change mode to read-only
-		f.Truncate(0)
-		f.Chmod(0o400)
-		f.Close()
-		f1, err := os.OpenFile(f.Name(), os.O_RDONLY, 0)
-		if err != nil {
-			t.Fatal(err)
-		}
-		t.Cleanup(func() { f1.Close() })
-		*f = *f1
-	})
-	defer reset()
-
-	err = PutBytes(c, d, "hello, world")
-	if err == nil {
-		t.Fatal("expected error")
-	}
-	checkNotExists(t, c, d)
-	reset()
-}
-
-func TestImport(t *testing.T) {
-	c, _ := openTester(t)
-
-	checkEntry := entryChecker(t, c)
-
-	want := mkdigest("x")
-	got, err := c.Import(strings.NewReader("x"), 1)
-	if err != nil {
-		t.Fatal(err)
-	}
-	if want != got {
-		t.Fatalf("digest = %v, want %v", got, want)
-	}
-	checkEntry(want, 1, epoch)
-
-	got, err = c.Import(strings.NewReader("x"), 1)
-	if err != nil {
-		t.Fatal(err)
-	}
-	if want != got {
-		t.Fatalf("digest = %v, want %v", got, want)
-	}
-	checkEntry(want, 1, epoch)
-}
-
-func (c *DiskCache) setTestHookBeforeFinalWrite(h func(*os.File)) (reset func()) {
-	old := c.testHookBeforeFinalWrite
-	c.testHookBeforeFinalWrite = h
-	return func() { c.testHookBeforeFinalWrite = old }
-}
-
-func TestPutGetZero(t *testing.T) {
-	c, sleep := openTester(t)
-
-	check := testutil.Checker(t)
-	checkEntry := entryChecker(t, c)
-
-	d := mkdigest("x")
-	err := PutBytes(c, d, "x")
-	check(err)
-	checkEntry(d, 1, sleep(0))
-
-	err = os.Truncate(c.GetFile(d), 0)
-	check(err)
-
-	_, err = c.Get(d)
-	if !errors.Is(err, fs.ErrNotExist) {
-		t.Fatalf("err = %v, want fs.ErrNotExist", err)
-	}
-}
-
-func TestPutZero(t *testing.T) {
-	c, _ := openTester(t)
-	d := mkdigest("x")
-	err := c.Put(d, strings.NewReader("x"), 0) // size == 0 (not size of content)
-	testutil.Check(t, err)
-	checkNotExists(t, c, d)
-}
-
-func TestCommit(t *testing.T) {
-	check := testutil.Checker(t)
-
-	c, err := Open(t.TempDir())
-	if err != nil {
-		t.Fatal(err)
-	}
-	checkEntry := entryChecker(t, c)
-
-	now := epoch
-	c.now = func() time.Time { return now }
-
-	d1 := mkdigest("1")
-	err = c.Link("h/n/m:t", d1)
-	if !errors.Is(err, fs.ErrNotExist) {
-		t.Fatalf("err = %v, want fs.ErrNotExist", err)
-	}
-
-	err = PutBytes(c, d1, "1")
-	check(err)
-
-	err = c.Link("h/n/m:t", d1)
-	check(err)
-
-	got, err := c.Resolve("h/n/m:t")
-	check(err)
-	if got != d1 {
-		t.Fatalf("d = %v, want %v", got, d1)
-	}
-
-	// commit again, more than 1 byte
-	d2 := mkdigest("22")
-	err = PutBytes(c, d2, "22")
-	check(err)
-	err = c.Link("h/n/m:t", d2)
-	check(err)
-	checkEntry(d2, 2, now)
-
-	filename := must(c.manifestPath("h/n/m:t"))
-	data, err := os.ReadFile(filename)
-	check(err)
-	if string(data) != "22" {
-		t.Fatalf("data = %q, want %q", data, "22")
-	}
-
-	t0 := now
-	now = now.Add(1 * time.Hour)
-	err = c.Link("h/n/m:t", d2) // same contents; nop
-	check(err)
-	info, err := os.Stat(filename)
-	check(err)
-	testutil.CheckTime(t, info.ModTime(), t0)
-}
-
-func TestManifestInvalidBlob(t *testing.T) {
-	c, _ := openTester(t)
-	d := mkdigest("1")
-	err := c.Link("h/n/m:t", d)
-	if err == nil {
-		t.Fatal("expected error")
-	}
-	checkNotExists(t, c, d)
-
-	err = PutBytes(c, d, "1")
-	testutil.Check(t, err)
-	err = os.WriteFile(c.GetFile(d), []byte("invalid"), 0o666)
-	if err != nil {
-		t.Fatal(err)
-	}
-
-	err = c.Link("h/n/m:t", d)
-	if !strings.Contains(err.Error(), "underfoot") {
-		t.Fatalf("err = %v, want error to contain %q", err, "underfoot")
-	}
-}
-
-func TestManifestNameReuse(t *testing.T) {
-	t.Run("case-insensitive", func(t *testing.T) {
-		// This should run on all file system types.
-		testManifestNameReuse(t)
-	})
-	t.Run("case-sensitive", func(t *testing.T) {
-		useCaseInsensitiveTempDir(t)
-		testManifestNameReuse(t)
-	})
-}
-
-func testManifestNameReuse(t *testing.T) {
-	check := testutil.Checker(t)
-
-	c, _ := openTester(t)
-
-	d1 := mkdigest("1")
-	err := PutBytes(c, d1, "1")
-	check(err)
-	err = c.Link("h/n/m:t", d1)
-	check(err)
-
-	d2 := mkdigest("22")
-	err = PutBytes(c, d2, "22")
-	check(err)
-	err = c.Link("H/N/M:T", d2)
-	check(err)
-
-	var g [2]Digest
-	g[0], err = c.Resolve("h/n/m:t")
-	check(err)
-	g[1], err = c.Resolve("H/N/M:T")
-	check(err)
-
-	w := [2]Digest{d2, d2}
-	if g != w {
-		t.Fatalf("g = %v, want %v", g, w)
-	}
-
-	var got []string
-	for l, err := range c.links() {
-		if err != nil {
-			t.Fatal(err)
-		}
-		got = append(got, l)
-	}
-	want := []string{"manifests/h/n/m/t"}
-	if !slices.Equal(got, want) {
-		t.Fatalf("got = %v, want %v", got, want)
-	}
-
-	// relink with different case
-	unlinked, err := c.Unlink("h/n/m:t")
-	check(err)
-	if !unlinked {
-		t.Fatal("expected unlinked")
-	}
-	err = c.Link("h/n/m:T", d1)
-	check(err)
-
-	got = got[:0]
-	for l, err := range c.links() {
-		if err != nil {
-			t.Fatal(err)
-		}
-		got = append(got, l)
-	}
-
-	// we should have only one link that is same case as the last link
-	want = []string{"manifests/h/n/m/T"}
-	if !slices.Equal(got, want) {
-		t.Fatalf("got = %v, want %v", got, want)
-	}
-}
-
-func TestManifestFile(t *testing.T) {
-	cases := []struct {
-		in   string
-		want string
-	}{
-		{"", ""},
-
-		// valid names
-		{"h/n/m:t", "/manifests/h/n/m/t"},
-		{"hh/nn/mm:tt", "/manifests/hh/nn/mm/tt"},
-
-		{"%/%/%/%", ""},
-
-		// already a path
-		{"h/n/m/t", ""},
-
-		// refs are not names
-		{"h/n/m:t@sha256-1", ""},
-		{"m@sha256-1", ""},
-		{"n/m:t@sha256-1", ""},
-	}
-
-	c, _ := openTester(t)
-	for _, tt := range cases {
-		t.Run(tt.in, func(t *testing.T) {
-			got, err := c.manifestPath(tt.in)
-			if err != nil && tt.want != "" {
-				t.Fatalf("unexpected error: %v", err)
-			}
-			if err == nil && tt.want == "" {
-				t.Fatalf("expected error")
-			}
-			dir := filepath.ToSlash(c.dir)
-			got = filepath.ToSlash(got)
-			got = strings.TrimPrefix(got, dir)
-			if got != tt.want {
-				t.Fatalf("got = %q, want %q", got, tt.want)
-			}
-		})
-	}
-}
-
-func TestNames(t *testing.T) {
-	c, _ := openTester(t)
-	check := testutil.Checker(t)
-
-	check(PutBytes(c, mkdigest("1"), "1"))
-	check(PutBytes(c, mkdigest("2"), "2"))
-
-	check(c.Link("h/n/m:t", mkdigest("1")))
-	check(c.Link("h/n/m:u", mkdigest("2")))
-
-	var got []string
-	for l, err := range c.Links() {
-		if err != nil {
-			t.Fatal(err)
-		}
-		got = append(got, l)
-	}
-	want := []string{"h/n/m:t", "h/n/m:u"}
-	if !slices.Equal(got, want) {
-		t.Fatalf("got = %v, want %v", got, want)
-	}
-}
-
-func mkdigest(s string) Digest {
-	return Digest{sha256.Sum256([]byte(s))}
-}
-
-func checkNotExists(t *testing.T, c *DiskCache, d Digest) {
-	t.Helper()
-	_, err := c.Get(d)
-	if !errors.Is(err, fs.ErrNotExist) {
-		t.Fatalf("err = %v, want fs.ErrNotExist", err)
-	}
-}
-
-func entryChecker(t *testing.T, c *DiskCache) func(Digest, int64, time.Time) {
-	t.Helper()
-	return func(d Digest, size int64, mod time.Time) {
-		t.Helper()
-		t.Run("checkEntry:"+d.String(), func(t *testing.T) {
-			t.Helper()
-
-			defer func() {
-				if t.Failed() {
-					dumpCacheContents(t, c)
-				}
-			}()
-
-			e, err := c.Get(d)
-			if size == 0 && errors.Is(err, fs.ErrNotExist) {
-				err = nil
-			}
-			if err != nil {
-				t.Fatal(err)
-			}
-			if e.Digest != d {
-				t.Errorf("e.Digest = %v, want %v", e.Digest, d)
-			}
-			if e.Size != size {
-				t.Fatalf("e.Size = %v, want %v", e.Size, size)
-			}
-
-			testutil.CheckTime(t, e.Time, mod)
-			info, err := os.Stat(c.GetFile(d))
-			if err != nil {
-				t.Fatal(err)
-			}
-			if info.Size() != size {
-				t.Fatalf("info.Size = %v, want %v", info.Size(), size)
-			}
-			testutil.CheckTime(t, info.ModTime(), mod)
-		})
-	}
-}
-
-func must[T any](v T, err error) T {
-	if err != nil {
-		panic(err)
-	}
-	return v
-}
-
-func TestNameToPath(t *testing.T) {
-	_, err := nameToPath("h/n/m:t")
-	if err != nil {
-		t.Fatal(err)
-	}
-}
-
-type errOnBangReader struct {
-	s string
-	n int
-}
-
-func (e *errOnBangReader) Read(p []byte) (int, error) {
-	if len(p) < 1 {
-		return 0, io.ErrShortBuffer
-	}
-	if e.n >= len(p) {
-		return 0, io.EOF
-	}
-	if e.s[e.n] == '!' {
-		return 0, errors.New("bang")
-	}
-	p[0] = e.s[e.n]
-	e.n++
-	return 1, nil
-}
-
-func dumpCacheContents(t *testing.T, c *DiskCache) {
-	t.Helper()
-
-	var b strings.Builder
-	fsys := os.DirFS(c.dir)
-	fs.WalkDir(fsys, ".", func(path string, d fs.DirEntry, err error) error {
-		t.Helper()
-
-		if err != nil {
-			return err
-		}
-		info, err := d.Info()
-		if err != nil {
-			return err
-		}
-
-		// Format like ls:
-		//
-		// ; ls -la
-		// drwxr-xr-x  224 Jan 13 14:22 blob/sha256-123
-		// drwxr-xr-x  224 Jan 13 14:22 manifest/h/n/m
-
-		fmt.Fprintf(&b, "    %s % 4d %s %s\n",
-			info.Mode(),
-			info.Size(),
-			info.ModTime().Format("Jan 2 15:04"),
-			path,
-		)
-		return nil
-	})
-	t.Log()
-	t.Logf("cache contents:\n%s", b.String())
-}

server/internal/cache/blob/casecheck_test.go

@@ -1,93 +0,0 @@
-package blob
-
-import (
-	"fmt"
-	"os"
-	"path/filepath"
-	"runtime"
-	"strings"
-	"testing"
-)
-
-func isCaseSensitive(dir string) bool {
-	defer func() {
-		os.Remove(filepath.Join(dir, "_casecheck"))
-	}()
-
-	exists := func(file string) bool {
-		_, err := os.Stat(file)
-		return err == nil
-	}
-
-	file := filepath.Join(dir, "_casecheck")
-	FILE := filepath.Join(dir, "_CASECHECK")
-	if exists(file) || exists(FILE) {
-		panic(fmt.Sprintf("_casecheck already exists in %q; remove and try again.", dir))
-	}
-
-	err := os.WriteFile(file, nil, 0o666)
-	if err != nil {
-		panic(err)
-	}
-
-	return !exists(FILE)
-}
-
-func isCI() bool {
-	return os.Getenv("CI") != ""
-}
-
-const volumeHint = `
-
-	Unable to locate case-insensitive TMPDIR on darwin.
-
-	To run tests, create the case-insensitive volume /Volumes/data:
-
-		$ sudo diskutil apfs addVolume disk1 APFSX data -mountpoint /Volumes/data
-
-	or run with:
-
-		CI=1 go test ./...
-
-`
-
-// useCaseInsensitiveTempDir sets TMPDIR to a case-insensitive directory
-// can find one, otherwise it skips the test if the CI environment variable is
-// set, or GOOS is not darwin.
-func useCaseInsensitiveTempDir(t *testing.T) bool {
-	if isCaseSensitive(os.TempDir()) {
-		// Use the default temp dir if it is already case-sensitive.
-		return true
-	}
-	if runtime.GOOS == "darwin" {
-		// If darwin, check for the special case-sensitive volume and
-		// use it if available.
-		const volume = "/Volumes/data"
-		_, err := os.Stat(volume)
-		if err == nil {
-			tmpdir := filepath.Join(volume, "tmp")
-			os.MkdirAll(tmpdir, 0o700)
-			t.Setenv("TMPDIR", tmpdir)
-			return true
-		}
-		if isCI() {
-			// Special case darwin in CI; it is not case-sensitive
-			// by default, and we will be testing other platforms
-			// that are case-sensitive, so we'll have the test
-			// being skipped covered there.
-			t.Skip("Skipping test in CI for darwin; TMPDIR is not case-insensitive.")
-		}
-	}
-
-	if !isCI() {
-		// Require devs to always tests with a case-insensitive TMPDIR.
-
-		// TODO(bmizerany): Print platform-specific instructions or
-		// link to docs on that topic.
-		lines := strings.Split(volumeHint, "\n")
-		for _, line := range lines {
-			t.Skip(line)
-		}
-	}
-	return false
-}

server/internal/cache/blob/chunked.go

@@ -1,73 +0,0 @@
-package blob
-
-import (
-	"crypto/sha256"
-	"errors"
-	"io"
-	"os"
-)
-
-// Chunk represents a range of bytes in a blob.
-type Chunk struct {
-	Start int64
-	End   int64
-}
-
-// Size returns end minus start plus one.
-func (c Chunk) Size() int64 {
-	return c.End - c.Start + 1
-}
-
-// Chunker writes to a blob in chunks.
-// Its zero value is invalid. Use [DiskCache.Chunked] to create a new Chunker.
-type Chunker struct {
-	digest Digest
-	size   int64
-	f      *os.File // nil means pre-validated
-}
-
-// Chunked returns a new Chunker, ready for use storing a blob of the given
-// size in chunks.
-//
-// Use [Chunker.Put] to write data to the blob at specific offsets.
-func (c *DiskCache) Chunked(d Digest, size int64) (*Chunker, error) {
-	name := c.GetFile(d)
-	info, err := os.Stat(name)
-	if err == nil && info.Size() == size {
-		return &Chunker{}, nil
-	}
-	f, err := os.OpenFile(name, os.O_CREATE|os.O_WRONLY, 0o666)
-	if err != nil {
-		return nil, err
-	}
-	return &Chunker{digest: d, size: size, f: f}, nil
-}
-
-// Put copies chunk.Size() bytes from r to the blob at the given offset,
-// merging the data with the existing blob. It returns an error if any. As a
-// special case, if r has less than chunk.Size() bytes, Put returns
-// io.ErrUnexpectedEOF.
-func (c *Chunker) Put(chunk Chunk, d Digest, r io.Reader) error {
-	if c.f == nil {
-		return nil
-	}
-
-	cw := &checkWriter{
-		d:    d,
-		size: chunk.Size(),
-		h:    sha256.New(),
-		f:    c.f,
-		w:    io.NewOffsetWriter(c.f, chunk.Start),
-	}
-
-	_, err := io.CopyN(cw, r, chunk.Size())
-	if err != nil && errors.Is(err, io.EOF) {
-		return io.ErrUnexpectedEOF
-	}
-	return err
-}
-
-// Close closes the underlying file.
-func (c *Chunker) Close() error {
-	return c.f.Close()
-}

server/internal/cache/blob/digest.go

@@ -1,99 +0,0 @@
-package blob
-
-import (
-	"crypto/sha256"
-	"encoding/hex"
-	"errors"
-	"fmt"
-	"slices"
-	"strings"
-)
-
-var ErrInvalidDigest = errors.New("invalid digest")
-
-// Digest is a blob identifier that is the SHA-256 hash of a blob's content.
-//
-// It is comparable and can be used as a map key.
-type Digest struct {
-	sum [32]byte
-}
-
-// ParseDigest parses a digest from a string. If the string is not a valid
-// digest, a call to the returned digest's IsValid method will return false.
-//
-// The input string may be in one of two forms:
-//
-//   - ("sha256-<hex>"), where <hex> is a 64-character hexadecimal string.
-//   - ("sha256:<hex>"), where <hex> is a 64-character hexadecimal string.
-//
-// The [Digest.String] method will return the canonical form of the
-// digest, "sha256:<hex>".
-func ParseDigest[S ~[]byte | ~string](v S) (Digest, error) {
-	s := string(v)
-	i := strings.IndexAny(s, ":-")
-	var zero Digest
-	if i < 0 {
-		return zero, ErrInvalidDigest
-	}
-
-	prefix, sum := s[:i], s[i+1:]
-	if prefix != "sha256" || len(sum) != 64 {
-		return zero, ErrInvalidDigest
-	}
-
-	var d Digest
-	_, err := hex.Decode(d.sum[:], []byte(sum))
-	if err != nil {
-		return zero, ErrInvalidDigest
-	}
-	return d, nil
-}
-
-func DigestFromBytes[S ~[]byte | ~string](v S) Digest {
-	return Digest{sha256.Sum256([]byte(v))}
-}
-
-// String returns the string representation of the digest in the conventional
-// form "sha256:<hex>".
-func (d Digest) String() string {
-	return fmt.Sprintf("sha256:%x", d.sum[:])
-}
-
-func (d Digest) Short() string {
-	return fmt.Sprintf("%x", d.sum[:4])
-}
-
-func (d Digest) Sum() [32]byte {
-	return d.sum
-}
-
-func (d Digest) Compare(other Digest) int {
-	return slices.Compare(d.sum[:], other.sum[:])
-}
-
-// IsValid returns true if the digest is valid, i.e. if it is the SHA-256 hash
-// of some content.
-func (d Digest) IsValid() bool {
-	return d != (Digest{})
-}
-
-// MarshalText implements the encoding.TextMarshaler interface. It returns an
-// error if [Digest.IsValid] returns false.
-func (d Digest) MarshalText() ([]byte, error) {
-	return []byte(d.String()), nil
-}
-
-// UnmarshalText implements the encoding.TextUnmarshaler interface, and only
-// works for a zero digest. If [Digest.IsValid] returns true, it returns an
-// error.
-func (d *Digest) UnmarshalText(text []byte) error {
-	if *d != (Digest{}) {
-		return errors.New("digest: illegal UnmarshalText on valid digest")
-	}
-	v, err := ParseDigest(string(text))
-	if err != nil {
-		return err
-	}
-	*d = v
-	return nil
-}

server/internal/cache/blob/digest_test.go

@@ -1,63 +0,0 @@
-package blob
-
-import (
-	"encoding/json"
-	"testing"
-)
-
-func TestParseDigest(t *testing.T) {
-	cases := []struct {
-		in    string
-		valid bool
-	}{
-		{"sha256-0123456789abcdef0123456789abcdef0123456789abcdef0123456789abcdef", true},
-		{"sha256:0123456789abcdef0123456789abcdef0123456789abcdef0123456789abcdef", true},
-
-		// too short
-		{"sha256-0123456789abcdef0123456789abcdef0123456789abcdef0123456789abcde", false},
-		{"sha256:0123456789abcdef0123456789abcdef0123456789abcdef0123456789abcde", false},
-
-		// too long
-		{"sha256-0123456789abcdef0123456789abcdef0123456789abcdef0123456789abcdef0", false},
-		{"sha256:0123456789abcdef0123456789abcdef0123456789abcdef0123456789abcdef0", false},
-
-		// invalid prefix
-		{"sha255-0123456789abcdef0123456789abcdef0123456789abcdef0123456789abcdef", false},
-		{"sha255:0123456789abcdef0123456789abcdef0123456789abcdef0123456789abcdef", false},
-		{"sha256!0123456789abcdef0123456789abcdef0123456789abcdef0123456789abcdef", false},
-
-		// invalid hex
-		{"sha256-XXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXX", false},
-		{"sha256:XXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXX", false},
-	}
-
-	for _, tt := range cases {
-		got, err := ParseDigest(tt.in)
-		if tt.valid && err != nil {
-			t.Errorf("ParseDigest(%q) = %v, %v; want valid", tt.in, got, err)
-		}
-		want := "sha256:" + tt.in[7:]
-		if tt.valid && got.String() != want {
-			t.Errorf("ParseDigest(%q).String() = %q, want %q", tt.in, got.String(), want)
-		}
-	}
-}
-
-func TestDigestMarshalText(t *testing.T) {
-	const s = `"sha256-0123456789abcdef0123456789abcdef0123456789abcdef0123456789abcdef"`
-	var d Digest
-	if err := json.Unmarshal([]byte(s), &d); err != nil {
-		t.Errorf("json.Unmarshal: %v", err)
-	}
-	out, err := json.Marshal(d)
-	if err != nil {
-		t.Errorf("json.Marshal: %v", err)
-	}
-	want := `"sha256:0123456789abcdef0123456789abcdef0123456789abcdef0123456789abcdef"`
-	if string(out) != want {
-		t.Errorf("json.Marshal: got %s, want %s", out, want)
-	}
-	if err := json.Unmarshal([]byte(`"invalid"`), &Digest{}); err == nil {
-		t.Errorf("json.Unmarshal: expected error")
-	}
-}

server/internal/client/ollama/registry.go

@@ -1,1197 +0,0 @@
-// Package ollama provides a client for interacting with an Ollama registry
-// which pushes and pulls model manifests and layers as defined by the
-// [ollama.com/manifest].
-package ollama
-
-import (
-	"bufio"
-	"bytes"
-	"cmp"
-	"context"
-	"crypto"
-	"crypto/ed25519"
-	"crypto/sha256"
-	"crypto/tls"
-	"encoding/base64"
-	"encoding/hex"
-	"encoding/json"
-	"errors"
-	"fmt"
-	"io"
-	"io/fs"
-	"iter"
-	"log/slog"
-	"net/http"
-	"os"
-	"path/filepath"
-	"runtime"
-	"runtime/debug"
-	"slices"
-	"strconv"
-	"strings"
-	"sync"
-	"sync/atomic"
-	"time"
-
-	"golang.org/x/crypto/ssh"
-	"golang.org/x/sync/errgroup"
-
-	"github.com/ollama/ollama/server/internal/cache/blob"
-	"github.com/ollama/ollama/server/internal/internal/names"
-
-	_ "embed"
-)
-
-// Errors
-var (
-	// ErrModelNotFound is returned when a manifest is not found in the
-	// cache or registry.
-	ErrModelNotFound = errors.New("model not found")
-
-	// ErrManifestInvalid is returned when a manifest found in a local or
-	// remote cache is invalid.
-	ErrManifestInvalid = errors.New("invalid manifest")
-
-	// ErrMissingModel is returned when the model part of a name is missing
-	// or invalid.
-	ErrNameInvalid = errors.New("invalid or missing name")
-
-	// ErrCached is passed to [Trace.PushUpdate] when a layer already
-	// exists. It is a non-fatal error and is never returned by [Registry.Push].
-	ErrCached = errors.New("cached")
-
-	// ErrIncomplete is returned by [Registry.Pull] when a model pull was
-	// incomplete due to one or more layer download failures. Users that
-	// want specific errors should use [WithTrace].
-	ErrIncomplete = errors.New("incomplete")
-)
-
-// Defaults
-const (
-	// DefaultChunkingThreshold is the threshold at which a layer should be
-	// split up into chunks when downloading.
-	DefaultChunkingThreshold = 64 << 20
-)
-
-var defaultCache = sync.OnceValues(func() (*blob.DiskCache, error) {
-	dir := os.Getenv("OLLAMA_MODELS")
-	if dir == "" {
-		home, _ := os.UserHomeDir()
-		home = cmp.Or(home, ".")
-		dir = filepath.Join(home, ".ollama", "models")
-	}
-	return blob.Open(dir)
-})
-
-// DefaultCache returns the default cache used by the registry. It is
-// configured from the OLLAMA_MODELS environment variable, or defaults to
-// $HOME/.ollama/models, or, if an error occurs obtaining the home directory,
-// it uses the current working directory.
-func DefaultCache() (*blob.DiskCache, error) {
-	return defaultCache()
-}
-
-// Error is the standard error returned by Ollama APIs. It can represent a
-// single or multiple error response.
-//
-// Single error responses have the following format:
-//
-//	{"code": "optional_code","error":"error message"}
-//
-// Multiple error responses have the following format:
-//
-//	{"errors": [{"code": "optional_code","message":"error message"}]}
-//
-// Note, that the error field is used in single error responses, while the
-// message field is used in multiple error responses.
-//
-// In both cases, the code field is optional and may be empty.
-type Error struct {
-	status  int    `json:"-"` // TODO(bmizerany): remove this
-	Code    string `json:"code"`
-	Message string `json:"message"`
-}
-
-// Temporary reports if the error is temporary (e.g. 5xx status code).
-func (e *Error) Temporary() bool {
-	return e.status >= 500
-}
-
-func (e *Error) Error() string {
-	var b strings.Builder
-	b.WriteString("registry responded with status ")
-	b.WriteString(strconv.Itoa(e.status))
-	if e.Code != "" {
-		b.WriteString(": code ")
-		b.WriteString(e.Code)
-	}
-	if e.Message != "" {
-		b.WriteString(": ")
-		b.WriteString(e.Message)
-	}
-	return b.String()
-}
-
-func (e *Error) LogValue() slog.Value {
-	return slog.GroupValue(
-		slog.Int("status", e.status),
-		slog.String("code", e.Code),
-		slog.String("message", e.Message),
-	)
-}
-
-// UnmarshalJSON implements json.Unmarshaler.
-func (e *Error) UnmarshalJSON(b []byte) error {
-	type E Error
-	var v struct {
-		// Single error
-		Code  string
-		Error string
-
-		// Multiple errors
-		Errors []E
-	}
-	if err := json.Unmarshal(b, &v); err != nil {
-		return err
-	}
-	if v.Error != "" {
-		// Single error case
-		e.Code = v.Code
-		e.Message = v.Error
-		return nil
-	}
-	if len(v.Errors) == 0 {
-		return fmt.Errorf("no messages in error response: %s", string(b))
-	}
-	*e = Error(v.Errors[0]) // our registry only returns one error.
-	return nil
-}
-
-const DefaultMask = "registry.ollama.ai/library/_:latest"
-
-var defaultMask = func() names.Name {
-	n := names.Parse(DefaultMask)
-	if !n.IsFullyQualified() {
-		panic("default mask is not fully qualified")
-	}
-	return n
-}()
-
-// CompleteName returns a fully qualified name by merging the given name with
-// the default mask. If the name is already fully qualified, it is returned
-// unchanged.
-func CompleteName(name string) string {
-	return names.Merge(names.Parse(name), defaultMask).String()
-}
-
-// Registry is a client for performing push and pull operations against an
-// Ollama registry.
-type Registry struct {
-	// Cache is the cache used to store models. If nil, [DefaultCache] is
-	// used.
-	Cache *blob.DiskCache
-
-	// UserAgent is the User-Agent header to send with requests to the
-	// registry. If empty, the User-Agent is determined by HTTPClient.
-	UserAgent string
-
-	// Key is the key used to authenticate with the registry.
-	//
-	// Currently, only Ed25519 keys are supported.
-	Key crypto.PrivateKey
-
-	// HTTPClient is the HTTP client used to make requests to the registry.
-	//
-	// If nil, [http.DefaultClient] is used.
-	//
-	// As a quick note: If a Registry function that makes a call to a URL
-	// with the "https+insecure" scheme, the client will be cloned and the
-	// transport will be set to skip TLS verification, unless the client's
-	// Transport done not have a Clone method with the same signature as
-	// [http.Transport.Clone], which case, the call will fail.
-	HTTPClient *http.Client
-
-	// MaxStreams is the maximum number of concurrent streams to use when
-	// pushing or pulling models. If zero, the number of streams is
-	// determined by [runtime.GOMAXPROCS].
-	//
-	// A negative value means no limit.
-	MaxStreams int
-
-	// ChunkingThreshold is the maximum size of a layer to download in a single
-	// request. If zero, [DefaultChunkingThreshold] is used.
-	ChunkingThreshold int64
-
-	// Mask, if set, is the name used to convert non-fully qualified names
-	// to fully qualified names.
-	// If empty, [DefaultMask] is used.
-	Mask string
-
-	// ReadTimeout is the maximum duration for reading the entire request,
-	// including the body.
-	// A zero or negative value means there will be no timeout.
-	ReadTimeout time.Duration
-}
-
-func (r *Registry) readTimeout() time.Duration {
-	if r.ReadTimeout > 0 {
-		return r.ReadTimeout
-	}
-	return 1<<63 - 1 // no timeout, max int
-}
-
-func (r *Registry) cache() (*blob.DiskCache, error) {
-	if r.Cache != nil {
-		return r.Cache, nil
-	}
-	return defaultCache()
-}
-
-func (r *Registry) parseName(name string) (names.Name, error) {
-	mask := defaultMask
-	if r.Mask != "" {
-		mask = names.Parse(r.Mask)
-	}
-	n := names.Merge(names.Parse(name), mask)
-	if !n.IsFullyQualified() {
-		return names.Name{}, fmt.Errorf("%w: %q", ErrNameInvalid, name)
-	}
-	return n, nil
-}
-
-// DefaultRegistry returns a new Registry configured from the environment. The
-// key is read from $HOME/.ollama/id_ed25519, MaxStreams is set to the
-// value of OLLAMA_REGISTRY_MAXSTREAMS, and ReadTimeout is set to 30 seconds.
-//
-// It returns an error if any configuration in the environment is invalid.
-func DefaultRegistry() (*Registry, error) {
-	home, err := os.UserHomeDir()
-	if err != nil {
-		return nil, err
-	}
-	keyPEM, err := os.ReadFile(filepath.Join(home, ".ollama/id_ed25519"))
-	if err != nil && errors.Is(err, fs.ErrNotExist) {
-		return nil, err
-	}
-
-	var rc Registry
-	rc.ReadTimeout = 30 * time.Second
-	rc.UserAgent = UserAgent()
-	rc.Key, err = ssh.ParseRawPrivateKey(keyPEM)
-	if err != nil {
-		return nil, err
-	}
-	maxStreams := os.Getenv("OLLAMA_REGISTRY_MAXSTREAMS")
-	if maxStreams != "" {
-		var err error
-		rc.MaxStreams, err = strconv.Atoi(maxStreams)
-		if err != nil {
-			return nil, fmt.Errorf("invalid OLLAMA_REGISTRY_MAXSTREAMS: %w", err)
-		}
-	}
-	return &rc, nil
-}
-
-func UserAgent() string {
-	buildinfo, _ := debug.ReadBuildInfo()
-
-	version := buildinfo.Main.Version
-	if version == "(devel)" {
-		// When using `go run .` the version is "(devel)". This is seen
-		// as an invalid version by ollama.com and so it defaults to
-		// "needs upgrade" for some requests, such as pulls. These
-		// checks can be skipped by using the special version "v0.0.0",
-		// so we set it to that here.
-		version = "v0.0.0"
-	}
-
-	return fmt.Sprintf("ollama/%s (%s %s) Go/%s",
-		version,
-		runtime.GOARCH,
-		runtime.GOOS,
-		runtime.Version(),
-	)
-}
-
-func (r *Registry) maxStreams() int {
-	return cmp.Or(r.MaxStreams, runtime.GOMAXPROCS(0))
-}
-
-func (r *Registry) maxChunkingThreshold() int64 {
-	return cmp.Or(r.ChunkingThreshold, DefaultChunkingThreshold)
-}
-
-type PushParams struct {
-	// From is an optional destination name for the model. If empty, the
-	// destination name is the same as the source name.
-	From string
-}
-
-// Push pushes the model with the name in the cache to the remote registry.
-func (r *Registry) Push(ctx context.Context, name string, p *PushParams) error {
-	if p == nil {
-		p = &PushParams{}
-	}
-
-	c, err := r.cache()
-	if err != nil {
-		return err
-	}
-
-	m, err := r.ResolveLocal(cmp.Or(p.From, name))
-	if err != nil {
-		return err
-	}
-
-	// Before much else happens, check layers at not null, the blobs exist,
-	// and the sizes match. This prevents long uploads followed by
-	// disappointment.
-	for _, l := range m.Layers {
-		if l == nil {
-			return fmt.Errorf("%w: null layer", ErrManifestInvalid)
-		}
-		info, err := c.Get(l.Digest)
-		if err != nil {
-			return fmt.Errorf("error getting %s: %w", l.Digest.Short(), err)
-		}
-		if info.Size != l.Size {
-			return fmt.Errorf("size mismatch for %s: %d != %d", l.Digest.Short(), info.Size, l.Size)
-		}
-	}
-
-	t := traceFromContext(ctx)
-
-	scheme, n, _, err := r.parseNameExtended(name)
-	if err != nil {
-		// This should never happen since ResolveLocal should have
-		// already validated the name.
-		panic(err)
-	}
-
-	ctx, cancel := context.WithCancel(ctx)
-	defer cancel()
-
-	var g errgroup.Group
-	g.SetLimit(r.maxStreams())
-	for _, l := range m.Layers {
-		var progress atomic.Int64
-		g.Go(func() (err error) {
-			defer func() { t.update(l, progress.Load(), err) }()
-
-			t.update(l, 0, nil)
-
-			startURL := fmt.Sprintf("%s://%s/v2/%s/%s/blobs/uploads/?digest=%s",
-				scheme,
-				n.Host(),
-				n.Namespace(),
-				n.Model(),
-				l.Digest,
-			)
-			res, err := r.send(ctx, "POST", startURL, nil)
-			if err != nil {
-				return err
-			}
-			res.Body.Close()
-
-			f, err := os.Open(c.GetFile(l.Digest))
-			if err != nil {
-				return err
-			}
-			defer f.Close()
-
-			uploadURL := res.Header.Get("Location")
-			if uploadURL == "" {
-				t.update(l, l.Size, ErrCached)
-				return nil
-			}
-
-			req, err := r.newRequest(ctx, "PUT", uploadURL, f)
-			if err != nil {
-				return fmt.Errorf("invalid upload URL returned from registry: %q: %w", uploadURL, err)
-			}
-			req.ContentLength = l.Size
-
-			res, err = sendRequest(r.client(), req)
-			if err == nil {
-				res.Body.Close()
-			}
-			return err
-		})
-	}
-
-	if err := g.Wait(); err != nil {
-		return err
-	}
-
-	// Commit
-	path := fmt.Sprintf("%s://%s/v2/%s/%s/manifests/%s",
-		scheme,
-		n.Host(),
-		n.Namespace(),
-		n.Model(),
-		n.Tag(),
-	)
-	res, err := r.send(ctx, "PUT", path, bytes.NewReader(m.Data))
-	if err == nil {
-		res.Body.Close()
-	}
-	// TODO(bmizerany): add a "commit" trace event
-	return err
-}
-
-// trackingReader is an io.Reader that tracks the number of bytes read and
-// calls the update function with the layer, the number of bytes read.
-//
-// It always calls update with a nil error.
-type trackingReader struct {
-	r      io.Reader
-	update func(n int64, err error) // err is always nil
-}
-
-func (r *trackingReader) Read(p []byte) (n int, err error) {
-	n, err = r.r.Read(p)
-	r.update(int64(n), nil)
-	return
-}
-
-// Pull pulls the model with the given name from the remote registry into the
-// cache.
-//
-// For layers larger then [Registry.MaxChunkSize], the layer is downloaded in
-// chunks of the specified size, and then reassembled and verified. This is
-// typically slower than splitting the model up across layers, and is mostly
-// utilized for layers of type equal to "application/vnd.ollama.image".
-func (r *Registry) Pull(ctx context.Context, name string) error {
-	m, err := r.Resolve(ctx, name)
-	if err != nil {
-		return err
-	}
-
-	// TODO(bmizerany): decide if this should be considered valid. Maybe
-	// server-side we special case '{}' to have some special meaning? Maybe
-	// "archiving" a tag (which is how we reason about it in the registry
-	// already, just with a different twist).
-	if len(m.Layers) == 0 {
-		return fmt.Errorf("%w: no layers", ErrManifestInvalid)
-	}
-
-	c, err := r.cache()
-	if err != nil {
-		return err
-	}
-
-	// TODO(bmizerany): work to remove the need to do this
-	layers := m.Layers
-	if m.Config != nil && m.Config.Digest.IsValid() {
-		layers = append(layers, m.Config)
-	}
-
-	// Send initial layer trace events to allow clients to have an
-	// understanding of work to be done before work starts.
-	var expected int64
-	t := traceFromContext(ctx)
-	for _, l := range layers {
-		t.update(l, 0, nil)
-		expected += l.Size
-	}
-
-	var g errgroup.Group
-	g.SetLimit(r.maxStreams())
-
-	var completed atomic.Int64
-	for _, l := range layers {
-		var received atomic.Int64
-		update := func(n int64, err error) {
-			if n == 0 && err == nil {
-				// Clients expect an update with no progress and no error to mean "starting download".
-				// This is not the case here,
-				// so we don't want to send an update in this case.
-				return
-			}
-			completed.Add(n)
-			t.update(l, received.Add(n), err)
-		}
-
-		info, err := c.Get(l.Digest)
-		if err == nil && info.Size == l.Size {
-			update(l.Size, ErrCached)
-			continue
-		}
-
-		func() (err error) {
-			defer func() {
-				if err != nil {
-					update(0, err)
-				}
-			}()
-
-			var wg sync.WaitGroup
-			chunked, err := c.Chunked(l.Digest, l.Size)
-			if err != nil {
-				return err
-			}
-			defer func() {
-				// Close the chunked writer when all chunks are
-				// downloaded.
-				//
-				// This is done as a background task in the
-				// group to allow the next layer to start while
-				// we wait for the final chunk in this layer to
-				// complete. It also ensures this is done
-				// before we exit Pull.
-				g.Go(func() error {
-					wg.Wait()
-					chunked.Close()
-					return nil
-				})
-			}()
-
-			for cs, err := range r.chunksums(ctx, name, l) {
-				if err != nil {
-					// Note the chunksum stream
-					// interruption, but do not cancel
-					// in-flight downloads. We can still
-					// make progress on them. Once they are
-					// done, ErrIncomplete will be returned
-					// below.
-					update(0, err)
-					break
-				}
-
-				cacheKey := fmt.Sprintf(
-					"v1 pull chunksum %s %s %d-%d",
-					l.Digest,
-					cs.Digest,
-					cs.Chunk.Start,
-					cs.Chunk.End,
-				)
-				cacheKeyDigest := blob.DigestFromBytes(cacheKey)
-				_, err := c.Get(cacheKeyDigest)
-				if err == nil {
-					update(cs.Chunk.Size(), ErrCached)
-					continue
-				}
-
-				wg.Add(1)
-				g.Go(func() (err error) {
-					defer func() {
-						defer wg.Done()
-						if err != nil {
-							update(0, err)
-						}
-					}()
-
-					ctx, cancel := context.WithCancelCause(ctx)
-					defer cancel(nil)
-
-					timer := time.AfterFunc(r.readTimeout(), func() {
-						cancel(fmt.Errorf("%w: downloading %s %d-%d/%d",
-							context.DeadlineExceeded,
-							cs.Digest.Short(),
-							cs.Chunk.Start,
-							cs.Chunk.End,
-							l.Size,
-						))
-					})
-					defer timer.Stop()
-
-					req, err := http.NewRequestWithContext(ctx, "GET", cs.URL, nil)
-					if err != nil {
-						return err
-					}
-					req.Header.Set("Range", fmt.Sprintf("bytes=%d-%d", cs.Chunk.Start, cs.Chunk.End))
-					res, err := sendRequest(r.client(), req)
-					if err != nil {
-						return err
-					}
-					defer res.Body.Close()
-
-					tr := &trackingReader{
-						r: res.Body,
-						update: func(n int64, err error) {
-							timer.Reset(r.readTimeout())
-							update(n, err)
-						},
-					}
-					if err := chunked.Put(cs.Chunk, cs.Digest, tr); err != nil {
-						return err
-					}
-
-					// Record the downloading of this chunk.
-					return blob.PutBytes(c, cacheKeyDigest, cacheKey)
-				})
-			}
-
-			return nil
-		}()
-	}
-	if err := g.Wait(); err != nil {
-		return err
-	}
-	if recv := completed.Load(); recv != expected {
-		return fmt.Errorf("%w: received %d/%d bytes", ErrIncomplete, recv, expected)
-	}
-
-	md := blob.DigestFromBytes(m.Data)
-	if err := blob.PutBytes(c, md, m.Data); err != nil {
-		return err
-	}
-	return c.Link(m.Name, md)
-}
-
-// Unlink is like [blob.DiskCache.Unlink], but makes name fully qualified
-// before attempting to unlink the model.
-func (r *Registry) Unlink(name string) (ok bool, _ error) {
-	n, err := r.parseName(name)
-	if err != nil {
-		return false, err
-	}
-	c, err := r.cache()
-	if err != nil {
-		return false, err
-	}
-	return c.Unlink(n.String())
-}
-
-// Manifest represents a [ollama.com/manifest].
-type Manifest struct {
-	Name   string   `json:"-"` // the canonical name of the model
-	Data   []byte   `json:"-"` // the raw data of the manifest
-	Layers []*Layer `json:"layers"`
-
-	// For legacy reasons, we still have to download the config layer.
-	Config *Layer `json:"config"`
-}
-
-// Layer returns the layer with the given
-// digest, or nil if not found.
-func (m *Manifest) Layer(d blob.Digest) *Layer {
-	for _, l := range m.Layers {
-		if l.Digest == d {
-			return l
-		}
-	}
-	return nil
-}
-
-func (m *Manifest) All() iter.Seq[*Layer] {
-	return func(yield func(*Layer) bool) {
-		if !yield(m.Config) {
-			return
-		}
-		for _, l := range m.Layers {
-			if !yield(l) {
-				return
-			}
-		}
-	}
-}
-
-func (m *Manifest) Size() int64 {
-	var size int64
-	if m.Config != nil {
-		size += m.Config.Size
-	}
-	for _, l := range m.Layers {
-		size += l.Size
-	}
-	return size
-}
-
-// MarshalJSON implements json.Marshaler.
-//
-// NOTE: It adds an empty config object to the manifest, which is required by
-// the registry, but not used by the client. In the future, the config object
-// will not be required by the registry and this will should be removed.
-func (m Manifest) MarshalJSON() ([]byte, error) {
-	type M Manifest
-	v := struct {
-		M
-
-		// This is ignored, mostly, by the registry But, if not
-		// present, it will cause an error to be returned during the
-		// last phase of the commit which expects it, but does nothing
-		// with it. This will be fixed in a future release of
-		// ollama.com.
-		Config Layer `json:"config"`
-	}{
-		M: M(m),
-	}
-	return json.Marshal(v)
-}
-
-// unmarshalManifest unmarshals the data into a manifest, and sets the name
-// field to the string representation of the name.
-//
-// It panics if the name is not fully qualified. Callers should ensure the name
-// is fully qualified before calling this function.
-func unmarshalManifest(n names.Name, data []byte) (*Manifest, error) {
-	if !n.IsFullyQualified() {
-		panic(fmt.Sprintf("unmarshalManifest: name is not fully qualified: %s", n.String()))
-	}
-	var m Manifest
-	if err := json.Unmarshal(data, &m); err != nil {
-		return nil, err
-	}
-	m.Name = n.String()
-	m.Data = data
-	return &m, nil
-}
-
-// Layer is a layer in a model.
-type Layer struct {
-	Digest    blob.Digest `json:"digest"`
-	MediaType string      `json:"mediaType"`
-	Size      int64       `json:"size"`
-}
-
-// ResolveLocal resolves a name to a Manifest in the local cache.
-func (r *Registry) ResolveLocal(name string) (*Manifest, error) {
-	_, n, d, err := r.parseNameExtended(name)
-	if err != nil {
-		return nil, err
-	}
-	c, err := r.cache()
-	if err != nil {
-		return nil, err
-	}
-	if !d.IsValid() {
-		// No digest, so resolve the manifest by name.
-		d, err = c.Resolve(n.String())
-		if err != nil {
-			return nil, err
-		}
-	}
-	data, err := os.ReadFile(c.GetFile(d))
-	if err != nil {
-		if errors.Is(err, fs.ErrNotExist) {
-			return nil, fmt.Errorf("%w: %s", ErrModelNotFound, name)
-		}
-		return nil, err
-	}
-	m, err := unmarshalManifest(n, data)
-	if err != nil {
-		return nil, fmt.Errorf("%s: %w", name, errors.Join(ErrManifestInvalid, err))
-	}
-	return m, nil
-}
-
-// Resolve resolves a name to a Manifest in the remote registry.
-func (r *Registry) Resolve(ctx context.Context, name string) (*Manifest, error) {
-	scheme, n, d, err := r.parseNameExtended(name)
-	if err != nil {
-		return nil, err
-	}
-
-	manifestURL := fmt.Sprintf("%s://%s/v2/%s/%s/manifests/%s", scheme, n.Host(), n.Namespace(), n.Model(), n.Tag())
-	if d.IsValid() {
-		manifestURL = fmt.Sprintf("%s://%s/v2/%s/%s/blobs/%s", scheme, n.Host(), n.Namespace(), n.Model(), d)
-	}
-
-	res, err := r.send(ctx, "GET", manifestURL, nil)
-	if err != nil {
-		return nil, err
-	}
-	defer res.Body.Close()
-	data, err := io.ReadAll(res.Body)
-	if err != nil {
-		return nil, err
-	}
-	// TODO(bmizerany): return digest here
-	m, err := unmarshalManifest(n, data)
-	if err != nil {
-		return nil, fmt.Errorf("%s: %w", name, errors.Join(ErrManifestInvalid, err))
-	}
-	return m, nil
-}
-
-type chunksum struct {
-	URL    string
-	Chunk  blob.Chunk
-	Digest blob.Digest
-}
-
-// chunksums returns a sequence of chunksums for the given layer. If the layer is under the
-// chunking threshold, a single chunksum is returned that covers the entire layer. If the layer
-// is over the chunking threshold, the chunksums are read from the chunksums endpoint.
-func (r *Registry) chunksums(ctx context.Context, name string, l *Layer) iter.Seq2[chunksum, error] {
-	return func(yield func(chunksum, error) bool) {
-		scheme, n, _, err := r.parseNameExtended(name)
-		if err != nil {
-			yield(chunksum{}, err)
-			return
-		}
-
-		if l.Size < r.maxChunkingThreshold() {
-			// any layer under the threshold should be downloaded
-			// in one go.
-			cs := chunksum{
-				URL: fmt.Sprintf("%s://%s/v2/%s/%s/blobs/%s",
-					scheme,
-					n.Host(),
-					n.Namespace(),
-					n.Model(),
-					l.Digest,
-				),
-				Chunk:  blob.Chunk{Start: 0, End: l.Size - 1},
-				Digest: l.Digest,
-			}
-			yield(cs, nil)
-			return
-		}
-
-		// The response is a sequence of chunksums.
-		//
-		// Chunksums are chunks of a larger blob that can be
-		// downloaded and verified independently.
-		//
-		// The chunksums endpoint is a GET request that returns a
-		// sequence of chunksums in the following format:
-		//
-		//     > GET /v2/<namespace>/<model>/chunksums/<digest>
-		//
-		//     < HTTP/1.1 200 OK
-		//     < Content-Location: <blobURL>
-		//     <
-		//     < <digest> <start>-<end>
-		//     < ...
-		//
-		// The <blobURL> is the URL to download the chunks from and
-		// each <digest> is the digest of the chunk, and <start>-<end>
-		// is the range the chunk in the blob.
-		//
-		// Ranges may be used directly in Range headers like
-		// "bytes=<start>-<end>".
-		//
-		// The chunksums returned are guaranteed to be contiguous and
-		// include all bytes of the layer. If the stream is cut short,
-		// clients should retry.
-
-		chunksumsURL := fmt.Sprintf("%s://%s/v2/%s/%s/chunksums/%s",
-			scheme,
-			n.Host(),
-			n.Namespace(),
-			n.Model(),
-			l.Digest,
-		)
-
-		req, err := r.newRequest(ctx, "GET", chunksumsURL, nil)
-		if err != nil {
-			yield(chunksum{}, err)
-			return
-		}
-		res, err := sendRequest(r.client(), req)
-		if err != nil {
-			yield(chunksum{}, err)
-			return
-		}
-		defer res.Body.Close()
-		if res.StatusCode != 200 {
-			err := fmt.Errorf("chunksums: unexpected status code %d", res.StatusCode)
-			yield(chunksum{}, err)
-			return
-		}
-		blobURL := res.Header.Get("Content-Location")
-
-		s := bufio.NewScanner(res.Body)
-		s.Split(bufio.ScanWords)
-		for {
-			if !s.Scan() {
-				if s.Err() != nil {
-					yield(chunksum{}, s.Err())
-				}
-				return
-			}
-			d, err := blob.ParseDigest(s.Bytes())
-			if err != nil {
-				yield(chunksum{}, fmt.Errorf("invalid digest: %q", s.Bytes()))
-				return
-			}
-
-			if !s.Scan() {
-				err := s.Err()
-				if err == nil {
-					err = fmt.Errorf("missing chunk range for digest %s", d)
-				}
-				yield(chunksum{}, err)
-				return
-			}
-			chunk, err := parseChunk(s.Bytes())
-			if err != nil {
-				yield(chunksum{}, fmt.Errorf("invalid chunk range for digest %s: %q", d, s.Bytes()))
-				return
-			}
-
-			cs := chunksum{
-				URL:    blobURL,
-				Chunk:  chunk,
-				Digest: d,
-			}
-			if !yield(cs, nil) {
-				return
-			}
-		}
-	}
-}
-
-func (r *Registry) client() *http.Client {
-	if r.HTTPClient != nil {
-		return r.HTTPClient
-	}
-	return http.DefaultClient
-}
-
-// newRequest constructs a new request, ready to use, with the given method,
-// url, and body, pre-signed with client [Key] and [UserAgent].
-func (r *Registry) newRequest(ctx context.Context, method, url string, body io.Reader) (*http.Request, error) {
-	req, err := http.NewRequestWithContext(ctx, method, url, body)
-	if err != nil {
-		return nil, err
-	}
-	if r.UserAgent != "" {
-		req.Header.Set("User-Agent", r.UserAgent)
-	}
-	if r.Key != nil {
-		token, err := makeAuthToken(r.Key)
-		if err != nil {
-			return nil, err
-		}
-		req.Header.Set("Authorization", "Bearer "+token)
-	}
-	return req, nil
-}
-
-// sendRequest makes a request with the given client and request, and returns the
-// response if the status code is 200. If the status code is not 200, an Error
-// is parsed from the response body and returned. If any other error occurs, it
-// is returned.
-func sendRequest(c *http.Client, r *http.Request) (_ *http.Response, err error) {
-	if r.URL.Scheme == "https+insecure" {
-		// TODO(bmizerany): clone client.Transport, set
-		// InsecureSkipVerify, etc.
-
-		type cloner interface {
-			Clone() *http.Transport
-		}
-
-		// Attempt to configure the transport to skip TLS verification
-		// if we can clone it, otherwise fall through and let the http
-		// client complain and the scheme being invalid.
-		x, ok := cmp.Or(c.Transport, http.DefaultTransport).(cloner)
-		if ok {
-			tr := x.Clone()
-			tr.TLSClientConfig = cmp.Or(tr.TLSClientConfig, &tls.Config{})
-			tr.TLSClientConfig.InsecureSkipVerify = true
-
-			cc := *c // shallow copy
-			cc.Transport = tr
-			c = &cc
-
-			r = r.Clone(r.Context())
-			r.URL.Scheme = "https"
-
-			// fall through
-		}
-	}
-
-	res, err := c.Do(r)
-	if err != nil {
-		return nil, err
-	}
-	if res.StatusCode/100 != 2 {
-		out, err := io.ReadAll(res.Body)
-		if err != nil {
-			return nil, err
-		}
-		var re Error
-		if err := json.Unmarshal(out, &re); err != nil {
-			// Use the raw body if we can't parse it as an error object.
-			re.Message = string(out)
-		}
-
-		// coerce MANIFEST_UNKNOWN to ErrManifestNotFound
-		if strings.EqualFold(re.Code, "MANIFEST_UNKNOWN") {
-			return nil, ErrModelNotFound
-		}
-
-		re.status = res.StatusCode
-		return nil, &re
-	}
-	return res, nil
-}
-
-// send is a convenience method for making a request with newRequest and
-// passing it to send with r.client().
-func (r *Registry) send(ctx context.Context, method, path string, body io.Reader) (*http.Response, error) {
-	req, err := r.newRequest(ctx, method, path, body)
-	if err != nil {
-		return nil, err
-	}
-	return sendRequest(r.client(), req)
-}
-
-// makeAuthToken creates an Ollama auth token for the given private key.
-//
-// NOTE: This format is OLD, overly complex, and should be replaced. We're
-// inheriting it from the original Ollama client and ollama.com
-// implementations, so we need to support it for now.
-func makeAuthToken(key crypto.PrivateKey) (string, error) {
-	privKey, _ := key.(*ed25519.PrivateKey)
-	if privKey == nil {
-		return "", fmt.Errorf("unsupported private key type: %T", key)
-	}
-
-	url := fmt.Sprintf("https://ollama.com?ts=%d", time.Now().Unix())
-	// Part 1: the checkData (e.g. the URL with a timestamp)
-
-	// Part 2: the public key
-	pubKeyShort, err := func() ([]byte, error) {
-		sshPubKey, err := ssh.NewPublicKey(privKey.Public())
-		if err != nil {
-			return nil, err
-		}
-		pubKeyParts := bytes.Fields(ssh.MarshalAuthorizedKey(sshPubKey))
-		if len(pubKeyParts) < 2 {
-			return nil, fmt.Errorf("malformed public key: %q", pubKeyParts)
-		}
-		pubKeyShort := pubKeyParts[1]
-		return pubKeyShort, nil
-	}()
-	if err != nil {
-		return "", err
-	}
-
-	// Part 3: the signature
-	sig := ed25519.Sign(*privKey, []byte(checkData(url)))
-
-	// Assemble the token: <checkData>:<pubKey>:<signature>
-	var b strings.Builder
-	io.WriteString(&b, base64.StdEncoding.EncodeToString([]byte(url)))
-	b.WriteByte(':')
-	b.Write(pubKeyShort)
-	b.WriteByte(':')
-	io.WriteString(&b, base64.StdEncoding.EncodeToString(sig))
-
-	return b.String(), nil
-}
-
-// The original spec for Ollama tokens was to use the SHA256 of the zero
-// string as part of the signature. I'm not sure why that was, but we still
-// need it to verify the signature.
-var zeroSum = func() string {
-	sha256sum := sha256.Sum256(nil)
-	x := base64.StdEncoding.EncodeToString([]byte(hex.EncodeToString(sha256sum[:])))
-	return x
-}()
-
-// checkData takes a URL and creates the original string format of the
-// data signature that is used by the ollama client to sign requests
-func checkData(url string) string {
-	return fmt.Sprintf("GET,%s,%s", url, zeroSum)
-}
-
-type publicError struct {
-	wrapped error
-	message string
-}
-
-func withPublicMessagef(err error, message string, args ...any) error {
-	return publicError{wrapped: err, message: fmt.Sprintf(message, args...)}
-}
-
-func (e publicError) Error() string { return e.message }
-func (e publicError) Unwrap() error { return e.wrapped }
-
-var supportedSchemes = []string{
-	"http",
-	"https",
-	"https+insecure",
-}
-
-var supportedSchemesMessage = fmt.Sprintf("supported schemes are %v", strings.Join(supportedSchemes, ", "))
-
-// parseNameExtended parses and validates an extended name, returning the scheme, name,
-// and digest.
-//
-// If the scheme is empty, scheme will be "https". If an unsupported scheme is
-// given, [ErrNameInvalid] wrapped with a display friendly message is returned.
-//
-// If the digest is invalid, [ErrNameInvalid] wrapped with a display friendly
-// message is returned.
-//
-// If the name is not, once merged with the mask, fully qualified,
-// [ErrNameInvalid] wrapped with a display friendly message is returned.
-func (r *Registry) parseNameExtended(s string) (scheme string, _ names.Name, _ blob.Digest, _ error) {
-	scheme, name, digest := splitExtended(s)
-	scheme = cmp.Or(scheme, "https")
-	if !slices.Contains(supportedSchemes, scheme) {
-		err := withPublicMessagef(ErrNameInvalid, "unsupported scheme: %q: %s", scheme, supportedSchemesMessage)
-		return "", names.Name{}, blob.Digest{}, err
-	}
-
-	var d blob.Digest
-	if digest != "" {
-		var err error
-		d, err = blob.ParseDigest(digest)
-		if err != nil {
-			err = withPublicMessagef(ErrNameInvalid, "invalid digest: %q", digest)
-			return "", names.Name{}, blob.Digest{}, err
-		}
-		if name == "" {
-			// We have can resolve a manifest from a digest only,
-			// so skip name validation and return the scheme and
-			// digest.
-			return scheme, names.Name{}, d, nil
-		}
-	}
-
-	n, err := r.parseName(name)
-	if err != nil {
-		return "", names.Name{}, blob.Digest{}, err
-	}
-	return scheme, n, d, nil
-}
-
-// splitExtended splits an extended name string into its scheme, name, and digest
-// parts.
-//
-// Examples:
-//
-//	http://ollama.com/bmizerany/smol:latest@digest
-//	https://ollama.com/bmizerany/smol:latest
-//	ollama.com/bmizerany/smol:latest@digest // returns "https" scheme.
-//	model@digest
-//	@digest
-func splitExtended(s string) (scheme, name, digest string) {
-	i := strings.Index(s, "://")
-	if i >= 0 {
-		scheme = s[:i]
-		s = s[i+3:]
-	}
-	i = strings.LastIndex(s, "@")
-	if i >= 0 {
-		digest = s[i+1:]
-		s = s[:i]
-	}
-	return scheme, s, digest
-}
-
-// parseChunk parses a string in the form "start-end" and returns the Chunk.
-func parseChunk[S ~string | ~[]byte](s S) (blob.Chunk, error) {
-	startPart, endPart, found := strings.Cut(string(s), "-")
-	if !found {
-		return blob.Chunk{}, fmt.Errorf("chunks: invalid range %q: missing '-'", s)
-	}
-	start, err := strconv.ParseInt(startPart, 10, 64)
-	if err != nil {
-		return blob.Chunk{}, fmt.Errorf("chunks: invalid start to %q: %v", s, err)
-	}
-	end, err := strconv.ParseInt(endPart, 10, 64)
-	if err != nil {
-		return blob.Chunk{}, fmt.Errorf("chunks: invalid end to %q: %v", s, err)
-	}
-	if start > end {
-		return blob.Chunk{}, fmt.Errorf("chunks: invalid range %q: start > end", s)
-	}
-	return blob.Chunk{Start: start, End: end}, nil
-}

server/internal/client/ollama/registry_synctest_test.go

@@ -1,51 +0,0 @@
-// TODO: go:build goexperiment.synctest
-
-package ollama
-
-import (
-	"context"
-	"errors"
-	"io"
-	"net/http"
-	"strings"
-	"testing"
-	"time"
-)
-
-func TestPullDownloadTimeout(t *testing.T) {
-	rc, ctx := newRegistryClient(t, func(w http.ResponseWriter, r *http.Request) {
-		defer t.Log("upstream", r.Method, r.URL.Path)
-		switch {
-		case strings.HasPrefix(r.URL.Path, "/v2/library/smol/manifests/"):
-			io.WriteString(w, `{
-				"layers": [{"digest": "sha256:1111111111111111111111111111111111111111111111111111111111111111", "size": 3}]
-			}`)
-		case strings.HasPrefix(r.URL.Path, "/v2/library/smol/blobs/sha256:1111111111111111111111111111111111111111111111111111111111111111"):
-			// Get headers out to client and then hang on the response
-			w.WriteHeader(200)
-			w.(http.Flusher).Flush()
-
-			// Hang on the response and unblock when the client
-			// gives up
-			<-r.Context().Done()
-		default:
-			t.Fatalf("unexpected request: %s", r.URL.Path)
-		}
-	})
-	rc.ReadTimeout = 100 * time.Millisecond
-
-	done := make(chan error, 1)
-	go func() {
-		done <- rc.Pull(ctx, "http://example.com/library/smol")
-	}()
-
-	select {
-	case err := <-done:
-		want := context.DeadlineExceeded
-		if !errors.Is(err, want) {
-			t.Errorf("err = %v, want %v", err, want)
-		}
-	case <-time.After(3 * time.Second):
-		t.Error("timeout waiting for Pull to finish")
-	}
-}

server/internal/client/ollama/registry_test.go

@@ -1,953 +0,0 @@
-package ollama
-
-import (
-	"bytes"
-	"cmp"
-	"context"
-	"encoding/json"
-	"errors"
-	"fmt"
-	"io"
-	"io/fs"
-	"net"
-	"net/http"
-	"net/http/httptest"
-	"os"
-	"reflect"
-	"strings"
-	"sync/atomic"
-	"testing"
-
-	"github.com/ollama/ollama/server/internal/cache/blob"
-	"github.com/ollama/ollama/server/internal/testutil"
-)
-
-func ExampleRegistry_cancelOnFirstError() {
-	ctx, cancel := context.WithCancel(context.Background())
-	defer cancel()
-
-	ctx = WithTrace(ctx, &Trace{
-		Update: func(l *Layer, n int64, err error) {
-			if err != nil {
-				// Discontinue pulling layers if there is an
-				// error instead of continuing to pull more
-				// data.
-				cancel()
-			}
-		},
-	})
-
-	var r Registry
-	if err := r.Pull(ctx, "model"); err != nil {
-		// panic for demo purposes
-		panic(err)
-	}
-}
-
-func TestManifestMarshalJSON(t *testing.T) {
-	// All manifests should contain an "empty" config object.
-	var m Manifest
-	data, err := json.Marshal(m)
-	if err != nil {
-		t.Fatal(err)
-	}
-	if !bytes.Contains(data, []byte(`"config":{"digest":"sha256:`)) {
-		t.Error("expected manifest to contain empty config")
-		t.Fatalf("got:\n%s", string(data))
-	}
-}
-
-var errRoundTrip = errors.New("forced roundtrip error")
-
-type recordRoundTripper http.HandlerFunc
-
-func (rr recordRoundTripper) RoundTrip(req *http.Request) (*http.Response, error) {
-	w := httptest.NewRecorder()
-	rr(w, req)
-	if w.Code == 499 {
-		return nil, errRoundTrip
-	}
-	resp := w.Result()
-	// For some reason, Response.Request is not set by httptest.NewRecorder, so we
-	// set it manually.
-	resp.Request = req
-	return w.Result(), nil
-}
-
-// newClient constructs a cache with predefined manifests for testing. The manifests are:
-//
-//	empty:         no data
-//	zero:          no layers
-//	single:        one layer with the contents "exists"
-//	multiple:      two layers with the contents "exists" and "here"
-//	notfound:      a layer that does not exist in the cache
-//	null:          one null layer (e.g. [null])
-//	sizemismatch:  one valid layer, and one with a size mismatch (file size is less than the reported size)
-//	invalid:       a layer with invalid JSON data
-//
-// Tests that want to ensure the client does not communicate with the upstream
-// registry should pass a nil handler, which will cause a panic if
-// communication is attempted.
-//
-// To simulate a network error, pass a handler that returns a 499 status code.
-func newClient(t *testing.T, upstreamRegistry http.HandlerFunc) (*Registry, *blob.DiskCache) {
-	t.Helper()
-
-	c, err := blob.Open(t.TempDir())
-	if err != nil {
-		t.Fatal(err)
-	}
-
-	mklayer := func(data string) *Layer {
-		return &Layer{
-			Digest: importBytes(t, c, data),
-			Size:   int64(len(data)),
-		}
-	}
-
-	r := &Registry{
-		Cache: c,
-		HTTPClient: &http.Client{
-			Transport: recordRoundTripper(upstreamRegistry),
-		},
-	}
-
-	link := func(name string, manifest string) {
-		n, err := r.parseName(name)
-		if err != nil {
-			panic(err)
-		}
-		d, err := c.Import(bytes.NewReader([]byte(manifest)), int64(len(manifest)))
-		if err != nil {
-			panic(err)
-		}
-		if err := c.Link(n.String(), d); err != nil {
-			panic(err)
-		}
-	}
-
-	commit := func(name string, layers ...*Layer) {
-		t.Helper()
-		data, err := json.Marshal(&Manifest{Layers: layers})
-		if err != nil {
-			t.Fatal(err)
-		}
-		link(name, string(data))
-	}
-
-	link("empty", "")
-	commit("zero")
-	commit("single", mklayer("exists"))
-	commit("multiple", mklayer("exists"), mklayer("present"))
-	commit("notfound", &Layer{Digest: blob.DigestFromBytes("notfound"), Size: int64(len("notfound"))})
-	commit("null", nil)
-	commit("sizemismatch", mklayer("exists"), &Layer{Digest: blob.DigestFromBytes("present"), Size: 499})
-	link("invalid", "!!!!!")
-
-	return r, c
-}
-
-func okHandler(w http.ResponseWriter, r *http.Request) {
-	w.WriteHeader(http.StatusOK)
-}
-
-func checkErrCode(t *testing.T, err error, status int, code string) {
-	t.Helper()
-	var e *Error
-	if !errors.As(err, &e) || e.status != status || e.Code != code {
-		t.Errorf("err = %v; want %v %v", err, status, code)
-	}
-}
-
-func importBytes(t *testing.T, c *blob.DiskCache, data string) blob.Digest {
-	d, err := c.Import(strings.NewReader(data), int64(len(data)))
-	if err != nil {
-		t.Fatal(err)
-	}
-	return d
-}
-
-func withTraceUnexpected(ctx context.Context) (context.Context, *Trace) {
-	t := &Trace{Update: func(*Layer, int64, error) { panic("unexpected") }}
-	return WithTrace(ctx, t), t
-}
-
-func TestPushZero(t *testing.T) {
-	rc, _ := newClient(t, okHandler)
-	err := rc.Push(t.Context(), "empty", nil)
-	if !errors.Is(err, ErrManifestInvalid) {
-		t.Errorf("err = %v; want %v", err, ErrManifestInvalid)
-	}
-}
-
-func TestPushSingle(t *testing.T) {
-	rc, _ := newClient(t, okHandler)
-	err := rc.Push(t.Context(), "single", nil)
-	testutil.Check(t, err)
-}
-
-func TestPushMultiple(t *testing.T) {
-	rc, _ := newClient(t, okHandler)
-	err := rc.Push(t.Context(), "multiple", nil)
-	testutil.Check(t, err)
-}
-
-func TestPushNotFound(t *testing.T) {
-	rc, _ := newClient(t, func(w http.ResponseWriter, r *http.Request) {
-		t.Errorf("unexpected request: %v", r)
-	})
-	err := rc.Push(t.Context(), "notfound", nil)
-	if !errors.Is(err, fs.ErrNotExist) {
-		t.Errorf("err = %v; want %v", err, fs.ErrNotExist)
-	}
-}
-
-func TestPushNullLayer(t *testing.T) {
-	rc, _ := newClient(t, nil)
-	err := rc.Push(t.Context(), "null", nil)
-	if err == nil || !strings.Contains(err.Error(), "invalid manifest") {
-		t.Errorf("err = %v; want invalid manifest", err)
-	}
-}
-
-func TestPushSizeMismatch(t *testing.T) {
-	rc, _ := newClient(t, nil)
-	ctx, _ := withTraceUnexpected(t.Context())
-	got := rc.Push(ctx, "sizemismatch", nil)
-	if got == nil || !strings.Contains(got.Error(), "size mismatch") {
-		t.Errorf("err = %v; want size mismatch", got)
-	}
-}
-
-func TestPushInvalid(t *testing.T) {
-	rc, _ := newClient(t, nil)
-	err := rc.Push(t.Context(), "invalid", nil)
-	if err == nil || !strings.Contains(err.Error(), "invalid manifest") {
-		t.Errorf("err = %v; want invalid manifest", err)
-	}
-}
-
-func TestPushExistsAtRemote(t *testing.T) {
-	var pushed bool
-	rc, _ := newClient(t, func(w http.ResponseWriter, r *http.Request) {
-		if strings.Contains(r.URL.Path, "/uploads/") {
-			if !pushed {
-				// First push. Return an uploadURL.
-				pushed = true
-				w.Header().Set("Location", "http://blob.store/blobs/123")
-				return
-			}
-			w.WriteHeader(http.StatusAccepted)
-			return
-		}
-
-		io.Copy(io.Discard, r.Body)
-		w.WriteHeader(http.StatusOK)
-	})
-
-	rc.MaxStreams = 1 // prevent concurrent uploads
-
-	var errs []error
-	ctx := WithTrace(t.Context(), &Trace{
-		Update: func(_ *Layer, n int64, err error) {
-			// uploading one at a time so no need to lock
-			errs = append(errs, err)
-		},
-	})
-
-	check := testutil.Checker(t)
-
-	err := rc.Push(ctx, "single", nil)
-	check(err)
-
-	if !errors.Is(errors.Join(errs...), nil) {
-		t.Errorf("errs = %v; want %v", errs, []error{ErrCached})
-	}
-
-	err = rc.Push(ctx, "single", nil)
-	check(err)
-}
-
-func TestPushRemoteError(t *testing.T) {
-	rc, _ := newClient(t, func(w http.ResponseWriter, r *http.Request) {
-		if strings.Contains(r.URL.Path, "/blobs/") {
-			w.WriteHeader(500)
-			io.WriteString(w, `{"errors":[{"code":"blob_error"}]}`)
-			return
-		}
-	})
-	got := rc.Push(t.Context(), "single", nil)
-	checkErrCode(t, got, 500, "blob_error")
-}
-
-func TestPushLocationError(t *testing.T) {
-	rc, _ := newClient(t, func(w http.ResponseWriter, r *http.Request) {
-		w.Header().Set("Location", ":///x")
-		w.WriteHeader(http.StatusAccepted)
-	})
-	got := rc.Push(t.Context(), "single", nil)
-	wantContains := "invalid upload URL"
-	if got == nil || !strings.Contains(got.Error(), wantContains) {
-		t.Errorf("err = %v; want to contain %v", got, wantContains)
-	}
-}
-
-func TestPushUploadRoundtripError(t *testing.T) {
-	rc, _ := newClient(t, func(w http.ResponseWriter, r *http.Request) {
-		if r.Host == "blob.store" {
-			w.WriteHeader(499) // force RoundTrip error on upload
-			return
-		}
-		w.Header().Set("Location", "http://blob.store/blobs/123")
-	})
-	got := rc.Push(t.Context(), "single", nil)
-	if !errors.Is(got, errRoundTrip) {
-		t.Errorf("got = %v; want %v", got, errRoundTrip)
-	}
-}
-
-func TestPushUploadFileOpenError(t *testing.T) {
-	rc, c := newClient(t, okHandler)
-	ctx := WithTrace(t.Context(), &Trace{
-		Update: func(l *Layer, _ int64, err error) {
-			// Remove the file just before it is opened for upload,
-			// but after the initial Stat that happens before the
-			// upload starts
-			os.Remove(c.GetFile(l.Digest))
-		},
-	})
-	got := rc.Push(ctx, "single", nil)
-	if !errors.Is(got, fs.ErrNotExist) {
-		t.Errorf("got = %v; want fs.ErrNotExist", got)
-	}
-}
-
-func TestPushCommitRoundtripError(t *testing.T) {
-	rc, _ := newClient(t, func(w http.ResponseWriter, r *http.Request) {
-		if strings.Contains(r.URL.Path, "/blobs/") {
-			panic("unexpected")
-		}
-		w.WriteHeader(499) // force RoundTrip error
-	})
-	err := rc.Push(t.Context(), "zero", nil)
-	if !errors.Is(err, errRoundTrip) {
-		t.Errorf("err = %v; want %v", err, errRoundTrip)
-	}
-}
-
-func TestRegistryPullInvalidName(t *testing.T) {
-	rc, _ := newRegistryClient(t, nil)
-	err := rc.Pull(t.Context(), "://")
-	if !errors.Is(err, ErrNameInvalid) {
-		t.Errorf("err = %v; want %v", err, ErrNameInvalid)
-	}
-}
-
-func TestRegistryPullInvalidManifest(t *testing.T) {
-	cases := []string{
-		"",
-		"null",
-		"!!!",
-		`{"layers":[]}`,
-	}
-
-	for _, resp := range cases {
-		rc, _ := newRegistryClient(t, func(w http.ResponseWriter, r *http.Request) {
-			io.WriteString(w, resp)
-		})
-		err := rc.Pull(t.Context(), "http://example.com/a/b")
-		if !errors.Is(err, ErrManifestInvalid) {
-			t.Errorf("err = %v; want invalid manifest", err)
-		}
-	}
-}
-
-func TestRegistryResolveByDigest(t *testing.T) {
-	check := testutil.Checker(t)
-
-	exists := blob.DigestFromBytes("exists")
-	rc, _ := newClient(t, func(w http.ResponseWriter, r *http.Request) {
-		if r.URL.Path != "/v2/alice/palace/blobs/"+exists.String() {
-			w.WriteHeader(499) // should not hit manifest endpoint
-		}
-		fmt.Fprintf(w, `{"layers":[{"digest":%q,"size":5}]}`, exists)
-	})
-
-	_, err := rc.Resolve(t.Context(), "alice/palace@"+exists.String())
-	check(err)
-}
-
-func TestInsecureSkipVerify(t *testing.T) {
-	exists := blob.DigestFromBytes("exists")
-
-	s := httptest.NewTLSServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
-		fmt.Fprintf(w, `{"layers":[{"digest":%q,"size":5}]}`, exists)
-	}))
-	defer s.Close()
-
-	const name = "library/insecure"
-
-	var rc Registry
-	url := fmt.Sprintf("https://%s/%s", s.Listener.Addr(), name)
-	_, err := rc.Resolve(t.Context(), url)
-	if err == nil || !strings.Contains(err.Error(), "failed to verify") {
-		t.Errorf("err = %v; want cert verification failure", err)
-	}
-
-	url = fmt.Sprintf("https+insecure://%s/%s", s.Listener.Addr(), name)
-	_, err = rc.Resolve(t.Context(), url)
-	testutil.Check(t, err)
-}
-
-func TestErrorUnmarshal(t *testing.T) {
-	cases := []struct {
-		name    string
-		data    string
-		want    *Error
-		wantErr bool
-	}{
-		{
-			name:    "errors empty",
-			data:    `{"errors":[]}`,
-			wantErr: true,
-		},
-		{
-			name:    "errors empty",
-			data:    `{"errors":[]}`,
-			wantErr: true,
-		},
-		{
-			name: "errors single",
-			data: `{"errors":[{"code":"blob_unknown"}]}`,
-			want: &Error{Code: "blob_unknown", Message: ""},
-		},
-		{
-			name: "errors multiple",
-			data: `{"errors":[{"code":"blob_unknown"},{"code":"blob_error"}]}`,
-			want: &Error{Code: "blob_unknown", Message: ""},
-		},
-		{
-			name:    "error empty",
-			data:    `{"error":""}`,
-			wantErr: true,
-		},
-		{
-			name:    "error very empty",
-			data:    `{}`,
-			wantErr: true,
-		},
-		{
-			name: "error message",
-			data: `{"error":"message", "code":"code"}`,
-			want: &Error{Code: "code", Message: "message"},
-		},
-		{
-			name:    "invalid value",
-			data:    `{"error": 1}`,
-			wantErr: true,
-		},
-	}
-	for _, tt := range cases {
-		t.Run(tt.name, func(t *testing.T) {
-			var got Error
-			err := json.Unmarshal([]byte(tt.data), &got)
-			if err != nil {
-				if tt.wantErr {
-					return
-				}
-				t.Errorf("Unmarshal() error = %v", err)
-				// fallthrough and check got
-			}
-			if tt.want == nil {
-				tt.want = &Error{}
-			}
-			if !reflect.DeepEqual(got, *tt.want) {
-				t.Errorf("got = %v; want %v", got, *tt.want)
-			}
-		})
-	}
-}
-
-// TestParseNameExtendedErrors tests that parseName returns errors messages with enough
-// detail for users to debug naming issues they may encounter. Previous to this
-// test, the error messages were not very helpful and each problem was reported
-// as the same message.
-//
-// It is only for testing error messages, not that all invalids and valids are
-// covered. Those are in other tests for names.Name and blob.Digest.
-func TestParseNameExtendedErrors(t *testing.T) {
-	cases := []struct {
-		name string
-		err  error
-		want string
-	}{}
-
-	var r Registry
-	for _, tt := range cases {
-		_, _, _, err := r.parseNameExtended(tt.name)
-		if !errors.Is(err, tt.err) {
-			t.Errorf("[%s]: err = %v; want %v", tt.name, err, tt.err)
-		}
-		if err != nil && !strings.Contains(err.Error(), tt.want) {
-			t.Errorf("[%s]: err =\n\t%v\nwant\n\t%v", tt.name, err, tt.want)
-		}
-	}
-}
-
-func TestParseNameExtended(t *testing.T) {
-	cases := []struct {
-		in     string
-		scheme string
-		name   string
-		digest string
-		err    string
-	}{
-		{in: "http://m", scheme: "http", name: "m"},
-		{in: "https+insecure://m", scheme: "https+insecure", name: "m"},
-		{in: "http+insecure://m", err: "unsupported scheme"},
-
-		{in: "http://m@sha256:1111111111111111111111111111111111111111111111111111111111111111", scheme: "http", name: "m", digest: "sha256:1111111111111111111111111111111111111111111111111111111111111111"},
-
-		{in: "", err: "invalid or missing name"},
-		{in: "m", scheme: "https", name: "m"},
-		{in: "://", err: "invalid or missing name"},
-		{in: "@sha256:deadbeef", err: "invalid digest"},
-		{in: "@sha256:deadbeef@sha256:deadbeef", err: "invalid digest"},
-	}
-	for _, tt := range cases {
-		t.Run(tt.in, func(t *testing.T) {
-			var r Registry
-			scheme, n, digest, err := r.parseNameExtended(tt.in)
-			if err != nil {
-				if tt.err == "" {
-					t.Errorf("err = %v; want nil", err)
-				} else if !strings.Contains(err.Error(), tt.err) {
-					t.Errorf("err = %v; want %q", err, tt.err)
-				}
-			} else if tt.err != "" {
-				t.Errorf("err = nil; want %q", tt.err)
-			}
-			if err == nil && !n.IsFullyQualified() {
-				t.Errorf("name = %q; want fully qualified", n)
-			}
-
-			if scheme != tt.scheme {
-				t.Errorf("scheme = %q; want %q", scheme, tt.scheme)
-			}
-
-			// smoke-test name is superset of tt.name
-			if !strings.Contains(n.String(), tt.name) {
-				t.Errorf("name = %q; want %q", n, tt.name)
-			}
-
-			tt.digest = cmp.Or(tt.digest, (&blob.Digest{}).String())
-			if digest.String() != tt.digest {
-				t.Errorf("digest = %q; want %q", digest, tt.digest)
-			}
-		})
-	}
-}
-
-func TestUnlink(t *testing.T) {
-	t.Run("found by name", func(t *testing.T) {
-		check := testutil.Checker(t)
-
-		rc, _ := newRegistryClient(t, nil)
-		// make a blob and link it
-		d := blob.DigestFromBytes("{}")
-		err := blob.PutBytes(rc.Cache, d, "{}")
-		check(err)
-		err = rc.Cache.Link("registry.ollama.ai/library/single:latest", d)
-		check(err)
-
-		// confirm linked
-		_, err = rc.ResolveLocal("single")
-		check(err)
-
-		// unlink
-		_, err = rc.Unlink("single")
-		check(err)
-
-		// confirm unlinked
-		_, err = rc.ResolveLocal("single")
-		if !errors.Is(err, fs.ErrNotExist) {
-			t.Errorf("err = %v; want fs.ErrNotExist", err)
-		}
-	})
-	t.Run("not found by name", func(t *testing.T) {
-		rc, _ := newRegistryClient(t, nil)
-		ok, err := rc.Unlink("manifestNotFound")
-		if err != nil {
-			t.Fatal(err)
-		}
-		if ok {
-			t.Error("expected not found")
-		}
-	})
-}
-
-// Many tests from here out, in this file are based on a single blob, "abc",
-// with the checksum of its sha256 hash. The checksum is:
-//
-//	"abc" -> sha256:ba7816bf8f01cfea414140de5dae2223b00361a396177a9cb410ff61f20015ad
-//
-// Using the literal value instead of a constant with fmt.Xprintf calls proved
-// to be the most readable and maintainable approach. The sum is consistently
-// used in the tests and unique so searches do not yield false positives.
-
-func checkRequest(t *testing.T, req *http.Request, method, path string) {
-	t.Helper()
-	if got := req.URL.Path; got != path {
-		t.Errorf("URL = %q, want %q", got, path)
-	}
-	if req.Method != method {
-		t.Errorf("Method = %q, want %q", req.Method, method)
-	}
-}
-
-func newRegistryClient(t *testing.T, upstream http.HandlerFunc) (*Registry, context.Context) {
-	s := httptest.NewServer(upstream)
-	t.Cleanup(s.Close)
-	cache, err := blob.Open(t.TempDir())
-	if err != nil {
-		t.Fatal(err)
-	}
-
-	ctx := WithTrace(t.Context(), &Trace{
-		Update: func(l *Layer, n int64, err error) {
-			t.Log("trace:", l.Digest.Short(), n, err)
-		},
-	})
-
-	rc := &Registry{
-		Cache: cache,
-		HTTPClient: &http.Client{Transport: &http.Transport{
-			Dial: func(network, addr string) (net.Conn, error) {
-				return net.Dial(network, s.Listener.Addr().String())
-			},
-		}},
-	}
-	return rc, ctx
-}
-
-func TestPullChunked(t *testing.T) {
-	var steps atomic.Int64
-	c, ctx := newRegistryClient(t, func(w http.ResponseWriter, r *http.Request) {
-		switch steps.Add(1) {
-		case 1:
-			checkRequest(t, r, "GET", "/v2/library/abc/manifests/latest")
-			io.WriteString(w, `{"layers":[{"size":3,"digest":"sha256:ba7816bf8f01cfea414140de5dae2223b00361a396177a9cb410ff61f20015ad"}]}`)
-		case 2:
-			checkRequest(t, r, "GET", "/v2/library/abc/chunksums/sha256:ba7816bf8f01cfea414140de5dae2223b00361a396177a9cb410ff61f20015ad")
-			w.Header().Set("Content-Location", "http://blob.store/v2/library/abc/blobs/sha256:ba7816bf8f01cfea414140de5dae2223b00361a396177a9cb410ff61f20015ad")
-			fmt.Fprintf(w, "%s 0-1\n", blob.DigestFromBytes("ab"))
-			fmt.Fprintf(w, "%s 2-2\n", blob.DigestFromBytes("c"))
-		case 3, 4:
-			checkRequest(t, r, "GET", "/v2/library/abc/blobs/sha256:ba7816bf8f01cfea414140de5dae2223b00361a396177a9cb410ff61f20015ad")
-			switch rng := r.Header.Get("Range"); rng {
-			case "bytes=0-1":
-				io.WriteString(w, "ab")
-			case "bytes=2-2":
-				t.Logf("writing c")
-				io.WriteString(w, "c")
-			default:
-				t.Errorf("unexpected range %q", rng)
-			}
-		default:
-			t.Errorf("unexpected steps %d: %v", steps.Load(), r)
-			http.Error(w, "unexpected steps", http.StatusInternalServerError)
-		}
-	})
-
-	c.ChunkingThreshold = 1 // force chunking
-
-	err := c.Pull(ctx, "http://o.com/library/abc")
-	testutil.Check(t, err)
-
-	_, err = c.Cache.Resolve("o.com/library/abc:latest")
-	testutil.Check(t, err)
-
-	if g := steps.Load(); g != 4 {
-		t.Fatalf("got %d steps, want 4", g)
-	}
-}
-
-func TestPullCached(t *testing.T) {
-	c, ctx := newRegistryClient(t, func(w http.ResponseWriter, r *http.Request) {
-		checkRequest(t, r, "GET", "/v2/library/abc/manifests/latest")
-		io.WriteString(w, `{"layers":[{"size":3,"digest":"sha256:ba7816bf8f01cfea414140de5dae2223b00361a396177a9cb410ff61f20015ad"}]}`)
-	})
-
-	check := testutil.Checker(t)
-
-	// Premeptively cache the blob
-	d, err := blob.ParseDigest("sha256:ba7816bf8f01cfea414140de5dae2223b00361a396177a9cb410ff61f20015ad")
-	check(err)
-	err = blob.PutBytes(c.Cache, d, []byte("abc"))
-	check(err)
-
-	// Pull only the manifest, which should be enough to resolve the cached blob
-	err = c.Pull(ctx, "http://o.com/library/abc")
-	check(err)
-}
-
-func TestPullManifestError(t *testing.T) {
-	c, ctx := newRegistryClient(t, func(w http.ResponseWriter, r *http.Request) {
-		checkRequest(t, r, "GET", "/v2/library/abc/manifests/latest")
-		w.WriteHeader(http.StatusNotFound)
-		io.WriteString(w, `{"errors":[{"code":"MANIFEST_UNKNOWN"}]}`)
-	})
-
-	err := c.Pull(ctx, "http://o.com/library/abc")
-	if err == nil {
-		t.Fatalf("expected error")
-	}
-	var got *Error
-	if !errors.Is(err, ErrModelNotFound) {
-		t.Fatalf("err = %v, want %v", got, ErrModelNotFound)
-	}
-}
-
-func TestPullLayerError(t *testing.T) {
-	c, ctx := newRegistryClient(t, func(w http.ResponseWriter, r *http.Request) {
-		checkRequest(t, r, "GET", "/v2/library/abc/manifests/latest")
-		io.WriteString(w, `!`)
-	})
-
-	err := c.Pull(ctx, "http://o.com/library/abc")
-	if err == nil {
-		t.Fatalf("expected error")
-	}
-	var want *json.SyntaxError
-	if !errors.As(err, &want) {
-		t.Fatalf("err = %T, want %T", err, want)
-	}
-}
-
-func TestPullLayerChecksumError(t *testing.T) {
-	var step atomic.Int64
-	c, _ := newRegistryClient(t, func(w http.ResponseWriter, r *http.Request) {
-		switch step.Add(1) {
-		case 1:
-			checkRequest(t, r, "GET", "/v2/library/abc/manifests/latest")
-			io.WriteString(w, `{"layers":[{"size":3,"digest":"sha256:ba7816bf8f01cfea414140de5dae2223b00361a396177a9cb410ff61f20015ad"}]}`)
-		case 2:
-			checkRequest(t, r, "GET", "/v2/library/abc/chunksums/sha256:ba7816bf8f01cfea414140de5dae2223b00361a396177a9cb410ff61f20015ad")
-			w.Header().Set("Content-Location", "http://blob.store/v2/library/abc/blobs/sha256:ba7816bf8f01cfea414140de5dae2223b00361a396177a9cb410ff61f20015ad")
-			fmt.Fprintf(w, "%s 0-1\n", blob.DigestFromBytes("ab"))
-			fmt.Fprintf(w, "%s 2-2\n", blob.DigestFromBytes("c"))
-		case 3:
-			w.WriteHeader(http.StatusNotFound)
-			io.WriteString(w, `{"errors":[{"code":"BLOB_UNKNOWN"}]}`)
-		case 4:
-			io.WriteString(w, "c")
-		default:
-			t.Errorf("unexpected steps %d: %v", step.Load(), r)
-			http.Error(w, "unexpected steps", http.StatusInternalServerError)
-		}
-	})
-
-	c.MaxStreams = 1
-	c.ChunkingThreshold = 1 // force chunking
-
-	var written atomic.Int64
-	ctx := WithTrace(t.Context(), &Trace{
-		Update: func(l *Layer, n int64, err error) {
-			t.Log("trace:", l.Digest.Short(), n, err)
-			written.Add(n)
-		},
-	})
-
-	err := c.Pull(ctx, "http://o.com/library/abc")
-	var got *Error
-	if !errors.As(err, &got) || got.Code != "BLOB_UNKNOWN" {
-		t.Fatalf("err = %v, want %v", err, got)
-	}
-
-	if g := written.Load(); g != 1 {
-		t.Fatalf("wrote %d bytes, want 1", g)
-	}
-}
-
-func TestPullChunksumStreamError(t *testing.T) {
-	var step atomic.Int64
-	c, ctx := newRegistryClient(t, func(w http.ResponseWriter, r *http.Request) {
-		switch step.Add(1) {
-		case 1:
-			checkRequest(t, r, "GET", "/v2/library/abc/manifests/latest")
-			io.WriteString(w, `{"layers":[{"size":3,"digest":"sha256:ba7816bf8f01cfea414140de5dae2223b00361a396177a9cb410ff61f20015ad"}]}`)
-		case 2:
-			w.Header().Set("Content-Location", "http://blob.store/v2/library/abc/blobs/sha256:ba7816bf8f01cfea414140de5dae2223b00361a396177a9cb410ff61f20015ad")
-
-			// Write one valid chunksum and one invalid chunksum
-			fmt.Fprintf(w, "%s 0-1\n", blob.DigestFromBytes("ab")) // valid
-			fmt.Fprint(w, "sha256:!")                              // invalid
-		case 3:
-			io.WriteString(w, "ab")
-		default:
-			t.Errorf("unexpected steps %d: %v", step.Load(), r)
-			http.Error(w, "unexpected steps", http.StatusInternalServerError)
-		}
-	})
-
-	c.ChunkingThreshold = 1 // force chunking
-
-	got := c.Pull(ctx, "http://o.com/library/abc")
-	if !errors.Is(got, ErrIncomplete) {
-		t.Fatalf("err = %v, want %v", got, ErrIncomplete)
-	}
-}
-
-type flushAfterWriter struct {
-	w io.Writer
-}
-
-func (f *flushAfterWriter) Write(p []byte) (n int, err error) {
-	n, err = f.w.Write(p)
-	f.w.(http.Flusher).Flush() // panic if not a flusher
-	return
-}
-
-func TestPullChunksumStreaming(t *testing.T) {
-	csr, csw := io.Pipe()
-	defer csw.Close()
-
-	var step atomic.Int64
-	c, _ := newRegistryClient(t, func(w http.ResponseWriter, r *http.Request) {
-		switch step.Add(1) {
-		case 1:
-			checkRequest(t, r, "GET", "/v2/library/abc/manifests/latest")
-			io.WriteString(w, `{"layers":[{"size":3,"digest":"sha256:ba7816bf8f01cfea414140de5dae2223b00361a396177a9cb410ff61f20015ad"}]}`)
-		case 2:
-			w.Header().Set("Content-Location", "http://blob.store/v2/library/abc/blobs/sha256:ba7816bf8f01cfea414140de5dae2223b00361a396177a9cb410ff61f20015ad")
-			fw := &flushAfterWriter{w} // ensure client gets data as it arrives by aggressively flushing
-			_, err := io.Copy(fw, csr)
-			if err != nil {
-				t.Errorf("copy: %v", err)
-			}
-		case 3:
-			io.WriteString(w, "ab")
-		case 4:
-			io.WriteString(w, "c")
-		default:
-			t.Errorf("unexpected steps %d: %v", step.Load(), r)
-			http.Error(w, "unexpected steps", http.StatusInternalServerError)
-		}
-	})
-
-	c.ChunkingThreshold = 1 // force chunking
-
-	update := make(chan int64, 1)
-	ctx := WithTrace(t.Context(), &Trace{
-		Update: func(l *Layer, n int64, err error) {
-			t.Log("trace:", l.Digest.Short(), n, err)
-			if n > 0 {
-				update <- n
-			}
-		},
-	})
-
-	errc := make(chan error, 1)
-	go func() {
-		errc <- c.Pull(ctx, "http://o.com/library/abc")
-	}()
-
-	// Send first chunksum and ensure it kicks off work immediately
-	fmt.Fprintf(csw, "%s 0-1\n", blob.DigestFromBytes("ab"))
-	if g := <-update; g != 2 {
-		t.Fatalf("got %d, want 2", g)
-	}
-
-	// now send the second chunksum and ensure it kicks off work immediately
-	fmt.Fprintf(csw, "%s 2-2\n", blob.DigestFromBytes("c"))
-	if g := <-update; g != 3 {
-		t.Fatalf("got %d, want 3", g)
-	}
-	csw.Close()
-	testutil.Check(t, <-errc)
-}
-
-func TestPullChunksumsCached(t *testing.T) {
-	var step atomic.Int64
-	c, _ := newRegistryClient(t, func(w http.ResponseWriter, r *http.Request) {
-		switch step.Add(1) {
-		case 1:
-			checkRequest(t, r, "GET", "/v2/library/abc/manifests/latest")
-			io.WriteString(w, `{"layers":[{"size":3,"digest":"sha256:ba7816bf8f01cfea414140de5dae2223b00361a396177a9cb410ff61f20015ad"}]}`)
-		case 2:
-			w.Header().Set("Content-Location", "http://blob.store/v2/library/abc/blobs/sha256:ba7816bf8f01cfea414140de5dae2223b00361a396177a9cb410ff61f20015ad")
-			fmt.Fprintf(w, "%s 0-1\n", blob.DigestFromBytes("ab"))
-			fmt.Fprintf(w, "%s 2-2\n", blob.DigestFromBytes("c"))
-		case 3, 4:
-			switch rng := r.Header.Get("Range"); rng {
-			case "bytes=0-1":
-				io.WriteString(w, "ab")
-			case "bytes=2-2":
-				io.WriteString(w, "c")
-			default:
-				t.Errorf("unexpected range %q", rng)
-			}
-		default:
-			t.Errorf("unexpected steps %d: %v", step.Load(), r)
-			http.Error(w, "unexpected steps", http.StatusInternalServerError)
-		}
-	})
-
-	c.MaxStreams = 1        // force serial processing of chunksums
-	c.ChunkingThreshold = 1 // force chunking
-
-	ctx, cancel := context.WithCancel(t.Context())
-	defer cancel()
-
-	// Cancel the pull after the first chunksum is processed, but before
-	// the second chunksum is processed (which is waiting because
-	// MaxStreams=1). This should cause the second chunksum to error out
-	// leaving the blob incomplete.
-	ctx = WithTrace(ctx, &Trace{
-		Update: func(l *Layer, n int64, err error) {
-			if n > 0 {
-				cancel()
-			}
-		},
-	})
-	err := c.Pull(ctx, "http://o.com/library/abc")
-	if !errors.Is(err, context.Canceled) {
-		t.Fatalf("err = %v, want %v", err, context.Canceled)
-	}
-
-	_, err = c.Cache.Resolve("o.com/library/abc:latest")
-	if !errors.Is(err, fs.ErrNotExist) {
-		t.Fatalf("err = %v, want nil", err)
-	}
-
-	// Reset state and pull again to ensure the blob chunks that should
-	// have been cached are, and the remaining chunk was downloaded, making
-	// the blob complete.
-	step.Store(0)
-	var written atomic.Int64
-	var cached atomic.Int64
-	ctx = WithTrace(t.Context(), &Trace{
-		Update: func(l *Layer, n int64, err error) {
-			t.Log("trace:", l.Digest.Short(), n, err)
-			if errors.Is(err, ErrCached) {
-				cached.Add(n)
-			}
-			written.Add(n)
-		},
-	})
-
-	check := testutil.Checker(t)
-
-	err = c.Pull(ctx, "http://o.com/library/abc")
-	check(err)
-
-	_, err = c.Cache.Resolve("o.com/library/abc:latest")
-	check(err)
-
-	if g := written.Load(); g != 5 {
-		t.Fatalf("wrote %d bytes, want 3", g)
-	}
-	if g := cached.Load(); g != 2 { // "ab" should have been cached
-		t.Fatalf("cached %d bytes, want 5", g)
-	}
-}

server/internal/client/ollama/trace.go

@@ -1,72 +0,0 @@
-package ollama
-
-import (
-	"context"
-)
-
-// Trace is a set of functions that are called to report progress during blob
-// downloads and uploads.
-//
-// Use [WithTrace] to attach a Trace to a context for use with [Registry.Push]
-// and [Registry.Pull].
-type Trace struct {
-	// Update is called during [Registry.Push] and [Registry.Pull] to
-	// report the progress of blob uploads and downloads.
-	//
-	// The n argument is the number of bytes transferred so far, and err is
-	// any error that has occurred. If n == 0, and err is nil, the download
-	// or upload has just started. If err is [ErrCached], the download or
-	// upload has been skipped because the blob is already present in the
-	// local cache or remote registry, respectively. Otherwise, if err is
-	// non-nil, the download or upload has failed. When l.Size == n, and
-	// err is nil, the download or upload has completed.
-	//
-	// A function assigned must be safe for concurrent use. The function is
-	// called synchronously and so should not block or take long to run.
-	Update func(_ *Layer, n int64, _ error)
-}
-
-func (t *Trace) update(l *Layer, n int64, err error) {
-	if t.Update != nil {
-		t.Update(l, n, err)
-	}
-}
-
-type traceKey struct{}
-
-// WithTrace adds a trace to the context for transfer progress reporting.
-func WithTrace(ctx context.Context, t *Trace) context.Context {
-	old := traceFromContext(ctx)
-	if old == t {
-		// No change, return the original context. This also prevents
-		// infinite recursion below, if the caller passes the same
-		// Trace.
-		return ctx
-	}
-
-	// Create a new Trace that wraps the old one, if any. If we used the
-	// same pointer t, we end up with a recursive structure.
-	composed := &Trace{
-		Update: func(l *Layer, n int64, err error) {
-			if old != nil {
-				old.update(l, n, err)
-			}
-			t.update(l, n, err)
-		},
-	}
-	return context.WithValue(ctx, traceKey{}, composed)
-}
-
-var emptyTrace = &Trace{}
-
-// traceFromContext returns the Trace associated with ctx, or an empty Trace if
-// none is found.
-//
-// It never returns nil.
-func traceFromContext(ctx context.Context) *Trace {
-	t, _ := ctx.Value(traceKey{}).(*Trace)
-	if t == nil {
-		return emptyTrace
-	}
-	return t
-}

server/internal/internal/backoff/backoff.go

@@ -1,45 +0,0 @@
-package backoff
-
-import (
-	"context"
-	"iter"
-	"math/rand/v2"
-	"time"
-)
-
-func Loop(ctx context.Context, maxBackoff time.Duration) iter.Seq2[int, error] {
-	var n int
-	return func(yield func(int, error) bool) {
-		var t *time.Timer
-		for {
-			if ctx.Err() != nil {
-				yield(n, ctx.Err())
-				return
-			}
-
-			if !yield(n, nil) {
-				return
-			}
-
-			n++
-
-			// n^2 backoff timer is a little smoother than the
-			// common choice of 2^n.
-			d := min(time.Duration(n*n)*10*time.Millisecond, maxBackoff)
-			// Randomize the delay between 0.5-1.5 x msec, in order
-			// to prevent accidental "thundering herd" problems.
-			d = time.Duration(float64(d) * (rand.Float64() + 0.5))
-
-			if t == nil {
-				t = time.NewTimer(d)
-			} else {
-				t.Reset(d)
-			}
-			select {
-			case <-ctx.Done():
-				t.Stop()
-			case <-t.C:
-			}
-		}
-	}
-}

server/internal/internal/backoff/backoff_synctest_test.go

@@ -1,38 +0,0 @@
-package backoff
-
-import (
-	"context"
-	"errors"
-	"testing"
-	"testing/synctest"
-	"time"
-)
-
-func TestLoop(t *testing.T) {
-	synctest.Test(t, func(t *testing.T) {
-		last := -1
-
-		ctx, cancel := context.WithCancel(t.Context())
-		defer cancel()
-
-		for n, err := range Loop(ctx, 100*time.Millisecond) {
-			if !errors.Is(err, ctx.Err()) {
-				t.Errorf("err = %v, want nil", err)
-			}
-			if err != nil {
-				break
-			}
-			if n != last+1 {
-				t.Errorf("n = %d, want %d", n, last+1)
-			}
-			last = n
-			if n > 5 {
-				cancel()
-			}
-		}
-
-		if last != 6 {
-			t.Errorf("last = %d, want 6", last)
-		}
-	})
-}

server/internal/internal/backoff/backoff_test.go

@@ -1,24 +0,0 @@
-package backoff
-
-import (
-	"testing"
-)
-
-func TestLoopAllocs(t *testing.T) {
-	for i := range 3 {
-		got := testing.AllocsPerRun(1000, func() {
-			for tick := range Loop(t.Context(), 1) {
-				if tick >= i {
-					break
-				}
-			}
-		})
-		want := float64(0)
-		if i > 0 {
-			want = 3 // due to time.NewTimer
-		}
-		if got > want {
-			t.Errorf("[%d ticks]: allocs = %v, want 0", i, want)
-		}
-	}
-}

server/internal/internal/names/name.go

@@ -1,228 +0,0 @@
-package names
-
-import (
-	"cmp"
-	"fmt"
-	"strings"
-
-	"github.com/ollama/ollama/server/internal/internal/stringsx"
-)
-
-const MaxNameLength = 350 + 1 + 80 + 1 + 80 + 1 + 80 // <host>/<namespace>/<model>:<tag>
-
-type Name struct {
-	// Make incomparable to enfoce use of Compare / Equal for
-	// case-insensitive comparisons.
-	_ [0]func()
-
-	h string
-	n string
-	m string
-	t string
-}
-
-// Parse parses and assembles a Name from a name string. The
-// format of a valid name string is:
-//
-//	  s:
-//		  { host } "/" { namespace } "/" { model } ":" { tag }
-//		  { host } "/" { namespace } "/" { model }
-//		  { namespace } "/" { model } ":" { tag }
-//		  { namespace } "/" { model }
-//		  { model } ":" { tag }
-//		  { model }
-//	  host:
-//	      pattern: { alphanum | "_" } { alphanum | "_" | "-" | "." | ":" }*
-//	      length:  [1, 350]
-//	  namespace:
-//	      pattern: { alphanum | "_" } { alphanum | "-" | "_" }*
-//	      length:  [1, 80]
-//	  model:
-//	      pattern: { alphanum | "_" } { alphanum | "-" | "_" | "." }*
-//	      length:  [1, 80]
-//	  tag:
-//	      pattern: { alphanum | "_" } { alphanum | "-" | "_" | "." }*
-//	      length:  [1, 80]
-//
-// The name returned is not guaranteed to be valid. If it is not valid, the
-// field values are left in an undefined state. Use [Name.IsValid] to check
-// if the name is valid.
-func Parse(s string) Name {
-	if len(s) > MaxNameLength {
-		return Name{}
-	}
-
-	var n Name
-	var tail string
-	var c byte
-	for {
-		s, tail, c = cutLastAny(s, "/:")
-		switch c {
-		case ':':
-			n.t = tail
-			continue // look for model
-		case '/':
-			n.h, n.n, _ = cutLastAny(s, "/")
-			n.m = tail
-			return n
-		case 0:
-			n.m = tail
-			return n
-		}
-	}
-}
-
-// Split splits an extended name string into its scheme, name, and digest
-// parts.
-//
-// Examples:
-//
-//	http://ollama.com/bmizerany/smol:latest@digest
-//	https://ollama.com/bmizerany/smol:latest
-//	ollama.com/bmizerany/smol:latest@digest // returns "https" scheme.
-//	model@digest
-//	@digest
-func Split(s string) (scheme, name, digest string) {
-	i := strings.Index(s, "://")
-	if i >= 0 {
-		scheme = s[:i]
-		s = s[i+3:]
-	}
-	i = strings.LastIndex(s, "@")
-	if i >= 0 {
-		digest = s[i+1:]
-		s = s[:i]
-	}
-	return scheme, s, digest
-}
-
-// Merge merges two names into a single name. Non-empty host, namespace, and
-// tag parts of a take precedence over fields in b. The model field is left as
-// is.
-//
-// The returned name is not guaranteed to be valid. Use [Name.IsValid] to check
-// if the name is valid.
-func Merge(a, b Name) Name {
-	a.h = cmp.Or(a.h, b.h)
-	a.n = cmp.Or(a.n, b.n)
-	a.t = cmp.Or(a.t, b.t)
-	return a
-}
-
-// IsValid returns true if the name is valid.
-func (n Name) IsValid() bool {
-	if n.h != "" && !isValidPart(partHost, n.h) {
-		return false
-	}
-	if n.n != "" && !isValidPart(partNamespace, n.n) {
-		return false
-	}
-	if n.t != "" && !isValidPart(partTag, n.t) {
-		return false
-	}
-
-	// at bare minimum, model must be present and valid
-	return n.m != "" && isValidPart(partModel, n.m)
-}
-
-func (n Name) IsFullyQualified() bool {
-	return n.IsValid() && n.h != "" && n.n != "" && n.m != "" && n.t != ""
-}
-
-const (
-	partHost = iota
-	partNamespace
-	partModel
-	partTag
-)
-
-func isValidPart(kind int, s string) bool {
-	maxlen := 80
-	if kind == partHost {
-		maxlen = 350
-	}
-	if len(s) > maxlen {
-		return false
-	}
-
-	for i := range s {
-		if i == 0 {
-			if !isAlphanumericOrUnderscore(s[i]) {
-				return false
-			}
-			continue
-		}
-		switch s[i] {
-		case '_', '-':
-		case '.':
-			if kind == partNamespace {
-				return false
-			}
-		case ':':
-			if kind != partHost {
-				return false
-			}
-		default:
-			if !isAlphanumericOrUnderscore(s[i]) {
-				return false
-			}
-		}
-	}
-	return true
-}
-
-func isAlphanumericOrUnderscore(c byte) bool {
-	return c >= 'A' && c <= 'Z' || c >= 'a' && c <= 'z' || c >= '0' && c <= '9' || c == '_'
-}
-
-func (n Name) Host() string      { return n.h }
-func (n Name) Namespace() string { return n.n }
-func (n Name) Model() string     { return n.m }
-func (n Name) Tag() string       { return n.t }
-
-// Compare compares n and o case-insensitively. It returns 0 if n and o are
-// equal, -1 if n sorts before o, and 1 if n sorts after o.
-func (n Name) Compare(o Name) int {
-	return cmp.Or(
-		stringsx.CompareFold(n.h, o.h),
-		stringsx.CompareFold(n.n, o.n),
-		stringsx.CompareFold(n.m, o.m),
-		stringsx.CompareFold(n.t, o.t),
-	)
-}
-
-// String returns the fully qualified name in the format
-// <namespace>/<model>:<tag>.
-func (n Name) String() string {
-	var b strings.Builder
-	if n.h != "" {
-		b.WriteString(n.h)
-		b.WriteByte('/')
-	}
-	if n.n != "" {
-		b.WriteString(n.n)
-		b.WriteByte('/')
-	}
-	b.WriteString(n.m)
-	if n.t != "" {
-		b.WriteByte(':')
-		b.WriteString(n.t)
-	}
-	return b.String()
-}
-
-func (n Name) GoString() string {
-	return fmt.Sprintf("<Name %q %q %q %q>", n.h, n.n, n.m, n.t)
-}
-
-// cutLastAny is like strings.Cut but scans in reverse for the last character
-// in chars. If no character is found, before is the empty string and after is
-// s. The returned sep is the byte value of the character in chars if one was
-// found; otherwise it is 0.
-func cutLastAny(s, chars string) (before, after string, sep byte) {
-	i := strings.LastIndexAny(s, chars)
-	if i >= 0 {
-		return s[:i], s[i+1:], s[i]
-	}
-	return "", s, 0
-}

server/internal/internal/names/name_test.go

@@ -1,220 +0,0 @@
-package names
-
-import (
-	"strings"
-	"testing"
-)
-
-func TestParseName(t *testing.T) {
-	cases := []struct {
-		in   string
-		want Name
-	}{
-		{"", Name{}},
-		{"m:t", Name{m: "m", t: "t"}},
-		{"m", Name{m: "m"}},
-		{"/m", Name{m: "m"}},
-		{"/n/m:t", Name{n: "n", m: "m", t: "t"}},
-		{"n/m", Name{n: "n", m: "m"}},
-		{"n/m:t", Name{n: "n", m: "m", t: "t"}},
-		{"n/m", Name{n: "n", m: "m"}},
-		{"n/m", Name{n: "n", m: "m"}},
-		{strings.Repeat("m", MaxNameLength+1), Name{}},
-		{"h/n/m:t", Name{h: "h", n: "n", m: "m", t: "t"}},
-		{"ollama.com/library/_:latest", Name{h: "ollama.com", n: "library", m: "_", t: "latest"}},
-
-		// Invalids
-		// TODO: {"n:t/m:t", Name{}},
-		// TODO: {"/h/n/m:t", Name{}},
-	}
-
-	for _, tt := range cases {
-		t.Run(tt.in, func(t *testing.T) {
-			got := Parse(tt.in)
-			if got.Compare(tt.want) != 0 {
-				t.Errorf("parseName(%q) = %#v, want %q", tt.in, got, tt.want)
-			}
-		})
-	}
-}
-
-func TestString(t *testing.T) {
-	cases := []string{
-		"",
-		"m:t",
-		"m:t",
-		"m",
-		"n/m",
-		"n/m:t",
-		"n/m",
-		"n/m",
-		"h/n/m:t",
-		"ollama.com/library/_:latest",
-
-		// Special cased to "round trip" without the leading slash.
-		"/m",
-		"/n/m:t",
-	}
-	for _, s := range cases {
-		t.Run(s, func(t *testing.T) {
-			s = strings.TrimPrefix(s, "/")
-			if g := Parse(s).String(); g != s {
-				t.Errorf("parse(%q).String() = %q", s, g)
-			}
-		})
-	}
-}
-
-func TestParseExtended(t *testing.T) {
-	cases := []struct {
-		in string
-
-		wantScheme string
-		wantName   Name
-		wantDigest string
-	}{
-		{"", "", Name{}, ""},
-		{"m", "", Name{m: "m"}, ""},
-		{"http://m", "http", Name{m: "m"}, ""},
-		{"http+insecure://m", "http+insecure", Name{m: "m"}, ""},
-		{"http://m@sha256:deadbeef", "http", Name{m: "m"}, "sha256:deadbeef"},
-	}
-	for _, tt := range cases {
-		t.Run(tt.in, func(t *testing.T) {
-			scheme, name, digest := Split(tt.in)
-			n := Parse(name)
-			if scheme != tt.wantScheme || n.Compare(tt.wantName) != 0 || digest != tt.wantDigest {
-				t.Errorf("ParseExtended(%q) = %q, %#v, %q, want %q, %#v, %q", tt.in, scheme, name, digest, tt.wantScheme, tt.wantName, tt.wantDigest)
-			}
-		})
-	}
-}
-
-func TestMerge(t *testing.T) {
-	cases := []struct {
-		a, b string
-		want string
-	}{
-		{"", "", ""},
-		{"m", "", "m"},
-		{"", "m", ""},
-		{"x", "y", "x"},
-		{"o.com/n/m:t", "o.com/n/m:t", "o.com/n/m:t"},
-		{"o.com/n/m:t", "o.com/n/_:t", "o.com/n/m:t"},
-
-		{"bmizerany/smol", "ollama.com/library/_:latest", "ollama.com/bmizerany/smol:latest"},
-		{"localhost:8080/bmizerany/smol", "ollama.com/library/_:latest", "localhost:8080/bmizerany/smol:latest"},
-	}
-	for _, tt := range cases {
-		t.Run("", func(t *testing.T) {
-			a, b := Parse(tt.a), Parse(tt.b)
-			got := Merge(a, b)
-			if got.Compare(Parse(tt.want)) != 0 {
-				t.Errorf("merge(%q, %q) = %#v, want %q", tt.a, tt.b, got, tt.want)
-			}
-		})
-	}
-}
-
-func TestParseStringRoundTrip(t *testing.T) {
-	cases := []string{
-		"",
-		"m",
-		"m:t",
-		"n/m",
-		"n/m:t",
-		"n/m:t",
-		"n/m",
-		"n/m",
-		"h/n/m:t",
-		"ollama.com/library/_:latest",
-	}
-	for _, s := range cases {
-		t.Run(s, func(t *testing.T) {
-			if got := Parse(s).String(); got != s {
-				t.Errorf("parse(%q).String() = %q", s, got)
-			}
-		})
-	}
-}
-
-var junkName Name
-
-func BenchmarkParseName(b *testing.B) {
-	b.ReportAllocs()
-	for range b.N {
-		junkName = Parse("h/n/m:t")
-	}
-}
-
-const (
-	part80  = "88888888888888888888888888888888888888888888888888888888888888888888888888888888"
-	part350 = "33333333333333333333333333333333333333333333333333333333333333333333333333333333333333333333333333333333333333333333333333333333333333333333333333333333333333333333333333333333333333333333333333333333333333333333333333333333333333333333333333333333333333333333333333333333333333333333333333333333333333333333333333333333333333333333333333333333333333"
-)
-
-var testCases = map[string]bool{ // name -> valid
-	"": false,
-
-	"_why/_the/_lucky:_stiff": true,
-
-	// minimal
-	"h/n/m:t": true,
-
-	"host/namespace/model:tag": true,
-	"host/namespace/model":     true,
-	"namespace/model":          true,
-	"model":                    true,
-
-	// long (but valid)
-	part80 + "/" + part80 + "/" + part80 + ":" + part80:  true,
-	part350 + "/" + part80 + "/" + part80 + ":" + part80: true,
-
-	// too long
-	part80 + "/" + part80 + "/" + part80 + ":" + part350:       false,
-	"x" + part350 + "/" + part80 + "/" + part80 + ":" + part80: false,
-
-	"h/nn/mm:t": true, // bare minimum part sizes
-
-	// unqualified
-	"m":     true,
-	"n/m:":  true,
-	"h/n/m": true,
-	"@t":    false,
-	"m@d":   false,
-
-	// invalids
-	"^":      false,
-	"mm:":    true,
-	"/nn/mm": true,
-	"//":     false, // empty model
-	"//mm":   true,
-	"hh//":   false, // empty model
-	"//mm:@": false,
-	"00@":    false,
-	"@":      false,
-
-	// not starting with alphanum
-	"-hh/nn/mm:tt": false,
-	"hh/-nn/mm:tt": false,
-	"hh/nn/-mm:tt": false,
-	"hh/nn/mm:-tt": false,
-
-	// smells like a flag
-	"-h": false,
-
-	// hosts
-	"host:https/namespace/model:tag": true,
-
-	// colon in non-host part before tag
-	"host/name:space/model:tag": false,
-}
-
-func TestParseNameValidation(t *testing.T) {
-	for s, valid := range testCases {
-		got := Parse(s)
-		if got.IsValid() != valid {
-			t.Logf("got: %v", got)
-			t.Errorf("Parse(%q).IsValid() = %v; want !%[2]v", s, got.IsValid())
-		}
-	}
-}