diff --git a/.github/workflows/test.yaml b/.github/workflows/test.yaml index 0548ff286e8..9c6138ccc6f 100644 --- a/.github/workflows/test.yaml +++ b/.github/workflows/test.yaml @@ -115,7 +115,7 @@ jobs: superbuild_args: '-DOLLAMA_LLAMA_BACKENDS=vulkan' expected_payload: lib/ollama/vulkan/libggml-vulkan.so - preset: 'MLX CUDA 13' - container: nvidia/cuda:13.0.0-devel-ubuntu22.04 + container: nvidia/cuda:13.0.0-devel-ubuntu24.04 extra-packages: libcudnn9-dev-cuda-13 libopenblas-dev liblapack-dev liblapacke-dev git curl superbuild_target: ollama-mlx-cuda_v13 superbuild_dir: build/local-superbuild-mlx-cuda_v13 @@ -386,7 +386,6 @@ jobs: go-version-file: go.mod - name: Verify Go dependency licenses run: | - # See cmake/generate_go_license.cmake for special case handling. cmake -S . -B build/go-license \ -DOLLAMA_LLAMA_BACKENDS= \ -DOLLAMA_MLX_BACKENDS= \ diff --git a/.golangci.yaml b/.golangci.yaml index 5a4254132ba..2b843b6cea9 100644 --- a/.golangci.yaml +++ b/.golangci.yaml @@ -5,6 +5,7 @@ linters: - bidichk - bodyclose - containedctx + - depguard - gocheckcompilerdirectives - intrange - makezero @@ -20,6 +21,15 @@ linters: - errcheck - usestdlibvars settings: + depguard: + rules: + test-only-packages: + files: + - $all + - "!$test" + deny: + - pkg: github.com/ollama/ollama/internal/testutil + desc: test helpers may only be imported by test files govet: disable: - unusedresult diff --git a/Dockerfile b/Dockerfile index c18d0656aed..3a6abc33187 100644 --- a/Dockerfile +++ b/Dockerfile @@ -220,6 +220,7 @@ ENV CGO_LDFLAGS="-L/usr/local/cuda-13/lib64 -L/usr/local/cuda-13/targets/x86_64- WORKDIR /go/src/github.com/ollama/ollama COPY CMakeLists.txt CMakePresets.json . COPY cmake cmake +COPY mlx/compat mlx/compat COPY x/mlxrunner/mlx x/mlxrunner/mlx COPY x/mlxrunner/xgrammar/native x/mlxrunner/xgrammar/native COPY go.mod go.sum . diff --git a/LLAMA_CPP_VERSION b/LLAMA_CPP_VERSION index 322774ed9da..f533d35b8f5 100644 --- a/LLAMA_CPP_VERSION +++ b/LLAMA_CPP_VERSION @@ -1 +1 @@ -b10760 +b10864 diff --git a/MLX_C_VERSION b/MLX_C_VERSION index 3f98d19af6e..aaf310acb8c 100644 --- a/MLX_C_VERSION +++ b/MLX_C_VERSION @@ -1 +1 @@ -c74db5307cc8ce122f48d97ef951b30578674e7f \ No newline at end of file +ebc88f10caa1b625e6b581437a8dea6df8a70085 \ No newline at end of file diff --git a/MLX_VERSION b/MLX_VERSION index 58b7f7c0057..e96224e66b9 100644 --- a/MLX_VERSION +++ b/MLX_VERSION @@ -1 +1 @@ -37c26e5755da637255d57ea34b4879196a485301 \ No newline at end of file +d9add9d11f3154111a4c85f267ec2fd307ecd18e \ No newline at end of file diff --git a/agent/approval.go b/agent/approval.go deleted file mode 100644 index 47b68e413ea..00000000000 --- a/agent/approval.go +++ /dev/null @@ -1,198 +0,0 @@ -package agent - -import ( - "context" - "strings" - "sync" -) - -type ApprovalRequest struct { - WorkingDir string - Calls []ApprovalToolCall -} - -func (r *ApprovalRequest) AddToolCall(id, name, scope string, args map[string]any) { - r.Calls = append(r.Calls, ApprovalToolCall{ - ToolCallID: id, - ToolName: name, - Args: args, - ApprovalScope: scope, - }) -} - -type ApprovalToolCall struct { - ToolCallID string - ToolName string - Args map[string]any - ApprovalScope string -} - -type Approval struct { - Allow bool - AllowAll bool - AllowScopes []string - Reason string -} - -type ApprovalPrompter interface { - PromptApproval(context.Context, ApprovalRequest) (Approval, error) -} - -type ApprovalState struct { - mu sync.RWMutex - allowAll bool - scopes map[string]bool -} - -func (s *ApprovalState) Set(allowAll bool, scopes map[string]bool) { - if s == nil { - return - } - s.mu.Lock() - defer s.mu.Unlock() - s.allowAll = allowAll - s.scopes = cloneApprovalScopes(scopes) -} - -// GrantAll grants blanket approval for all future tool calls. -func (s *ApprovalState) GrantAll() { - if s == nil { - return - } - s.mu.Lock() - defer s.mu.Unlock() - s.allowAll = true -} - -// AllGranted reports whether blanket approval has been granted. -func (s *ApprovalState) AllGranted() bool { - if s == nil { - return false - } - s.mu.RLock() - defer s.mu.RUnlock() - return s.allowAll -} - -func (s *ApprovalState) Allows(scope string) bool { - if s == nil { - return false - } - s.mu.RLock() - defer s.mu.RUnlock() - return s.allowAll || s.scopes[scope] -} - -// Apply merges an approval's scopes and allow-all flag into the state. It -// returns true if the approval grants permission (allow-all or at least one -// scope). It does not mutate the approval; the caller sets Allow based on the -// returned value. -func (s *ApprovalState) Apply(result *Approval) bool { - if s == nil || result == nil { - return false - } - s.mu.Lock() - defer s.mu.Unlock() - granted := false - if result.AllowAll { - s.allowAll = true - granted = true - } - if len(result.AllowScopes) > 0 { - granted = true - s.grantScopesLocked(result.AllowScopes) - } - return granted -} - -// GrantScopes merges the given scopes into the state. -func (s *ApprovalState) GrantScopes(scopes []string) { - if s == nil { - return - } - s.mu.Lock() - defer s.mu.Unlock() - s.grantScopesLocked(scopes) -} - -// grantScopesLocked adds trimmed, non-empty scopes to the state. Caller must -// hold s.mu. -func (s *ApprovalState) grantScopesLocked(scopes []string) { - if s.scopes == nil { - s.scopes = make(map[string]bool, len(scopes)) - } - for _, scope := range scopes { - scope = strings.TrimSpace(scope) - if scope != "" { - s.scopes[scope] = true - } - } -} - -func cloneApprovalScopes(src map[string]bool) map[string]bool { - if len(src) == 0 { - return nil - } - dst := make(map[string]bool, len(src)) - for scope, allowed := range src { - if allowed { - dst[scope] = true - } - } - return dst -} - -func (s *Session) needsApproval(tool Tool, name string, args map[string]any) bool { - return ToolRequiresApproval(tool, args) && !s.allows(toolApprovalScope(tool, name, args)) -} - -// allows reports whether scope is permitted by the session's accumulated approval state. -func (s *Session) allows(scope string) bool { - if s == nil || s.ApprovalState == nil { - return false - } - return s.ApprovalState.Allows(scope) -} - -// applyApproval merges an approval result into the session's state and marks -// the result as allowed when scopes or allow-all were granted. -func (s *Session) applyApproval(result *Approval) { - if s == nil || result == nil { - return - } - if s.ApprovalState == nil { - s.ApprovalState = &ApprovalState{} - } - if s.ApprovalState.Apply(result) { - result.Allow = true - } -} - -func (s *Session) authorizeToolCalls(ctx context.Context, req ApprovalRequest) (Approval, error) { - if s == nil || len(req.Calls) == 0 || (s.ApprovalState != nil && s.ApprovalState.AllGranted()) { - return Approval{Allow: true}, nil - } - if s.ApprovalPrompter == nil { - return Approval{ - Reason: "Tool execution requires approval, but no approval prompter is available.", - }, nil - } - - result, err := s.ApprovalPrompter.PromptApproval(ctx, req) - if err != nil { - return Approval{}, err - } - s.applyApproval(&result) - return result, nil -} - -// toolApprovalScope returns the approval scope key for a tool invocation. -// If the tool implements ScopedTool, its ApprovalScope method determines the -// scope (e.g. shell tools scope to "\x00"). Otherwise the scope -// is the trimmed tool name. -func toolApprovalScope(tool Tool, toolName string, args map[string]any) string { - if scoped, ok := tool.(ScopedTool); ok { - return scoped.ApprovalScope(args) - } - return strings.TrimSpace(toolName) -} diff --git a/agent/approval_test.go b/agent/approval_test.go deleted file mode 100644 index 726a3890885..00000000000 --- a/agent/approval_test.go +++ /dev/null @@ -1,95 +0,0 @@ -package agent - -import ( - "context" - "strings" - "testing" - - "github.com/ollama/ollama/api" -) - -type mockTool struct { - name string -} - -func (m mockTool) Name() string { return m.name } -func (m mockTool) Description() string { return "" } -func (m mockTool) Schema() api.ToolFunction { - return api.ToolFunction{Name: m.name} -} - -func (m mockTool) Execute(context.Context, ToolContext, map[string]any) (ToolResult, error) { - return ToolResult{}, nil -} - -func TestToolApprovalScopeUsesScopedTool(t *testing.T) { - shellTool := mockScopedTool{ - mockTool: mockTool{name: "bash"}, - scope: func(args map[string]any) string { - if cmd, ok := args["command"].(string); ok { - cmd = strings.TrimSpace(cmd) - if cmd != "" { - return "bash\x00" + cmd - } - } - return "bash" - }, - } - plainTool := mockTool{name: "edit"} - - tests := []struct { - tool Tool - name string - args map[string]any - want string - }{ - {shellTool, "bash", map[string]any{"command": " pwd "}, "bash\x00pwd"}, - {shellTool, "bash", map[string]any{"command": "Get-ChildItem"}, "bash\x00Get-ChildItem"}, - {plainTool, "edit", map[string]any{"path": "README.md"}, "edit"}, - } - for _, tt := range tests { - if got := toolApprovalScope(tt.tool, tt.name, tt.args); got != tt.want { - t.Fatalf("toolApprovalScope(%q) = %q, want %q", tt.name, got, tt.want) - } - } -} - -type mockScopedTool struct { - mockTool - scope func(args map[string]any) string -} - -func (m mockScopedTool) ApprovalScope(args map[string]any) string { - return m.scope(args) -} - -func TestSessionApplyApprovalScopes(t *testing.T) { - session := &Session{} - result := Approval{AllowScopes: []string{"edit", "bash\x00pwd", " "}} - - session.applyApproval(&result) - - if !result.Allow { - t.Fatal("scoped approval should allow the current request") - } - if !session.allows("edit") || !session.allows("bash\x00pwd") { - t.Fatal("scoped approval was not saved") - } - if session.allows("bash") || session.allows("bash\x00ls") { - t.Fatal("shell approval was too broad") - } - if session.ApprovalState.AllGranted() { - t.Fatal("allow all = true, want false for scoped approval") - } -} - -func TestSessionApplyApprovalAllowAll(t *testing.T) { - session := &Session{} - result := Approval{AllowAll: true} - - session.applyApproval(&result) - - if !result.Allow || !session.allows("anything") { - t.Fatalf("allow all = %v result = %#v, want allow all", session.ApprovalState.AllGranted(), result) - } -} diff --git a/agent/compactor.go b/agent/compactor.go deleted file mode 100644 index e0f6ce454d1..00000000000 --- a/agent/compactor.go +++ /dev/null @@ -1,667 +0,0 @@ -package agent - -import ( - "context" - "encoding/json" - "errors" - "fmt" - "net/http" - "strings" - - "github.com/ollama/ollama/api" -) - -// Compaction wire-format. These constants and helpers are the single canonical -// definition of how a compacted turn is represented in message history. -const ( - CompactionSummaryMessagePrefix = "Conversation summary:\n" - CompactionToolName = "summary" - CompactionToolCallID = "ollama_compaction" - CompactionContinueInstruction = "continue the task in progress. the history has been compacted, do not mention compaction to the user" -) - -const ( - defaultCompactionContextWindowTokens = 32768 - defaultCompactionKeepUserTurns = 3 - defaultCompactionThreshold = 0.8 - compactOnlySummaryContextTokens = 16000 - - maxCompactionSummaryRunes = 16 * 1024 - - compactionSystemPrompt = "Summarize the archived part of an Ollama agent conversation. Preserve user goals, decisions, files, commands, tool results, and unresolved tasks needed to continue. Omit private reasoning and return only the summary." -) - -type Compactor interface { - MaybeCompact(context.Context, CompactionRequest) (CompactionResult, error) - - // ContextWindowTokens returns the effective context window size in - // tokens, resolving runtime options against configured defaults. - ContextWindowTokens(options map[string]any) int - - // Threshold returns the compaction threshold as a fraction of the - // context window (e.g. 0.8 means compact at 80% capacity). - Threshold() float64 - - // ShouldCompact reports whether a compaction should run and returns the - // trigger reason. An empty trigger means compaction is not needed. - ShouldCompact(req CompactionRequest) (trigger string, should bool) -} - -type CompactionOptions struct { - ContextWindowTokens int - KeepUserTurns int - Threshold float64 -} - -type CompactionRequest struct { - ChatID string - Model string - SystemPrompt string - Messages []api.Message - Tools api.Tools - Format string - Latest api.ChatResponse - Options map[string]any - KeepAlive *api.Duration - Think *api.ThinkValue - Force bool - ContinueTask bool - KeepUserTurns *int - Progress func(CompactionProgress) -} - -type CompactionProgress struct { - Tokens int -} - -type CompactionResult struct { - Messages []api.Message - Compacted bool - Due bool - Summary string - Reason string -} - -type SimpleCompactor struct { - Client ChatClient - Options CompactionOptions -} - -func (c *SimpleCompactor) MaybeCompact(ctx context.Context, req CompactionRequest) (CompactionResult, error) { - result := CompactionResult{Messages: req.Messages} - if c == nil { - return result, nil - } - - result.Due = req.Force || c.shouldCompact(req) - if !result.Due { - return result, nil - } - if c.Client == nil { - result.Reason = "compaction is unavailable" - return result, nil - } - - keepUserTurns := c.keepUserTurns(req.Options) - if req.KeepUserTurns != nil { - keepUserTurns = *req.KeepUserTurns - } - prefix, previousSummary, archive, suffix, _, ok := splitCompactionMessages(req.Messages, keepUserTurns) - if !ok || len(archive) == 0 { - result.Reason = "nothing to compact" - return result, nil - } - - summary, err := c.summarize(ctx, req, previousSummary, archive) - if err != nil { - result.Reason = err.Error() - return result, err - } - summary = truncateCompactionSummary(strings.TrimSpace(summary)) - if summary == "" { - summary, err = c.summarizeEmptyFallback(ctx, req, previousSummary, archive) - if err != nil { - result.Reason = err.Error() - return result, err - } - summary = truncateCompactionSummary(strings.TrimSpace(summary)) - } - if summary == "" { - result.Reason = "summary was empty" - return result, nil - } - - compacted := make([]api.Message, 0, len(prefix)+len(suffix)+2) - compacted = append(compacted, prefix...) - compacted = append(compacted, CompactionSummaryMessages(summary, req.ContinueTask)...) - compacted = append(compacted, suffix...) - result.Messages = compacted - result.Compacted = true - result.Summary = summary - return result, nil -} - -func (c *SimpleCompactor) shouldCompact(req CompactionRequest) bool { - contextWindow := c.contextWindowTokens(req.Options) - threshold := int(float64(contextWindow) * c.threshold()) - if threshold <= 0 { - return false - } - if req.Latest.PromptEvalCount > 0 && req.Latest.PromptEvalCount >= threshold { - return true - } - return estimateCompactionRequestTokens(req) >= threshold -} - -func (c *SimpleCompactor) contextWindowTokens(options map[string]any) int { - return ResolveContextWindowTokens(options, c.Options.ContextWindowTokens) -} - -// ContextWindowTokens resolves the effective context window from runtime -// options or configured defaults. Satisfies the Compactor interface. -func (c *SimpleCompactor) ContextWindowTokens(options map[string]any) int { - if c == nil { - return 0 - } - return c.contextWindowTokens(options) -} - -func (c *SimpleCompactor) threshold() float64 { - return ResolveCompactionThreshold(c.Options.Threshold) -} - -// Threshold returns the configured compaction threshold fraction. Satisfies -// the Compactor interface. -func (c *SimpleCompactor) Threshold() float64 { - if c == nil { - return 0 - } - return c.threshold() -} - -// ShouldCompact reports whether compaction is due and the trigger reason. -// Satisfies the Compactor interface. -func (c *SimpleCompactor) ShouldCompact(req CompactionRequest) (string, bool) { - if c == nil { - return "", false - } - if req.Force { - return "force", true - } - if c.shouldCompact(req) { - contextWindow := c.contextWindowTokens(req.Options) - threshold := int(float64(contextWindow) * c.threshold()) - if req.Latest.PromptEvalCount > 0 && req.Latest.PromptEvalCount >= threshold { - return "prompt_eval", true - } - return "estimate", true - } - return "", false -} - -func (c *SimpleCompactor) keepUserTurns(options map[string]any) int { - contextWindow := c.contextWindowTokens(options) - if contextWindow > 0 && contextWindow < compactOnlySummaryContextTokens { - return 0 - } - if c.Options.KeepUserTurns > 0 { - return c.Options.KeepUserTurns - } - return defaultCompactionKeepUserTurns -} - -func ResolveContextWindowTokens(options map[string]any, configured int) int { - if n := intOption(options, "num_ctx"); n > 0 { - return n - } - if configured > 0 { - return configured - } - return defaultCompactionContextWindowTokens -} - -func ResolveCompactionThreshold(configured float64) float64 { - if configured > 0 { - return configured - } - return defaultCompactionThreshold -} - -func (c *SimpleCompactor) summarize(ctx context.Context, req CompactionRequest, previousSummary string, archive []api.Message) (string, error) { - body, err := compactionPrompt(previousSummary, archive, c.compactionPromptBodyBudgetTokens(req.Options)) - if err != nil { - return "", err - } - - chatReq := &api.ChatRequest{ - Model: req.Model, - Messages: []api.Message{ - { - Role: "system", - Content: compactionSystemPrompt, - }, - { - Role: "user", - Content: body, - }, - }, - Options: req.Options, - Think: req.Think, - } - if req.KeepAlive != nil { - chatReq.KeepAlive = req.KeepAlive - } - - var summary strings.Builder - if err := c.Client.Chat(ctx, chatReq, func(response api.ChatResponse) error { - summary.WriteString(response.Message.Content) - if req.Progress != nil { - tokens := response.EvalCount - if tokens <= 0 { - tokens = estimateCompactionTokens(summary.String()) - } - req.Progress(CompactionProgress{Tokens: tokens}) - } - return nil - }); err != nil { - return "", err - } - return summary.String(), nil -} - -func (c *SimpleCompactor) summarizeEmptyFallback(ctx context.Context, req CompactionRequest, previousSummary string, archive []api.Message) (string, error) { - retry := req - retry.Think = &api.ThinkValue{Value: false} - summary, err := c.summarize(ctx, retry, previousSummary, archive) - if err == nil { - return summary, nil - } - if !isUnsupportedCompactionThinkError(err) { - return "", err - } - if req.Think == nil { - return "", nil - } - retry.Think = nil - return c.summarize(ctx, retry, previousSummary, archive) -} - -func isUnsupportedCompactionThinkError(err error) bool { - if err == nil { - return false - } - text := strings.ToLower(err.Error()) - if !strings.Contains(text, "think") { - return false - } - var statusErr api.StatusError - if errors.As(err, &statusErr) && statusErr.StatusCode != 0 { - return statusErr.StatusCode == http.StatusBadRequest - } - return strings.Contains(text, "does not support") || strings.Contains(text, "not supported") || strings.Contains(text, "unsupported") -} - -// compactionSummaryMessageForTask renders a compaction summary as the content -// string stored on the synthetic tool-result message. -func compactionSummaryMessageForTask(summary string, continueTask bool) string { - content := CompactionSummaryMessagePrefix + strings.TrimSpace(summary) - if continueTask { - content = strings.TrimSpace(content) + "\n\n" + CompactionContinueInstruction - } - return content -} - -// CompactionSummaryMessages renders a compaction summary as the assistant -// tool-call plus tool-result pair that represents a compacted turn in the -// message history. -func CompactionSummaryMessages(summary string, continueTask bool) []api.Message { - return []api.Message{ - { - Role: "assistant", - ToolCalls: []api.ToolCall{{ - ID: CompactionToolCallID, - Function: api.ToolCallFunction{ - Name: CompactionToolName, - }, - }}, - }, - { - Role: "tool", - ToolName: CompactionToolName, - ToolCallID: CompactionToolCallID, - Content: compactionSummaryMessageForTask(summary, continueTask), - }, - } -} - -func (c *SimpleCompactor) compactionPromptBodyBudgetTokens(options map[string]any) int { - contextWindow := c.contextWindowTokens(options) - threshold := int(float64(contextWindow) * c.threshold()) - if threshold <= 0 { - return 0 - } - systemTokens := estimateCompactionTokens("system") + estimateCompactionTokens(compactionSystemPrompt) - userRoleTokens := estimateCompactionTokens("user") - budget := threshold - systemTokens - userRoleTokens - if budget <= 0 { - return 0 - } - return budget -} - -func truncateCompactionSummary(summary string) string { - return Truncate(summary, TruncateConfig{ - MaxRunes: maxCompactionSummaryRunes, - Label: "summary", - }) -} - -func estimateCompactionTokens(text string) int { - text = strings.TrimSpace(text) - if text == "" { - return 0 - } - return ApproximateTokens(len([]rune(text))) -} - -func estimateMessagesTokens(messages []api.Message) int { - var total int - for _, msg := range messages { - total += estimateCompactionTokens(msg.Role) - total += estimateCompactionTokens(msg.Content) - total += estimateCompactionTokens(msg.Thinking) - total += estimateCompactionTokens(msg.ToolName) - total += estimateCompactionTokens(msg.ToolCallID) - for _, call := range msg.ToolCalls { - total += estimateCompactionTokens(call.Function.Name) - total += estimateCompactionTokens(call.Function.Arguments.String()) - } - } - return total -} - -func estimateCompactionRequestTokens(req CompactionRequest) int { - requestMessages := sanitizeMessagesForEstimate(req.Messages) - if strings.TrimSpace(req.SystemPrompt) != "" { - requestMessages = make([]api.Message, 0, len(req.Messages)+1) - requestMessages = append(requestMessages, api.Message{Role: "system", Content: strings.TrimSpace(req.SystemPrompt)}) - requestMessages = append(requestMessages, sanitizeMessagesForEstimate(req.Messages)...) - } - - payload := struct { - Messages []api.Message `json:"messages,omitempty"` - Tools api.Tools `json:"tools,omitempty"` - Format json.RawMessage `json:"format,omitempty"` - }{ - Messages: requestMessages, - Tools: req.Tools, - } - if rawFormat, ok := compactionFormatForEstimate(req.Format); ok { - payload.Format = rawFormat - } - if data, err := json.Marshal(payload); err == nil { - return estimateCompactionTokens(string(data)) - } - - total := estimateMessagesTokens(requestMessages) - total += estimateCompactionTokens(req.Tools.String()) - total += estimateCompactionTokens(req.Format) - return total -} - -func (s *Session) estimateRunPromptTokens(opts RunOptions, messages []api.Message) int { - return estimateCompactionRequestTokens(CompactionRequest{ - SystemPrompt: opts.SystemPrompt, - Messages: messages, - Tools: s.availableTools(), - Format: opts.Format, - Options: opts.Options, - }) -} - -func (s *Session) checkPreflightPromptBudget(opts RunOptions, messages []api.Message) error { - contextWindow := s.contextWindowTokens(opts) - if contextWindow <= 0 { - return nil - } - estimated := s.estimateRunPromptTokens(opts, messages) - if estimated < contextWindow { - return nil - } - return fmt.Errorf("prompt is too large for the current context (~%d/%d tokens). Reduce the system prompt or message history, compact the conversation, or use a model with a larger context", estimated, contextWindow) -} - -func (s *Session) checkPostCompactionPromptBudget(opts RunOptions, messages []api.Message) error { - contextWindow := s.contextWindowTokens(opts) - if contextWindow <= 0 { - return nil - } - estimated := s.estimateRunPromptTokens(opts, messages) - if estimated < contextWindow { - return nil - } - return fmt.Errorf("history is still too large after compaction (~%d/%d tokens). Start a fresh request, reduce the system prompt or history, or use a model with a larger context", estimated, contextWindow) -} - -func sanitizeMessagesForEstimate(messages []api.Message) []api.Message { - requestMessages := sanitizeMessagesForRequest(messages) - for i := range requestMessages { - // Image token accounting is model-specific. Without the active model's - // tokenizer and vision accounting, raw image bytes/base64 make the - // estimate look much larger than the prompt the model actually sees. - requestMessages[i].Images = nil - } - return requestMessages -} - -func compactionFormatForEstimate(format string) (json.RawMessage, bool) { - format = strings.TrimSpace(format) - if format == "" { - return nil, false - } - if format == "json" { - return json.RawMessage(`"json"`), true - } - if !json.Valid([]byte(format)) { - return nil, false - } - return json.RawMessage(format), true -} - -func compactionPrompt(previousSummary string, archive []api.Message, maxTokens int) (string, error) { - messages := make([]api.Message, 0, len(archive)) - for _, msg := range archive { - msg.Thinking = "" - msg.Images = nil - messages = append(messages, msg) - } - return renderCompactionPrompt(previousSummary, fitCompactionMessagesToBudget(previousSummary, messages, maxTokens)) -} - -func renderCompactionPrompt(previousSummary string, messages []api.Message) (string, error) { - payload, err := json.MarshalIndent(messages, "", " ") - if err != nil { - return "", fmt.Errorf("marshal compaction messages: %w", err) - } - - var b strings.Builder - if strings.TrimSpace(previousSummary) != "" { - b.WriteString("Previous summary:\n") - b.WriteString(strings.TrimSpace(previousSummary)) - b.WriteString("\n\n") - } - b.WriteString("Messages to archive as JSON:\n") - b.Write(payload) - return b.String(), nil -} - -func fitCompactionMessagesToBudget(previousSummary string, messages []api.Message, maxTokens int) []api.Message { - if maxTokens <= 0 { - return messages - } - fitted := append([]api.Message(nil), messages...) - for range 16 { - body, err := renderCompactionPrompt(previousSummary, fitted) - if err != nil || estimateCompactionTokens(body) <= maxTokens { - return fitted - } - - idx := largestCompactionContentMessage(fitted) - if idx < 0 { - return fitted - } - overageTokens := estimateCompactionTokens(body) - maxTokens - currentRunes := len([]rune(fitted[idx].Content)) - nextRunes := currentRunes - overageTokens*4 - 256 - if nextRunes >= currentRunes { - nextRunes = currentRunes / 2 - } - fitted[idx].Content = truncateToolResultContentTo(fitted[idx].Content, nextRunes) - } - return fitted -} - -func largestCompactionContentMessage(messages []api.Message) int { - idx := -1 - size := 0 - for i, msg := range messages { - n := len([]rune(msg.Content)) - if n > size { - idx = i - size = n - } - } - return idx -} - -func splitCompactionMessages(messages []api.Message, keepUserTurns int) (prefix []api.Message, previousSummary string, archive []api.Message, suffix []api.Message, keptUserTurns int, ok bool) { - if keepUserTurns < 0 { - keepUserTurns = defaultCompactionKeepUserTurns - } - - start := 0 - for start < len(messages) && messages[start].Role == "system" && !isCompactionSummary(messages[start]) { - prefix = append(prefix, messages[start]) - start++ - } - - candidates := make([]api.Message, 0, len(messages)-start) - for i := start; i < len(messages); i++ { - msg := messages[i] - if isCompactionSummary(msg) { - previousSummary = CompactionSummaryText(msg.Content) - continue - } - if isCompactionToolCall(msg) { - if i+1 < len(messages) && isCompactionSummary(messages[i+1]) { - previousSummary = CompactionSummaryText(messages[i+1].Content) - i++ - } - continue - } - candidates = append(candidates, msg) - } - - userTurnIndexes := make([]int, 0, keepUserTurns) - for i := len(candidates) - 1; i >= 0; i-- { - if candidates[i].Role == "user" { - userTurnIndexes = append(userTurnIndexes, i) - } - } - keptUserTurns = keepUserTurns - if len(userTurnIndexes) <= keptUserTurns { - keptUserTurns = len(userTurnIndexes) - 1 - } - if keptUserTurns < 0 { - keptUserTurns = 0 - } - - suffixStart := len(candidates) - if keptUserTurns > 0 { - suffixStart = userTurnIndexes[keptUserTurns-1] - } - if suffixStart <= 0 || len(candidates[:suffixStart]) == 0 { - return prefix, previousSummary, nil, nil, keptUserTurns, false - } - - return prefix, previousSummary, candidates[:suffixStart], candidates[suffixStart:], keptUserTurns, true -} - -func isCompactionToolName(name string) bool { - return name == CompactionToolName -} - -func isCompactionSummary(msg api.Message) bool { - return (msg.Role == "user" || msg.Role == "system" || (msg.Role == "tool" && isCompactionToolName(msg.ToolName))) && - strings.HasPrefix(msg.Content, CompactionSummaryMessagePrefix) -} - -// IsCompactionSummary reports whether msg uses the canonical compaction -// summary message representation. -func IsCompactionSummary(msg api.Message) bool { - return isCompactionSummary(msg) -} - -// CompactionSummaryContent returns the user-visible summary from msg when it -// is a canonical compaction summary. -func CompactionSummaryContent(msg api.Message) (string, bool) { - if !isCompactionSummary(msg) { - return "", false - } - return CompactionSummaryText(msg.Content), true -} - -// IsCompactionToolResult reports whether msg is the synthetic tool result used -// to represent compaction in message history. -func IsCompactionToolResult(msg api.Message) bool { - return msg.Role == "tool" && (isCompactionToolName(msg.ToolName) || msg.ToolCallID == CompactionToolCallID) -} - -// IsCompactionToolCall reports whether msg is the synthetic assistant tool -// call paired with a compaction summary result. -func IsCompactionToolCall(msg api.Message) bool { - return isCompactionToolCall(msg) -} - -func isCompactionToolCall(msg api.Message) bool { - if msg.Role != "assistant" { - return false - } - for _, call := range msg.ToolCalls { - if isCompactionToolName(call.Function.Name) { - return true - } - } - return false -} - -// CompactionSummaryText reverses CompactionSummaryMessages, returning the -// user-visible summary text with the prefix and any continuation instruction -// removed. -func CompactionSummaryText(content string) string { - return strings.TrimSpace(strings.TrimSuffix( - strings.TrimSpace(strings.TrimPrefix(content, CompactionSummaryMessagePrefix)), - CompactionContinueInstruction, - )) -} - -func intOption(options map[string]any, key string) int { - if options == nil { - return 0 - } - switch v := options[key].(type) { - case int: - return v - case int64: - return int(v) - case float64: - return int(v) - case float32: - return int(v) - case json.Number: - n, _ := v.Int64() - return int(n) - default: - return 0 - } -} diff --git a/agent/compactor_test.go b/agent/compactor_test.go deleted file mode 100644 index 79eba1f371d..00000000000 --- a/agent/compactor_test.go +++ /dev/null @@ -1,773 +0,0 @@ -package agent - -import ( - "context" - "net/http" - "strings" - "testing" - - "github.com/ollama/ollama/api" -) - -type scriptedCompactionClient struct { - responses [][]api.ChatResponse - errs []error - requests []*api.ChatRequest -} - -func (c *scriptedCompactionClient) Chat(_ context.Context, req *api.ChatRequest, fn api.ChatResponseFunc) error { - c.requests = append(c.requests, req) - i := len(c.requests) - 1 - if i < len(c.responses) { - for _, response := range c.responses[i] { - if err := fn(response); err != nil { - return err - } - } - } - if i < len(c.errs) { - return c.errs[i] - } - return nil -} - -func assertCompactionSummaryPair(t *testing.T, messages []api.Message) { - t.Helper() - if len(messages) != 2 { - t.Fatalf("compaction summary pair len = %d, want 2: %#v", len(messages), messages) - } - if messages[0].Role != "assistant" || len(messages[0].ToolCalls) != 1 || messages[0].ToolCalls[0].Function.Name != CompactionToolName { - t.Fatalf("compaction assistant message = %#v", messages[0]) - } - if messages[0].ToolCalls[0].Function.Arguments.Len() != 0 { - t.Fatalf("compaction summary tool call should not have arguments: %#v", messages[0].ToolCalls[0].Function.Arguments.ToMap()) - } - if messages[1].Role != "tool" || messages[1].ToolName != CompactionToolName || messages[1].ToolCallID != messages[0].ToolCalls[0].ID { - t.Fatalf("compaction tool result = %#v", messages[1]) - } - if !strings.HasPrefix(messages[1].Content, CompactionSummaryMessagePrefix) { - t.Fatalf("compaction tool result missing summary prefix: %#v", messages[1]) - } -} - -func TestSimpleCompactorSummarizesOldMessages(t *testing.T) { - client := &fakeClient{ - responses: [][]api.ChatResponse{{ - {Message: api.Message{Role: "assistant", Content: "summary"}}, - }}, - } - compactor := &SimpleCompactor{Client: client, Options: CompactionOptions{ - ContextWindowTokens: 16000, - KeepUserTurns: 2, - Threshold: 0.5, - }} - - messages := []api.Message{ - {Role: "system", Content: "stay pinned"}, - {Role: "user", Content: "old request"}, - {Role: "assistant", Content: "old answer", Thinking: "hidden"}, - {Role: "user", Content: "recent one"}, - {Role: "assistant", Content: "recent answer"}, - {Role: "user", Content: "recent two"}, - } - - result, err := compactor.MaybeCompact(context.Background(), CompactionRequest{ - ChatID: "chat-1", - Model: "model", - Messages: messages, - Latest: api.ChatResponse{Metrics: api.Metrics{PromptEvalCount: 12000}}, - }) - if err != nil { - t.Fatal(err) - } - if !result.Compacted { - t.Fatal("expected compaction") - } - compacted := result.Messages - if len(compacted) != 6 { - t.Fatalf("compacted messages = %d, want 6", len(compacted)) - } - if compacted[0].Content != "stay pinned" { - t.Fatalf("first message = %#v", compacted[0]) - } - if result.Summary != "summary" { - t.Fatalf("result summary = %q", result.Summary) - } - assertCompactionSummaryPair(t, compacted[1:3]) - if compacted[3].Content != "recent one" || compacted[5].Content != "recent two" { - t.Fatalf("recent turns were not kept: %#v", compacted) - } - if len(client.requests) != 1 { - t.Fatalf("summary requests = %d, want 1", len(client.requests)) - } - if strings.Contains(client.requests[0].Messages[1].Content, "hidden") { - t.Fatal("compaction prompt should omit thinking") - } -} - -func TestSimpleCompactorKeepsOnlySummaryForSmallContext(t *testing.T) { - client := &fakeClient{ - responses: [][]api.ChatResponse{{ - {Message: api.Message{Role: "assistant", Content: "small context summary"}}, - }}, - } - compactor := &SimpleCompactor{Client: client, Options: CompactionOptions{ - ContextWindowTokens: compactOnlySummaryContextTokens - 1, - KeepUserTurns: 3, - Threshold: 0.5, - }} - - result, err := compactor.MaybeCompact(context.Background(), CompactionRequest{ - ChatID: "chat-1", - Model: "model", - ContinueTask: true, - Messages: []api.Message{ - {Role: "system", Content: "pinned"}, - {Role: "user", Content: "old request"}, - {Role: "assistant", Content: "old answer"}, - {Role: "user", Content: "latest request"}, - }, - Latest: api.ChatResponse{Metrics: api.Metrics{PromptEvalCount: 12000}}, - }) - if err != nil { - t.Fatal(err) - } - if !result.Compacted { - t.Fatal("expected compaction") - } - if len(result.Messages) != 3 { - t.Fatalf("messages = %#v, want system plus compaction summary pair", result.Messages) - } - if result.Messages[0].Content != "pinned" { - t.Fatalf("leading system message not kept: %#v", result.Messages) - } - assertCompactionSummaryPair(t, result.Messages[1:]) - if !strings.Contains(result.Messages[2].Content, CompactionContinueInstruction) { - t.Fatalf("tool result missing continue instruction: %q", result.Messages[2].Content) - } -} - -func TestSimpleCompactorAddsContinueTaskInstructionOnlyToToolResult(t *testing.T) { - client := &fakeClient{ - responses: [][]api.ChatResponse{{ - {Message: api.Message{Role: "assistant", Content: "summary"}}, - }}, - } - compactor := &SimpleCompactor{Client: client, Options: CompactionOptions{ - ContextWindowTokens: 100, - KeepUserTurns: 1, - Threshold: 0.5, - }} - - result, err := compactor.MaybeCompact(context.Background(), CompactionRequest{ - ChatID: "chat-1", - Model: "model", - ContinueTask: true, - Messages: []api.Message{ - {Role: "user", Content: "old request"}, - {Role: "assistant", Content: "old answer"}, - {Role: "user", Content: "recent request"}, - }, - Latest: api.ChatResponse{Metrics: api.Metrics{PromptEvalCount: 75}}, - }) - if err != nil { - t.Fatal(err) - } - if result.Summary != "summary" { - t.Fatalf("result summary = %q", result.Summary) - } - content := result.Messages[1].Content - if !strings.Contains(content, CompactionContinueInstruction) { - t.Fatalf("tool result missing continue instruction: %q", content) - } - if got := CompactionSummaryText(content); got != "summary" { - t.Fatalf("visible summary text = %q", got) - } -} - -func TestSimpleCompactorTruncatesOversizedSummary(t *testing.T) { - longSummary := strings.Repeat("x", maxCompactionSummaryRunes+1024) - client := &fakeClient{ - responses: [][]api.ChatResponse{{ - {Message: api.Message{Role: "assistant", Content: longSummary}}, - }}, - } - compactor := &SimpleCompactor{Client: client, Options: CompactionOptions{ - ContextWindowTokens: 100, - KeepUserTurns: 1, - Threshold: 0.5, - }} - - result, err := compactor.MaybeCompact(context.Background(), CompactionRequest{ - ChatID: "chat-1", - Model: "model", - Messages: []api.Message{ - {Role: "user", Content: "old one"}, - {Role: "assistant", Content: "old answer"}, - {Role: "user", Content: "recent one"}, - }, - Latest: api.ChatResponse{Metrics: api.Metrics{PromptEvalCount: 75}}, - }) - if err != nil { - t.Fatal(err) - } - if !result.Compacted { - t.Fatal("expected compaction") - } - if runeCount := len([]rune(result.Summary)); runeCount > maxCompactionSummaryRunes+200 { - t.Fatalf("summary runes = %d, want <= %d (plus marker)", runeCount, maxCompactionSummaryRunes) - } - if !strings.Contains(result.Summary, "[summary truncated:") { - t.Fatalf("summary missing truncation marker: %q", result.Summary) - } - if !strings.Contains(result.Messages[1].Content, "[summary truncated:") { - t.Fatalf("compacted message missing truncation marker: %#v", result.Messages) - } -} - -func TestSimpleCompactorRetriesEmptySummaryWithThinkFalse(t *testing.T) { - client := &scriptedCompactionClient{ - responses: [][]api.ChatResponse{ - {{Message: api.Message{Role: "assistant", Thinking: "internal summary plan"}}}, - {{Message: api.Message{Role: "assistant", Content: "fallback summary"}}}, - }, - } - compactor := &SimpleCompactor{Client: client, Options: CompactionOptions{ - ContextWindowTokens: 100, - KeepUserTurns: 1, - Threshold: 0.5, - }} - - result, err := compactor.MaybeCompact(context.Background(), CompactionRequest{ - Model: "model", - Messages: []api.Message{ - {Role: "user", Content: "old request"}, - {Role: "assistant", Content: "old answer"}, - {Role: "user", Content: "recent request"}, - }, - Force: true, - }) - if err != nil { - t.Fatal(err) - } - if !result.Compacted || result.Summary != "fallback summary" { - t.Fatalf("compaction result = %#v", result) - } - if len(client.requests) != 2 { - t.Fatalf("summary requests = %d, want 2", len(client.requests)) - } - if client.requests[0].Think != nil { - t.Fatalf("first summary request think = %#v, want nil", client.requests[0].Think) - } - if client.requests[1].Think == nil || client.requests[1].Think.Value != false { - t.Fatalf("fallback summary request think = %#v, want false", client.requests[1].Think) - } -} - -func TestSimpleCompactorIgnoresUnsupportedThinkFalseFallback(t *testing.T) { - client := &scriptedCompactionClient{ - responses: [][]api.ChatResponse{ - {{Message: api.Message{Role: "assistant", Thinking: "internal summary plan"}}}, - nil, - }, - errs: []error{ - nil, - api.StatusError{StatusCode: http.StatusBadRequest, ErrorMessage: "model does not support thinking"}, - }, - } - compactor := &SimpleCompactor{Client: client, Options: CompactionOptions{ - ContextWindowTokens: 100, - KeepUserTurns: 1, - Threshold: 0.5, - }} - - result, err := compactor.MaybeCompact(context.Background(), CompactionRequest{ - Model: "model", - Messages: []api.Message{ - {Role: "user", Content: "old request"}, - {Role: "assistant", Content: "old answer"}, - {Role: "user", Content: "recent request"}, - }, - Force: true, - }) - if err != nil { - t.Fatal(err) - } - if result.Compacted || result.Reason != "summary was empty" { - t.Fatalf("compaction result = %#v", result) - } - if len(client.requests) != 2 { - t.Fatalf("summary requests = %d, want 2", len(client.requests)) - } - if client.requests[1].Think == nil || client.requests[1].Think.Value != false { - t.Fatalf("fallback summary request think = %#v, want false", client.requests[1].Think) - } -} - -func TestSimpleCompactorFallsBackToUnsetThinkWhenThinkFalseUnsupported(t *testing.T) { - client := &scriptedCompactionClient{ - responses: [][]api.ChatResponse{ - {{Message: api.Message{Role: "assistant", Thinking: "internal summary plan"}}}, - nil, - {{Message: api.Message{Role: "assistant", Content: "unset think summary"}}}, - }, - errs: []error{ - nil, - api.StatusError{StatusCode: http.StatusBadRequest, ErrorMessage: "think level is not supported"}, - nil, - }, - } - compactor := &SimpleCompactor{Client: client, Options: CompactionOptions{ - ContextWindowTokens: 100, - KeepUserTurns: 1, - Threshold: 0.5, - }} - thinkHigh := &api.ThinkValue{Value: "high"} - - result, err := compactor.MaybeCompact(context.Background(), CompactionRequest{ - Model: "model", - Messages: []api.Message{ - {Role: "user", Content: "old request"}, - {Role: "assistant", Content: "old answer"}, - {Role: "user", Content: "recent request"}, - }, - Think: thinkHigh, - Force: true, - }) - if err != nil { - t.Fatal(err) - } - if !result.Compacted || result.Summary != "unset think summary" { - t.Fatalf("compaction result = %#v", result) - } - if len(client.requests) != 3 { - t.Fatalf("summary requests = %d, want 3", len(client.requests)) - } - if client.requests[0].Think != thinkHigh { - t.Fatalf("first summary request think = %#v, want original", client.requests[0].Think) - } - if client.requests[1].Think == nil || client.requests[1].Think.Value != false { - t.Fatalf("fallback summary request think = %#v, want false", client.requests[1].Think) - } - if client.requests[2].Think != nil { - t.Fatalf("unsupported fallback retry think = %#v, want nil", client.requests[2].Think) - } -} - -func TestSimpleCompactorKeepsFewerTurnsForShortChats(t *testing.T) { - client := &fakeClient{ - responses: [][]api.ChatResponse{{ - {Message: api.Message{Role: "assistant", Content: "short summary"}}, - }}, - } - compactor := &SimpleCompactor{Client: client, Options: CompactionOptions{ - ContextWindowTokens: 16000, - KeepUserTurns: 3, - Threshold: 0.5, - }} - - result, err := compactor.MaybeCompact(context.Background(), CompactionRequest{ - ChatID: "chat-1", - Model: "model", - Messages: []api.Message{ - {Role: "user", Content: "old request"}, - {Role: "assistant", Content: "old answer"}, - {Role: "user", Content: "latest request"}, - }, - Latest: api.ChatResponse{Metrics: api.Metrics{PromptEvalCount: 12000}}, - }) - if err != nil { - t.Fatal(err) - } - if !result.Compacted { - t.Fatal("expected compaction") - } - if len(result.Messages) != 3 { - t.Fatalf("messages = %#v, want compaction tool pair plus latest request", result.Messages) - } - assertCompactionSummaryPair(t, result.Messages[:2]) - if result.Messages[2].Content != "latest request" { - t.Fatalf("latest turn was not kept: %#v", result.Messages) - } -} - -func TestSimpleCompactorCanArchiveWholeShortChat(t *testing.T) { - client := &fakeClient{ - responses: [][]api.ChatResponse{{ - {Message: api.Message{Role: "assistant", Content: "whole summary"}}, - }}, - } - compactor := &SimpleCompactor{Client: client, Options: CompactionOptions{ - ContextWindowTokens: 100, - KeepUserTurns: 3, - Threshold: 0.5, - }} - - result, err := compactor.MaybeCompact(context.Background(), CompactionRequest{ - ChatID: "chat-1", - Model: "model", - Messages: []api.Message{ - {Role: "user", Content: "only request"}, - {Role: "assistant", Content: "only answer"}, - }, - Latest: api.ChatResponse{Metrics: api.Metrics{PromptEvalCount: 75}}, - }) - if err != nil { - t.Fatal(err) - } - if !result.Compacted { - t.Fatal("expected compaction") - } - if len(result.Messages) != 2 { - t.Fatalf("messages = %#v, want only compaction tool pair", result.Messages) - } - assertCompactionSummaryPair(t, result.Messages) -} - -func TestSimpleCompactorSkipsBelowThreshold(t *testing.T) { - client := &fakeClient{} - compactor := &SimpleCompactor{Client: client, Options: CompactionOptions{ - ContextWindowTokens: 100, - Threshold: 0.8, - }} - - messages := []api.Message{ - {Role: "user", Content: "one"}, - {Role: "user", Content: "two"}, - {Role: "user", Content: "three"}, - {Role: "user", Content: "four"}, - {Role: "user", Content: "five"}, - } - result, err := compactor.MaybeCompact(context.Background(), CompactionRequest{ - Model: "model", - Messages: messages, - Latest: api.ChatResponse{Metrics: api.Metrics{PromptEvalCount: 50}}, - }) - if err != nil { - t.Fatal(err) - } - if result.Compacted { - t.Fatal("did not expect compaction") - } - if result.Due { - t.Fatal("below-threshold compaction should not be due") - } - if len(result.Messages) != len(messages) { - t.Fatalf("messages changed below threshold: %#v", result.Messages) - } - if len(client.requests) != 0 { - t.Fatalf("summary requests = %d, want 0", len(client.requests)) - } -} - -func TestSimpleCompactorUsesEstimatedMessagesWhenPromptEvalMissing(t *testing.T) { - client := &fakeClient{ - responses: [][]api.ChatResponse{{ - {Message: api.Message{Role: "assistant", Content: "estimated summary"}}, - }}, - } - compactor := &SimpleCompactor{Client: client, Options: CompactionOptions{ - ContextWindowTokens: 100, - KeepUserTurns: 1, - Threshold: 0.8, - }} - - result, err := compactor.MaybeCompact(context.Background(), CompactionRequest{ - Model: "model", - Messages: []api.Message{ - {Role: "user", Content: "old request"}, - {Role: "assistant", Content: "old answer"}, - {Role: "user", Content: "read large output"}, - {Role: "assistant", ToolCalls: []api.ToolCall{{ - ID: "call-1", - Function: api.ToolCallFunction{ - Name: "read", - }, - }}}, - {Role: "tool", ToolName: "read", ToolCallID: "call-1", Content: strings.Repeat("x", 360)}, - }, - }) - if err != nil { - t.Fatal(err) - } - if !result.Due || !result.Compacted { - t.Fatalf("expected estimate-driven compaction, got %#v", result) - } - if result.Summary != "estimated summary" { - t.Fatalf("summary = %q", result.Summary) - } -} - -func TestSimpleCompactorEstimateIncludesRequestPreamble(t *testing.T) { - compactor := &SimpleCompactor{Client: nil, Options: CompactionOptions{ - ContextWindowTokens: 100, - Threshold: 0.8, - }} - - if !compactor.shouldCompact(CompactionRequest{ - SystemPrompt: strings.Repeat("system ", 360), - Messages: []api.Message{{Role: "user", Content: "tiny"}}, - }) { - t.Fatal("system prompt should count toward compaction estimate") - } - - if !compactor.shouldCompact(CompactionRequest{ - Messages: []api.Message{{Role: "user", Content: "tiny"}}, - Tools: api.Tools{{ - Type: "function", - Function: api.ToolFunction{ - Name: "verbose_tool", - Description: strings.Repeat("description ", 360), - }, - }}, - }) { - t.Fatal("tool definitions should count toward compaction estimate") - } -} - -func TestCompactionPromptFitsBudgetByTruncatingLargeToolOutput(t *testing.T) { - largeToolOutput := strings.Repeat("x", 10_000) - body, err := compactionPrompt("", []api.Message{ - {Role: "user", Content: "what changed?"}, - {Role: "assistant", ToolCalls: []api.ToolCall{{ - ID: "call-1", - Function: api.ToolCallFunction{ - Name: "bash", - }, - }}}, - {Role: "tool", ToolName: "bash", ToolCallID: "call-1", Content: largeToolOutput}, - }, 300) - if err != nil { - t.Fatal(err) - } - if estimateCompactionTokens(body) > 300 { - t.Fatalf("compaction prompt tokens = %d, want <= 300", estimateCompactionTokens(body)) - } - if strings.Count(body, "x") >= len(largeToolOutput) { - t.Fatal("large tool output was not truncated") - } - if !strings.Contains(body, "[tool output truncated: showing first ~") { - t.Fatalf("truncation marker missing from compaction prompt: %q", body) - } -} - -func TestCompactionPromptRetruncatesAlreadyTruncatedToolOutput(t *testing.T) { - alreadyTruncated := strings.Repeat("x", 7000) + "\n\n[tool output truncated: showing first ~100 tokens and last ~100 tokens; omitted ~99999 tokens. Use a narrower command, line range, or search query if more detail is needed.]\n\n" + strings.Repeat("y", 7000) - body, err := compactionPrompt("", []api.Message{ - {Role: "user", Content: "what changed?"}, - {Role: "assistant", ToolCalls: []api.ToolCall{{ - ID: "call-1", - Function: api.ToolCallFunction{ - Name: "bash", - }, - }}}, - {Role: "tool", ToolName: "bash", ToolCallID: "call-1", Content: alreadyTruncated}, - }, 300) - if err != nil { - t.Fatal(err) - } - if estimateCompactionTokens(body) > 300 { - t.Fatalf("compaction prompt tokens = %d, want <= 300", estimateCompactionTokens(body)) - } - if strings.Count(body, "x")+strings.Count(body, "y") >= 14_000 { - t.Fatal("already-truncated tool output was not truncated again") - } - if !strings.Contains(body, "[tool output truncated: showing first ~") { - t.Fatalf("truncation marker missing from compaction prompt: %q", body) - } -} - -func TestCompactionSummaryTextStripsPrefix(t *testing.T) { - content := compactionSummaryMessageForTask("worked on branch changes", false) - if got := CompactionSummaryText(content); got != "worked on branch changes" { - t.Fatalf("summary text = %q", got) - } -} - -func TestCompactionSummaryCanTellModelToContinueTask(t *testing.T) { - content := compactionSummaryMessageForTask("worked on branch changes", true) - if !strings.Contains(content, CompactionContinueInstruction) { - t.Fatalf("summary message missing continue instruction: %q", content) - } - if got := CompactionSummaryText(content); got != "worked on branch changes" { - t.Fatalf("summary text = %q", got) - } -} - -func TestResolveContextWindowTokensPrefersExplicitNumCtx(t *testing.T) { - tests := []struct { - name string - options map[string]any - configured int - want int - }{ - { - name: "explicit smaller num ctx", - options: map[string]any{"num_ctx": 4096}, - configured: 8192, - want: 4096, - }, - { - name: "explicit num ctx can exceed configured metadata", - options: map[string]any{"num_ctx": 131072}, - configured: 8192, - want: 131072, - }, - { - name: "metadata without explicit num ctx", - configured: 32768, - want: 32768, - }, - } - - for _, tt := range tests { - t.Run(tt.name, func(t *testing.T) { - if got := ResolveContextWindowTokens(tt.options, tt.configured); got != tt.want { - t.Fatalf("ResolveContextWindowTokens() = %d, want %d", got, tt.want) - } - }) - } -} - -func TestSimpleCompactorForceCompactsWithoutPromptEvalCount(t *testing.T) { - client := &fakeClient{ - responses: [][]api.ChatResponse{{ - {Message: api.Message{Role: "assistant", Content: "forced summary"}}, - }}, - } - compactor := &SimpleCompactor{Client: client, Options: CompactionOptions{ - ContextWindowTokens: 100, - KeepUserTurns: 1, - Threshold: 0.8, - }} - - result, err := compactor.MaybeCompact(context.Background(), CompactionRequest{ - Model: "model", - Messages: []api.Message{ - {Role: "user", Content: "old"}, - {Role: "assistant", Content: "old answer"}, - {Role: "user", Content: "recent"}, - }, - Force: true, - }) - if err != nil { - t.Fatal(err) - } - if !result.Due || !result.Compacted { - t.Fatalf("forced compaction result = %#v", result) - } - if result.Summary != "forced summary" { - t.Fatalf("summary = %q", result.Summary) - } -} - -func TestSimpleCompactorDefaultsToKeepingThreeUserTurns(t *testing.T) { - client := &fakeClient{ - responses: [][]api.ChatResponse{{ - {Message: api.Message{Role: "assistant", Content: "summary"}}, - }}, - } - compactor := &SimpleCompactor{Client: client, Options: CompactionOptions{ - ContextWindowTokens: 16000, - Threshold: 0.5, - }} - - result, err := compactor.MaybeCompact(context.Background(), CompactionRequest{ - ChatID: "chat-1", - Model: "model", - Messages: []api.Message{ - {Role: "user", Content: "old"}, - {Role: "assistant", Content: "old answer"}, - {Role: "user", Content: "one"}, - {Role: "assistant", Content: "one answer"}, - {Role: "user", Content: "two"}, - {Role: "assistant", Content: "two answer"}, - {Role: "user", Content: "three"}, - }, - Latest: api.ChatResponse{Metrics: api.Metrics{PromptEvalCount: 12000}}, - }) - if err != nil { - t.Fatal(err) - } - if !result.Compacted { - t.Fatal("expected compaction") - } - assertCompactionSummaryPair(t, result.Messages[:2]) - if got := result.Messages[2].Content; got != "one" { - t.Fatalf("first kept turn = %q, want one", got) - } -} - -func TestSimpleCompactorCarriesPreviousSummary(t *testing.T) { - client := &fakeClient{ - responses: [][]api.ChatResponse{{ - {Message: api.Message{Role: "assistant", Content: "new summary"}}, - }}, - } - compactor := &SimpleCompactor{Client: client, Options: CompactionOptions{ - ContextWindowTokens: 16000, - KeepUserTurns: 1, - Threshold: 0.5, - }} - - result, err := compactor.MaybeCompact(context.Background(), CompactionRequest{ - Model: "model", - Messages: []api.Message{ - {Role: "system", Content: CompactionSummaryMessagePrefix + "old summary"}, - {Role: "user", Content: "old"}, - {Role: "assistant", Content: "old answer"}, - {Role: "user", Content: "recent"}, - }, - Latest: api.ChatResponse{Metrics: api.Metrics{PromptEvalCount: 12000}}, - }) - if err != nil { - t.Fatal(err) - } - if !result.Compacted { - t.Fatal("expected compaction") - } - if !strings.Contains(client.requests[0].Messages[1].Content, "Previous summary:\nold summary") { - t.Fatalf("previous summary missing from request: %q", client.requests[0].Messages[1].Content) - } -} - -func TestSimpleCompactorCarriesPreviousToolSummaryAndPlacesNewSummaryBeforeKeptSuffix(t *testing.T) { - client := &fakeClient{ - responses: [][]api.ChatResponse{{ - {Message: api.Message{Role: "assistant", Content: "new summary"}}, - }}, - } - compactor := &SimpleCompactor{Client: client, Options: CompactionOptions{ - ContextWindowTokens: 16000, - KeepUserTurns: 1, - Threshold: 0.5, - }} - - messages := []api.Message{ - {Role: "user", Content: "kept before old summary"}, - CompactionSummaryMessages("old summary", false)[0], - CompactionSummaryMessages("old summary", false)[1], - {Role: "user", Content: "latest request"}, - } - result, err := compactor.MaybeCompact(context.Background(), CompactionRequest{ - Model: "model", - Messages: messages, - Latest: api.ChatResponse{Metrics: api.Metrics{PromptEvalCount: 12000}}, - }) - if err != nil { - t.Fatal(err) - } - if !result.Compacted { - t.Fatal("expected compaction") - } - if !strings.Contains(client.requests[0].Messages[1].Content, "Previous summary:\nold summary") { - t.Fatalf("previous summary missing from request: %q", client.requests[0].Messages[1].Content) - } - if len(result.Messages) != 3 { - t.Fatalf("messages = %#v, want compaction pair plus latest request", result.Messages) - } - assertCompactionSummaryPair(t, result.Messages[:2]) - if result.Messages[2].Content != "latest request" { - t.Fatalf("kept suffix = %#v", result.Messages) - } -} diff --git a/agent/events.go b/agent/events.go deleted file mode 100644 index 6d2c4a881eb..00000000000 --- a/agent/events.go +++ /dev/null @@ -1,177 +0,0 @@ -package agent - -import ( - "context" - "errors" - - "github.com/ollama/ollama/api" -) - -type EventType string - -const ( - EventMessageDelta EventType = "message_delta" - EventThinkingDelta EventType = "thinking_delta" - EventToolCallDetected EventType = "tool_call_detected" - EventToolStarted EventType = "tool_started" - EventToolFinished EventType = "tool_finished" - EventCompactionStarted EventType = "compaction_started" - EventCompactionProgress EventType = "compaction_progress" - EventCompacted EventType = "compacted" - EventCompactionSkipped EventType = "compaction_skipped" - EventRunFinished EventType = "run_finished" - EventError EventType = "error" -) - -// ToolStatus is the typed lifecycle state for a tool call, carried on -// Event.ToolStatus for tool events. -type ToolStatus string - -const ( - ToolStatusRunning ToolStatus = "running" - ToolStatusDone ToolStatus = "done" - ToolStatusFailed ToolStatus = "failed" - ToolStatusDenied ToolStatus = "denied" - ToolStatusDisabled ToolStatus = "disabled" - ToolStatusSkipped ToolStatus = "skipped" -) - -// RunStatus is the typed terminal outcome of a run, carried on Event.Status for -// run_finished events. -type RunStatus string - -const ( - RunStatusDone RunStatus = "done" - RunStatusDenied RunStatus = "denied" - RunStatusCanceled RunStatus = "canceled" -) - -// CompactionTrigger is the typed reason a compaction ran or was attempted, -// carried on Event.CompactionTrigger for compaction events. -type CompactionTrigger string - -const ( - CompactionTriggerForce CompactionTrigger = "force" - CompactionTriggerPromptEval CompactionTrigger = "prompt_eval" - CompactionTriggerEstimate CompactionTrigger = "estimate" - CompactionTriggerToolOutput CompactionTrigger = "tool_output" - CompactionTriggerError CompactionTrigger = "error" - CompactionTriggerDue CompactionTrigger = "due" -) - -type Event struct { - Type EventType `json:"type"` - RunID string `json:"runId,omitempty"` - ChatID string `json:"chatId,omitempty"` - Model string `json:"model,omitempty"` - Status RunStatus `json:"status,omitempty"` - ToolStatus ToolStatus `json:"toolStatus,omitempty"` - CompactionTrigger CompactionTrigger `json:"compactionTrigger,omitempty"` - ToolCallID string `json:"toolCallId,omitempty"` - ToolName string `json:"toolName,omitempty"` - WorkingDir string `json:"workingDir,omitempty"` - Content string `json:"content,omitempty"` - Thinking string `json:"thinking,omitempty"` - ToolCalls []api.ToolCall `json:"toolCalls,omitempty"` - Messages []api.Message `json:"messages,omitempty"` - Args map[string]any `json:"args,omitempty"` - Tokens int `json:"tokens,omitempty"` - Error string `json:"error,omitempty"` -} - -type EventSink interface { - Emit(Event) error -} - -type EventSinkFunc func(Event) error - -func (fn EventSinkFunc) Emit(event Event) error { - if fn == nil { - return nil - } - return fn(event) -} - -// eventMetadata carries the run identification fields shared by all events. -type eventMetadata struct { - runID string - chatID string - model string -} - -func newEventMetadata(runID string, opts RunOptions) eventMetadata { - return eventMetadata{runID: runID, chatID: opts.ChatID, model: opts.Model} -} - -func newMessageDelta(m eventMetadata, content string) Event { - return Event{Type: EventMessageDelta, RunID: m.runID, ChatID: m.chatID, Model: m.model, Content: content} -} - -func newThinkingDelta(m eventMetadata, thinking string) Event { - return Event{Type: EventThinkingDelta, RunID: m.runID, ChatID: m.chatID, Model: m.model, Thinking: thinking} -} - -func newToolCallDetected(m eventMetadata, calls []api.ToolCall) Event { - return Event{Type: EventToolCallDetected, RunID: m.runID, ChatID: m.chatID, Model: m.model, ToolCalls: calls} -} - -func newToolStarted(m eventMetadata, callID, toolName, workingDir string, args map[string]any) Event { - return Event{Type: EventToolStarted, RunID: m.runID, ChatID: m.chatID, Model: m.model, ToolStatus: ToolStatusRunning, ToolCallID: callID, ToolName: toolName, WorkingDir: workingDir, Args: args} -} - -func newToolFinished(m eventMetadata, status ToolStatus, callID, toolName, workingDir string, args map[string]any, content, errMsg string) Event { - ev := Event{Type: EventToolFinished, RunID: m.runID, ChatID: m.chatID, Model: m.model, ToolStatus: status, ToolCallID: callID, ToolName: toolName, WorkingDir: workingDir, Args: args, Content: content} - if errMsg != "" { - ev.Error = errMsg - } - return ev -} - -func newRunFinished(m eventMetadata, status RunStatus) Event { - return Event{Type: EventRunFinished, RunID: m.runID, ChatID: m.chatID, Model: m.model, Status: status} -} - -func newErrorEvent(m eventMetadata, errMsg string) Event { - return Event{Type: EventError, RunID: m.runID, ChatID: m.chatID, Model: m.model, Error: errMsg} -} - -func newCompactionProgress(m eventMetadata, tokens int) Event { - return Event{Type: EventCompactionProgress, RunID: m.runID, ChatID: m.chatID, Model: m.model, Tokens: tokens} -} - -func newCompactionStarted(m eventMetadata, trigger CompactionTrigger) Event { - return Event{Type: EventCompactionStarted, RunID: m.runID, ChatID: m.chatID, Model: m.model, CompactionTrigger: trigger} -} - -func newCompactionSkipped(m eventMetadata, trigger CompactionTrigger, content string) Event { - return Event{Type: EventCompactionSkipped, RunID: m.runID, ChatID: m.chatID, Model: m.model, CompactionTrigger: trigger, Content: content} -} - -func newCompacted(m eventMetadata, messages []api.Message, trigger CompactionTrigger, content string) Event { - return Event{Type: EventCompacted, RunID: m.runID, ChatID: m.chatID, Model: m.model, CompactionTrigger: trigger, Content: content, Messages: messages} -} - -func (s *Session) emit(event Event) error { - if s == nil { - return nil - } - var errs []error - for _, sink := range s.EventSinks { - if sink == nil { - continue - } - if err := sink.Emit(event); err != nil { - errs = append(errs, err) - } - } - return errors.Join(errs...) -} - -func (s *Session) emitIgnoringCanceled(ctx context.Context, event Event) error { - err := s.emit(event) - if err != nil && ctx != nil && ctx.Err() != nil { - //nolint:nilerr // Event sinks may close during cancellation; cancellation is not a user-facing emit failure. - return nil - } - return err -} diff --git a/agent/registry.go b/agent/registry.go deleted file mode 100644 index 377efa8141e..00000000000 --- a/agent/registry.go +++ /dev/null @@ -1,104 +0,0 @@ -package agent - -import ( - "context" - "fmt" - "sort" - - "github.com/ollama/ollama/api" -) - -type ToolContext struct { - WorkingDir string -} - -type ToolResult struct { - Content string - WorkingDir string -} - -type Tool interface { - Name() string - Description() string - Schema() api.ToolFunction - Execute(context.Context, ToolContext, map[string]any) (ToolResult, error) -} - -type ApprovalRequired interface { - RequiresApproval(map[string]any) bool -} - -// ScopedTool is implemented by tools that need per-invocation approval -// scoping beyond the tool name (e.g. shell commands scoped to the exact -// command string). Tools that don't implement this are scoped by name only. -type ScopedTool interface { - ApprovalScope(args map[string]any) string -} - -type Registry struct { - tools map[string]Tool -} - -func (r *Registry) Register(tool Tool) { - if r == nil || tool == nil { - return - } - if r.tools == nil { - r.tools = make(map[string]Tool) - } - r.tools[tool.Name()] = tool -} - -func (r *Registry) Get(name string) (Tool, bool) { - if r == nil { - return nil, false - } - tool, ok := r.tools[name] - return tool, ok -} - -func (r *Registry) Names() []string { - if r == nil { - return nil - } - names := make([]string, 0, len(r.tools)) - for name := range r.tools { - names = append(names, name) - } - sort.Strings(names) - return names -} - -func (r *Registry) Tools() api.Tools { - if r == nil { - return nil - } - names := r.Names() - apiTools := make(api.Tools, 0, len(names)) - for _, name := range names { - tool := r.tools[name] - apiTools = append(apiTools, api.Tool{ - Type: "function", - Function: tool.Schema(), - }) - } - return apiTools -} - -func (r *Registry) Execute(ctx context.Context, toolCtx ToolContext, call api.ToolCall) (ToolResult, error) { - tool, ok := r.Get(call.Function.Name) - if !ok { - return ToolResult{}, fmt.Errorf("unknown tool: %s", call.Function.Name) - } - return tool.Execute(ctx, toolCtx, call.Function.Arguments.ToMap()) -} - -func ToolRequiresApproval(tool Tool, args map[string]any) bool { - if tool == nil { - return false - } - if t, ok := tool.(ApprovalRequired); ok { - return t.RequiresApproval(args) - } - return false -} diff --git a/agent/session.go b/agent/session.go deleted file mode 100644 index 46969c8f931..00000000000 --- a/agent/session.go +++ /dev/null @@ -1,1092 +0,0 @@ -package agent - -import ( - "context" - "encoding/json" - "errors" - "fmt" - "os" - "path/filepath" - "strings" - - "github.com/google/uuid" - - "github.com/ollama/ollama/api" - "github.com/ollama/ollama/internal/modelref" -) - -type ChatClient interface { - Chat(context.Context, *api.ChatRequest, api.ChatResponseFunc) error -} - -type Session struct { - Client ChatClient - EventSinks []EventSink - Tools *Registry - Skills *SkillCatalog - DisableTools bool - ApprovalPrompter ApprovalPrompter - ApprovalState *ApprovalState - WorkingDir string - Compactor Compactor -} - -type RunOptions struct { - ChatID string - Model string - SystemPrompt string - Messages []api.Message - NewMessages []api.Message - Format string - Options map[string]any - Think *api.ThinkValue - KeepAlive *api.Duration - // SkillName loads a catalog skill as an ordered synthetic tool call/result - // before the first model request for this run. - SkillName string - // MaxToolRounds limits consecutive model/tool cycles. A positive value is - // an explicit limit. Zero selects the model-specific default: local models - // use the default guard and cloud models are unlimited. A negative value - // disables the guard for tests or special callers. - MaxToolRounds int -} - -type RunResult struct { - Messages []api.Message - Latest api.ChatResponse - WorkingDir string -} - -const ( - defaultMaxToolRounds = 100 - maxToolResultRunes = 60000 - smallContextToolResultRunes = 6000 - tinyContextToolResultRunes = 3200 - smallContextToolResultTokenWindow = 8192 - tinyContextToolResultTokenWindow = 4096 - toolTruncationMarkerReserveTokens = 64 - toolOutputFullOmissionPrefix = "[tool output truncated: output omitted because the context is full;" -) - -type toolOutputOverflow struct { - toolName string - toolCallID string - content string -} - -type toolBatchResult struct { - messages []api.Message - stop toolExecutionStop - overflows []toolOutputOverflow -} - -// toolExecutionStop is the batch-level outcome for a group of tool calls, -// distinct from per-call Event.Status values. The values overlap with -// runFinish.status ("denied", "canceled") because a denied or canceled -// batch also terminates the run with the matching status. -type toolExecutionStop string - -const ( - toolExecutionDenied toolExecutionStop = "denied" - toolExecutionCanceled toolExecutionStop = "canceled" -) - -const toolExecutionDisabledMessage = "Tool execution disabled." - -type runPhase int - -const ( - runPhaseModel runPhase = iota - runPhaseTools - runPhaseCompact - runPhaseDone -) - -type runState struct { - runID string - opts RunOptions - - phase runPhase - - messages []api.Message - latest api.ChatResponse - - assistant api.Message - pendingToolCalls []api.ToolCall - canceled bool - - toolBatch *toolBatchResult - - consecutiveModelErrors int - toolRounds int - maxToolRounds int - compactionSkipNotified bool - - finish runFinish -} - -type runFinish struct { - status RunStatus - ignoreCanceled bool - err error -} - -func (st *runState) finishDone() { - st.finish = runFinish{status: RunStatusDone} - st.phase = runPhaseDone -} - -func (st *runState) finishDenied() { - st.finish = runFinish{status: RunStatusDenied} - st.phase = runPhaseDone -} - -func (st *runState) finishCanceled() { - st.finish = runFinish{status: RunStatusCanceled, ignoreCanceled: true} - st.phase = runPhaseDone -} - -func (st *runState) finishError(err error) { - st.finish = runFinish{err: err} - st.phase = runPhaseDone -} - -func (s *Session) Run(ctx context.Context, opts RunOptions) (*RunResult, error) { - if err := s.validateRun(opts); err != nil { - return nil, err - } - if s.ApprovalState == nil { - s.ApprovalState = &ApprovalState{} - } - runID := uuid.NewString() - messages, err := s.buildRunMessages(ctx, runID, opts) - if err != nil { - return nil, err - } - activatedSkill, err := s.activateSkill(ctx, runID, opts) - if err != nil { - s.emit(newErrorEvent(newEventMetadata(runID, opts), err.Error())) - return nil, err - } - if len(activatedSkill) > 0 { - messages = append(messages, activatedSkill...) - if err := s.checkPreflightPromptBudget(opts, messages); err != nil { - s.emit(newErrorEvent(newEventMetadata(runID, opts), err.Error())) - return nil, err - } - } - - st := runState{ - runID: runID, - opts: opts, - phase: runPhaseModel, - messages: messages, - maxToolRounds: resolvedMaxToolRounds(opts.Model, opts.MaxToolRounds), - } - for { - switch st.phase { - case runPhaseModel: - if err := s.runModelStep(ctx, &st); err != nil { - return nil, err - } - case runPhaseTools: - if err := s.runToolStep(ctx, &st); err != nil { - return nil, err - } - case runPhaseCompact: - if err := s.runCompactionStep(ctx, &st); err != nil { - return nil, err - } - case runPhaseDone: - return s.finishRun(ctx, &st) - } - } -} - -// validateRun checks the preconditions for a run. -func (s *Session) validateRun(opts RunOptions) error { - if s == nil { - return errors.New("nil session") - } - if s.Client == nil { - return errors.New("agent session requires a chat client") - } - if opts.Model == "" { - return errors.New("agent session requires a model") - } - return nil -} - -// buildRunMessages sanitizes the provided message history, runs the preflight -// prompt-budget check, and returns the initial message list for the run. It -// emits an EventError and returns it if the preflight check fails. -func (s *Session) buildRunMessages(ctx context.Context, runID string, opts RunOptions) ([]api.Message, error) { - messages := make([]api.Message, 0, len(opts.Messages)+len(opts.NewMessages)) - for _, msg := range opts.Messages { - messages = append(messages, sanitizeMessageForRun(msg)) - } - for _, msg := range opts.NewMessages { - msg = sanitizeMessageForRun(msg) - messages = append(messages, msg) - } - - if err := s.checkPreflightPromptBudget(opts, messages); err != nil { - s.emit(newErrorEvent(newEventMetadata(runID, opts), err.Error())) - return nil, err - } - return messages, nil -} - -func (s *Session) runModelStep(ctx context.Context, st *runState) error { - opts := st.opts - meta := newEventMetadata(st.runID, opts) - - assistant, pendingToolCalls, canceled, err := s.chatRound(ctx, st.runID, opts, st.messages, &st.latest) - if err != nil { - var statusErr api.StatusError - if errors.As(err, &statusErr) && statusErr.StatusCode >= 500 && st.consecutiveModelErrors < 2 { - st.consecutiveModelErrors++ - st.messages = append(st.messages, api.Message{ - Role: "user", - Content: fmt.Sprintf("Your previous response caused an error: %s\n\nPlease try again with a valid response.", statusErr.ErrorMessage), - }) - return nil - } - s.emit(newErrorEvent(meta, err.Error())) - return err - } - st.consecutiveModelErrors = 0 - st.assistant = assistant - st.pendingToolCalls = pendingToolCalls - st.canceled = canceled - - if !messageEmpty(assistant) { - st.messages = append(st.messages, assistant) - } - - if len(pendingToolCalls) == 0 { - st.toolBatch = nil - st.phase = runPhaseCompact - return nil - } - - if canceled { - skipped, skipErr := s.skipToolCalls(ctx, st.runID, opts, pendingToolCalls, "Tool execution skipped because the run was canceled.") - if skipErr != nil { - s.emit(newErrorEvent(meta, skipErr.Error())) - return skipErr - } - st.messages = append(st.messages, skipped...) - st.finishCanceled() - return nil - } - - if s.DisableTools { - batch, skipErr := s.disabledToolCalls(ctx, st.runID, opts, st.messages, pendingToolCalls) - if skipErr != nil { - s.emit(newErrorEvent(meta, skipErr.Error())) - return skipErr - } - st.messages = append(st.messages, batch.messages...) - st.toolBatch = &batch - st.phase = runPhaseCompact - return nil - } - - if s.Tools == nil { - st.finishDone() - return nil - } - - if st.maxToolRounds >= 0 && st.toolRounds >= st.maxToolRounds { - content := fmt.Sprintf("Tool execution skipped because the max tool-round limit of %d was reached. Send another message to continue.", st.maxToolRounds) - toolMessages, skipErr := s.skipToolCalls(ctx, st.runID, opts, pendingToolCalls, content) - if skipErr != nil { - s.emit(newErrorEvent(meta, skipErr.Error())) - return skipErr - } - st.messages = append(st.messages, toolMessages...) - err := fmt.Errorf("tool round limit reached after %d rounds; send another message to continue", st.maxToolRounds) - s.emit(newErrorEvent(meta, err.Error())) - st.finishError(err) - return nil - } - - st.phase = runPhaseTools - return nil -} - -func (s *Session) runToolStep(ctx context.Context, st *runState) error { - batch, err := s.executeToolCalls(ctx, st.runID, st.opts, st.messages, st.pendingToolCalls) - if err != nil { - s.emit(newErrorEvent(newEventMetadata(st.runID, st.opts), err.Error())) - return err - } - - st.messages = append(st.messages, batch.messages...) - st.toolBatch = &batch - st.phase = runPhaseCompact - return nil -} - -func (s *Session) runCompactionStep(ctx context.Context, st *runState) error { - opts := st.opts - meta := newEventMetadata(st.runID, opts) - var err error - if st.toolBatch != nil && len(st.toolBatch.overflows) > 0 { - st.messages, st.compactionSkipNotified, err = s.compactForToolOutputOverflow(ctx, st.runID, opts, st.messages, st.latest, st.assistant, st.toolBatch.messages, st.toolBatch.overflows, st.compactionSkipNotified) - } else { - st.messages, st.compactionSkipNotified, err = s.maybeCompact(ctx, st.runID, opts, st.messages, st.latest, st.compactionSkipNotified) - } - if err != nil { - s.emit(newErrorEvent(meta, err.Error())) - st.finishError(err) - return nil - } - - if st.toolBatch == nil { - if st.canceled { - st.finishCanceled() - } else { - st.finishDone() - } - return nil - } - - switch st.toolBatch.stop { - case toolExecutionDenied: - st.finishDenied() - case toolExecutionCanceled: - st.finishCanceled() - default: - st.toolRounds++ - st.assistant = api.Message{} - st.pendingToolCalls = nil - st.toolBatch = nil - st.phase = runPhaseModel - } - return nil -} - -func (s *Session) finishRun(ctx context.Context, st *runState) (*RunResult, error) { - if st.finish.status != "" { - event := newRunFinished(newEventMetadata(st.runID, st.opts), st.finish.status) - var err error - if st.finish.ignoreCanceled { - err = s.emitIgnoringCanceled(ctx, event) - } else { - err = s.emit(event) - } - if err != nil { - return nil, err - } - } - return &RunResult{Messages: st.messages, Latest: st.latest, WorkingDir: s.WorkingDir}, st.finish.err -} - -func (s *Session) chatRound(ctx context.Context, runID string, opts RunOptions, messages []api.Message, latest *api.ChatResponse) (api.Message, []api.ToolCall, bool, error) { - meta := newEventMetadata(runID, opts) - var tools api.Tools - if !s.DisableTools { - tools = s.availableTools() - } - req := buildChatRequest(opts, messages, tools) - - assistant := api.Message{Role: "assistant"} - var pendingToolCalls []api.ToolCall - - err := s.Client.Chat(ctx, &req, func(response api.ChatResponse) error { - if response.Message.Role != "" { - assistant.Role = response.Message.Role - } - - if messageEmpty(response.Message) { - *latest = response - return nil - } - - if response.Message.Thinking != "" { - assistant.Thinking += response.Message.Thinking - if err := s.emit(newThinkingDelta(meta, response.Message.Thinking)); err != nil { - return err - } - } - - if response.Message.Content != "" { - assistant.Content += response.Message.Content - if err := s.emit(newMessageDelta(meta, response.Message.Content)); err != nil { - return err - } - } - - if len(response.Message.ToolCalls) > 0 { - assistant.ToolCalls = append(assistant.ToolCalls, response.Message.ToolCalls...) - pendingToolCalls = append(pendingToolCalls, response.Message.ToolCalls...) - if err := s.emit(newToolCallDetected(meta, response.Message.ToolCalls)); err != nil { - return err - } - } - - *latest = response - return nil - }) - if err != nil { - if isContextCanceledError(ctx, err) { - return assistant, pendingToolCalls, true, nil - } - return assistant, pendingToolCalls, false, err - } - - return assistant, pendingToolCalls, false, nil -} - -func buildChatRequest(opts RunOptions, messages []api.Message, tools api.Tools) api.ChatRequest { - requestMessages := sanitizeMessagesForRequest(messages) - if strings.TrimSpace(opts.SystemPrompt) != "" { - withSystem := make([]api.Message, 0, len(requestMessages)+1) - withSystem = append(withSystem, api.Message{Role: "system", Content: opts.SystemPrompt}) - requestMessages = append(withSystem, requestMessages...) - } - - format := opts.Format - if format == "json" { - format = `"` + format + `"` - } - - req := api.ChatRequest{ - Model: opts.Model, - Messages: requestMessages, - Format: json.RawMessage(format), - Options: opts.Options, - Think: opts.Think, - } - if opts.KeepAlive != nil { - req.KeepAlive = opts.KeepAlive - } - if len(tools) > 0 { - req.Tools = tools - } - return req -} - -func (s *Session) executeToolCalls(ctx context.Context, runID string, opts RunOptions, messages []api.Message, calls []api.ToolCall) (toolBatchResult, error) { - meta := newEventMetadata(runID, opts) - batch := toolBatchResult{ - messages: make([]api.Message, 0, len(calls)), - } - // Pre-compute the full-history token estimate once per batch instead of - // re-marshaling the entire history for each tool call. Per-call deltas - // (tool messages already appended this batch) are tracked in batchTokens - // and added to historyTokens for a lightweight running total. - historyTokens := s.estimateRunPromptTokens(opts, messages) - batchTokens := 0 - - type plannedToolCall struct { - call api.ToolCall - tool Tool - toolName string - args map[string]any - workingDir string - } - plans := make([]plannedToolCall, 0, len(calls)) - batchWorkingDir := s.currentWorkingDir() - approvalReq := ApprovalRequest{WorkingDir: batchWorkingDir} - for _, call := range calls { - toolName := call.Function.Name - args := call.Function.Arguments.ToMap() - tool, ok := s.Tools.Get(toolName) - plans = append(plans, plannedToolCall{ - call: call, - tool: tool, - toolName: toolName, - args: args, - workingDir: batchWorkingDir, - }) - if ok && s.needsApproval(tool, toolName, args) { - approvalReq.AddToolCall(call.ID, toolName, toolApprovalScope(tool, toolName, args), args) - } - } - - if len(approvalReq.Calls) > 0 { - approvalResult, err := s.authorizeToolCalls(ctx, approvalReq) - if err != nil { - if ctx.Err() != nil { - skipped, skipErr := s.skipToolCalls(ctx, runID, opts, calls, "Tool execution skipped because the run was canceled.") - if skipErr != nil { - return toolBatchResult{}, skipErr - } - batch.messages = append(batch.messages, skipped...) - batch.stop = toolExecutionCanceled - return batch, nil - } - return toolBatchResult{}, err - } - if !approvalResult.Allow { - content := approvalResult.Reason - if content == "" { - content = "Tool execution denied." - } - for _, plan := range plans { - msg := s.toolMessageForContext(plan.toolName, plan.call.ID, content, opts, historyTokens+batchTokens) - batch.messages = append(batch.messages, msg) - batchTokens += estimateMessagesTokens([]api.Message{msg}) - deniedContent := msg.Content - if emitErr := s.emit(newToolFinished(meta, "denied", plan.call.ID, plan.toolName, "", plan.args, deniedContent, deniedContent)); emitErr != nil { - return toolBatchResult{}, emitErr - } - } - batch.stop = toolExecutionDenied - return batch, nil - } - } - - for i, plan := range plans { - call := plan.call - toolName := plan.toolName - args := plan.args - if ctx.Err() != nil { - skipped, skipErr := s.skipToolCalls(ctx, runID, opts, calls[i:], "Tool execution skipped because the run was canceled.") - if skipErr != nil { - return toolBatchResult{}, skipErr - } - batch.messages = append(batch.messages, skipped...) - batch.stop = toolExecutionCanceled - return batch, nil - } - if plan.tool == nil { - content := fmt.Sprintf("Error: unknown tool: %s", toolName) - msg := s.toolMessageForContext(toolName, call.ID, content, opts, historyTokens+batchTokens) - batch.messages = append(batch.messages, msg) - batchTokens += estimateMessagesTokens([]api.Message{msg}) - content = msg.Content - if toolOutputFullyOmitted(content) { - batch.overflows = append(batch.overflows, toolOutputOverflow{toolName: toolName, toolCallID: call.ID, content: fmt.Sprintf("Error: unknown tool: %s", toolName)}) - } - if emitErr := s.emit(newToolFinished(meta, "failed", call.ID, toolName, "", args, content, fmt.Sprintf("unknown tool: %s", toolName))); emitErr != nil { - return toolBatchResult{}, emitErr - } - continue - } - - if err := s.emit(newToolStarted(meta, call.ID, toolName, plan.workingDir, args)); err != nil { - return toolBatchResult{}, err - } - - result, err := s.Tools.Execute(ctx, ToolContext{WorkingDir: plan.workingDir}, call) - if err != nil { - rawContent := fmt.Sprintf("Error: %v", err) - msg := s.toolMessageForContext(toolName, call.ID, rawContent, opts, historyTokens+batchTokens) - batch.messages = append(batch.messages, msg) - batchTokens += estimateMessagesTokens([]api.Message{msg}) - content := msg.Content - if toolOutputFullyOmitted(content) { - batch.overflows = append(batch.overflows, toolOutputOverflow{toolName: toolName, toolCallID: call.ID, content: rawContent}) - } - if emitErr := s.emitIgnoringCanceled(ctx, newToolFinished(meta, "failed", call.ID, toolName, "", args, content, err.Error())); emitErr != nil { - return toolBatchResult{}, emitErr - } - if ctx.Err() != nil { - skipped, skipErr := s.skipToolCalls(ctx, runID, opts, calls[i+1:], "Tool execution skipped because the run was canceled.") - if skipErr != nil { - return toolBatchResult{}, skipErr - } - batch.messages = append(batch.messages, skipped...) - batch.stop = toolExecutionCanceled - return batch, nil - } - continue - } - - eventWorkingDir := plan.workingDir - if s.applyToolWorkingDir(result.WorkingDir) { - eventWorkingDir = s.WorkingDir - } - rawContent := result.Content - - msg := s.toolMessageForContext(toolName, call.ID, rawContent, opts, historyTokens+batchTokens) - batch.messages = append(batch.messages, msg) - batchTokens += estimateMessagesTokens([]api.Message{msg}) - content := msg.Content - - if toolOutputFullyOmitted(content) { - batch.overflows = append(batch.overflows, toolOutputOverflow{toolName: toolName, toolCallID: call.ID, content: rawContent}) - } - if err := s.emitIgnoringCanceled(ctx, newToolFinished(meta, "done", call.ID, toolName, eventWorkingDir, args, content, "")); err != nil { - return toolBatchResult{}, err - } - if ctx.Err() != nil { - skipped, skipErr := s.skipToolCalls(ctx, runID, opts, calls[i+1:], "Tool execution skipped because the run was canceled.") - if skipErr != nil { - return toolBatchResult{}, skipErr - } - batch.messages = append(batch.messages, skipped...) - batch.stop = toolExecutionCanceled - return batch, nil - } - } - return batch, nil -} - -func (s *Session) disabledToolCalls(ctx context.Context, runID string, opts RunOptions, messages []api.Message, calls []api.ToolCall) (toolBatchResult, error) { - meta := newEventMetadata(runID, opts) - batch := toolBatchResult{ - messages: make([]api.Message, 0, len(calls)), - } - historyTokens := s.estimateRunPromptTokens(opts, messages) - batchTokens := 0 - for _, call := range calls { - toolName := call.Function.Name - args := call.Function.Arguments.ToMap() - msg := s.toolMessageForContext(toolName, call.ID, toolExecutionDisabledMessage, opts, historyTokens+batchTokens) - batch.messages = append(batch.messages, msg) - batchTokens += estimateMessagesTokens([]api.Message{msg}) - if emitErr := s.emitIgnoringCanceled(ctx, newToolFinished(meta, "disabled", call.ID, toolName, "", args, msg.Content, msg.Content)); emitErr != nil { - return toolBatchResult{}, emitErr - } - } - return batch, nil -} - -func (s *Session) skipToolCalls(ctx context.Context, runID string, opts RunOptions, calls []api.ToolCall, content string) ([]api.Message, error) { - meta := newEventMetadata(runID, opts) - toolMessages := make([]api.Message, 0, len(calls)) - for _, call := range calls { - toolName := call.Function.Name - args := call.Function.Arguments.ToMap() - msg := toolMessage(toolName, call.ID, content) - toolMessages = append(toolMessages, msg) - if emitErr := s.emitIgnoringCanceled(ctx, newToolFinished(meta, "skipped", call.ID, toolName, "", args, msg.Content, msg.Content)); emitErr != nil { - return nil, emitErr - } - } - return toolMessages, nil -} - -func (s *Session) currentWorkingDir() string { - if s.WorkingDir != "" { - return s.WorkingDir - } - wd, err := os.Getwd() - if err != nil { - return "" - } - s.WorkingDir = wd - return s.WorkingDir -} - -func (s *Session) applyToolWorkingDir(next string) bool { - next = strings.TrimSpace(next) - if next == "" { - return false - } - current := s.currentWorkingDir() - nextAbs, err := canonicalSessionPath(next) - if err != nil { - return false - } - if current == nextAbs { - return false - } - s.WorkingDir = nextAbs - return true -} - -func canonicalSessionPath(path string) (string, error) { - abs, err := filepath.Abs(path) - if err != nil { - return "", err - } - resolved, err := filepath.EvalSymlinks(abs) - if err == nil { - return resolved, nil - } - return abs, nil -} - -func isContextCanceledError(ctx context.Context, err error) bool { - if err == nil { - return false - } - if errors.Is(err, context.Canceled) { - return true - } - return ctx != nil && errors.Is(ctx.Err(), context.Canceled) && strings.Contains(err.Error(), "context canceled") -} - -func (s *Session) maybeCompact(ctx context.Context, runID string, opts RunOptions, messages []api.Message, latest api.ChatResponse, skipNotified bool) ([]api.Message, bool, error) { - if s.Compactor == nil { - return messages, skipNotified, nil - } - req := s.compactionRequest(runID, opts, messages, latest) - trigger := s.autoCompactionTrigger(req) - if trigger != "" { - s.emitCompactionStarted(runID, opts, trigger) - } - result, err := s.Compactor.MaybeCompact(ctx, req) - if err != nil { - if result.Due && !skipNotified { - if trigger == "" { - trigger = CompactionTriggerError - } - s.emitCompactionSkipped(runID, opts, trigger, result.Reason) - skipNotified = true - } - return messages, skipNotified, nil - } - if !result.Compacted { - if result.Due && !skipNotified { - if trigger == "" { - trigger = CompactionTriggerDue - } - s.emitCompactionSkipped(runID, opts, trigger, result.Reason) - skipNotified = true - } - return messages, skipNotified, nil - } - s.emitCompacted(runID, opts, result.Messages, trigger, result.Summary) - if err := s.checkPostCompactionPromptBudget(opts, result.Messages); err != nil { - return result.Messages, skipNotified, err - } - return result.Messages, skipNotified, nil -} - -func (s *Session) compactForToolOutputOverflow(ctx context.Context, runID string, opts RunOptions, messages []api.Message, latest api.ChatResponse, assistant api.Message, toolMessages []api.Message, overflows []toolOutputOverflow, skipNotified bool) ([]api.Message, bool, error) { - if s.Compactor == nil { - return messages, skipNotified, nil - } - - keepUserTurns := 0 - req := s.compactionRequest(runID, opts, messages, latest) - req.Force = true - req.KeepUserTurns = &keepUserTurns - s.emitCompactionStarted(runID, opts, CompactionTriggerToolOutput) - - result, err := s.Compactor.MaybeCompact(ctx, req) - if err != nil { - if result.Due && !skipNotified { - s.emitCompactionSkipped(runID, opts, CompactionTriggerToolOutput, result.Reason) - skipNotified = true - } - return messages, skipNotified, nil - } - if !result.Compacted { - if result.Due && !skipNotified { - s.emitCompactionSkipped(runID, opts, CompactionTriggerToolOutput, result.Reason) - skipNotified = true - } - return messages, skipNotified, nil - } - - overflowByID := make(map[string]toolOutputOverflow, len(overflows)) - for _, overflow := range overflows { - overflowByID[overflow.toolCallID] = overflow - } - - compacted := append([]api.Message(nil), result.Messages...) - if !messageEmpty(assistant) { - compacted = append(compacted, assistant) - } - - historyTokens := s.estimateRunPromptTokens(opts, compacted) - batchTokens := 0 - for _, msg := range toolMessages { - content := msg.Content - toolName := msg.ToolName - if overflow, ok := overflowByID[msg.ToolCallID]; ok { - content = overflow.content - if overflow.toolName != "" { - toolName = overflow.toolName - } - } - refit := s.toolMessageForPostCompactionContext(toolName, msg.ToolCallID, content, opts, historyTokens+batchTokens) - compacted = append(compacted, refit) - batchTokens += estimateMessagesTokens([]api.Message{refit}) - } - - s.emitCompacted(runID, opts, compacted, CompactionTriggerToolOutput, result.Summary) - if err := s.checkPostCompactionPromptBudget(opts, compacted); err != nil { - return compacted, skipNotified, err - } - return compacted, skipNotified, nil -} - -func (s *Session) compactionRequest(runID string, opts RunOptions, messages []api.Message, latest api.ChatResponse) CompactionRequest { - meta := newEventMetadata(runID, opts) - return CompactionRequest{ - ChatID: opts.ChatID, - Model: opts.Model, - SystemPrompt: opts.SystemPrompt, - Messages: messages, - Tools: s.availableTools(), - Format: opts.Format, - Latest: latest, - Options: opts.Options, - KeepAlive: opts.KeepAlive, - Think: opts.Think, - ContinueTask: true, - Progress: func(progress CompactionProgress) { - _ = s.emit(newCompactionProgress(meta, progress.Tokens)) - }, - } -} - -func (s *Session) emitCompactionStarted(runID string, opts RunOptions, trigger CompactionTrigger) { - _ = s.emit(newCompactionStarted(newEventMetadata(runID, opts), trigger)) -} - -func (s *Session) emitCompactionSkipped(runID string, opts RunOptions, trigger CompactionTrigger, reason string) { - _ = s.emit(newCompactionSkipped(newEventMetadata(runID, opts), trigger, CompactionSkippedMessage(reason))) -} - -func (s *Session) emitCompacted(runID string, opts RunOptions, messages []api.Message, trigger CompactionTrigger, summary string) { - _ = s.emit(newCompacted(newEventMetadata(runID, opts), messages, trigger, summary)) -} - -func (s *Session) autoCompactionTrigger(req CompactionRequest) CompactionTrigger { - if s.Compactor == nil { - return "" - } - trigger, should := s.Compactor.ShouldCompact(req) - if should { - return CompactionTrigger(trigger) - } - return "" -} - -func CompactionSkippedMessage(reason string) string { - reason = strings.TrimSpace(reason) - if reason == "" { - reason = "compaction could not run" - } - return reason -} - -func resolvedMaxToolRounds(model string, value int) int { - if value != 0 { - return value - } - if modelref.HasExplicitCloudSource(model) { - return -1 - } - return defaultMaxToolRounds -} - -// toolMessageWithBudget sizes a tool result message to fit within a token -// budget (compaction threshold or context window). baseTokens is the -// pre-computed estimate of everything before this message; budgetTokens is -// the ceiling. If the message already fits, it is returned with only the -// small-context rune cap applied. -func (s *Session) toolMessageWithBudget(toolName, toolCallID, content string, opts RunOptions, baseTokens, budgetTokens int) api.Message { - maxRunes := maxToolResultRunes - if limit := smallContextToolResultLimitRunes(s.contextWindowTokens(opts)); limit > 0 { - maxRunes = min(maxRunes, limit) - } - - if budgetTokens <= 0 { - return toolMessageWithLimit(toolName, toolCallID, content, maxRunes) - } - - msg := toolMessageWithLimit(toolName, toolCallID, content, maxRunes) - projectedTokens := baseTokens + estimateMessagesTokens([]api.Message{msg}) - if projectedTokens < budgetTokens { - return msg - } - - overheadTokens := estimateMessagesTokens([]api.Message{{ - Role: "tool", - ToolName: toolName, - ToolCallID: toolCallID, - }}) - // Keep oversized tool output below the budget before it is appended to - // history. This is especially important for <=8k contexts: the next step - // must still have enough room to compact and continue the same user - // request instead of asking the user to prompt again. - availableRunes := (budgetTokens - baseTokens - overheadTokens - toolTruncationMarkerReserveTokens) * 4 - maxRunes = min(maxRunes, max(0, availableRunes)) - msg.Content = truncateToolResultContentTo(content, maxRunes) - return msg -} - -func (s *Session) toolMessageForContext(toolName, toolCallID, content string, opts RunOptions, baseTokens int) api.Message { - return s.toolMessageWithBudget(toolName, toolCallID, content, opts, baseTokens, s.compactionThresholdTokens(opts)) -} - -func (s *Session) toolMessageForPostCompactionContext(toolName, toolCallID, content string, opts RunOptions, baseTokens int) api.Message { - return s.toolMessageWithBudget(toolName, toolCallID, content, opts, baseTokens, s.contextWindowTokens(opts)) -} - -func toolMessageWithLimit(toolName, toolCallID, content string, maxRunes int) api.Message { - return api.Message{ - Role: "tool", - Content: truncateToolResultContentTo(content, maxRunes), - ToolName: toolName, - ToolCallID: toolCallID, - } -} - -func smallContextToolResultLimitRunes(contextWindow int) int { - switch { - case contextWindow > 0 && contextWindow <= tinyContextToolResultTokenWindow: - return tinyContextToolResultRunes - case contextWindow > 0 && contextWindow <= smallContextToolResultTokenWindow: - return smallContextToolResultRunes - default: - return 0 - } -} - -func (s *Session) availableTools() api.Tools { - if s == nil || s.Tools == nil { - return nil - } - return s.Tools.Tools() -} - -func (s *Session) compactionThresholdTokens(opts RunOptions) int { - contextWindow := s.contextWindowTokens(opts) - if contextWindow <= 0 { - return 0 - } - - configuredThreshold := 0.0 - if s.Compactor != nil { - configuredThreshold = s.Compactor.Threshold() - } - - threshold := int(float64(contextWindow) * ResolveCompactionThreshold(configuredThreshold)) - if threshold <= 0 { - return 0 - } - return threshold -} - -func (s *Session) contextWindowTokens(opts RunOptions) int { - if s.Compactor == nil { - return 0 - } - return s.Compactor.ContextWindowTokens(opts.Options) -} - -func toolMessage(toolName, toolCallID, content string) api.Message { - return toolMessageWithLimit(toolName, toolCallID, content, maxToolResultRunes) -} - -func sanitizeMessageForRun(msg api.Message) api.Message { - if msg.Role == "tool" { - msg.Content = truncateToolResultContent(msg.Content) - } - return msg -} - -func sanitizeMessagesForRequest(messages []api.Message) []api.Message { - if len(messages) == 0 { - return nil - } - sanitized := make([]api.Message, len(messages)) - for i, msg := range messages { - sanitized[i] = sanitizeMessageForRun(msg) - } - return sanitized -} - -func truncateToolResultContent(content string) string { - return truncateToolResultContentTo(content, maxToolResultRunes) -} - -func truncateToolResultContentTo(content string, maxRunes int) string { - return Truncate(content, TruncateConfig{ - MaxRunes: maxRunes, - HeadTail: true, - HeadPct: 75, - Label: "tool output", - Hint: "Use a narrower command, line range, or search query if more detail is needed.", - FullOmissionPrefix: toolOutputFullOmissionPrefix, - }) -} - -// TruncateConfig configures content truncation via Truncate. -type TruncateConfig struct { - MaxRunes int // rune limit; <= 0 means full omission - HeadTail bool // true = head + tail split; false = head only - HeadPct int // percentage of MaxRunes for head (e.g. 75); tail gets the rest - Label string // e.g. "tool output", "summary", "stdout" - Hint string // guidance text appended to marker (optional) - FullOmissionPrefix string // marker prefix when MaxRunes <= 0 -} - -// Truncate truncates content to at most cfg.MaxRunes runes. When HeadTail is -// true, it preserves the first HeadPct% and last (100-HeadPct)% of the budget -// with a marker between; otherwise it keeps only the head. MaxRunes <= 0 -// triggers full omission using FullOmissionPrefix. All token counts in -// markers use ApproximateTokens. -func Truncate(content string, cfg TruncateConfig) string { - runes := []rune(content) - total := len(runes) - - if cfg.MaxRunes <= 0 { - return fmt.Sprintf("%s omitted ~%d tokens.%s]", cfg.FullOmissionPrefix, ApproximateTokens(total), truncHint(cfg.Hint)) - } - if total <= cfg.MaxRunes { - return content - } - - if !cfg.HeadTail { - head := cfg.MaxRunes - omitted := total - head - return string(runes[:head]) + TruncMarker(cfg.Label, head, 0, omitted, false, cfg.Hint) - } - - head := cfg.MaxRunes * cfg.HeadPct / 100 - tail := cfg.MaxRunes - head - omitted := total - head - tail - return string(runes[:head]) + TruncMarker(cfg.Label, head, tail, omitted, true, cfg.Hint) + string(runes[len(runes)-tail:]) -} - -func truncHint(hint string) string { - hint = strings.TrimSpace(hint) - if hint == "" { - return "" - } - if !strings.HasSuffix(hint, ".") { - hint += "." - } - return " " + hint -} - -// TruncMarker formats a truncation marker with consistent wording. head and -// tail are rune counts; omitted is the count of runes removed. headTail -// selects the head+tail vs head-only format. hint is optional guidance text. -func TruncMarker(label string, head, tail, omitted int, headTail bool, hint string) string { - var b strings.Builder - b.WriteString("\n\n[") - b.WriteString(label) - b.WriteString(" truncated: ") - if headTail { - fmt.Fprintf(&b, "showing first ~%d tokens and last ~%d tokens; ", ApproximateTokens(head), ApproximateTokens(tail)) - } else { - fmt.Fprintf(&b, "showing first ~%d tokens; ", ApproximateTokens(head)) - } - fmt.Fprintf(&b, "omitted ~%d tokens.%s]", ApproximateTokens(omitted), truncHint(hint)) - if headTail { - b.WriteString("\n\n") - } - return b.String() -} - -func toolOutputFullyOmitted(content string) bool { - return strings.HasPrefix(content, toolOutputFullOmissionPrefix) -} - -// ApproximateTokens estimates token count from a character/byte count using -// the standard ~4 chars-per-token heuristic. It is intentionally rough; all -// callers use it only for sizing/truncation decisions, not billing. -func ApproximateTokens(n int) int { - if n <= 0 { - return 0 - } - return max(1, (n+3)/4) -} - -func messageEmpty(msg api.Message) bool { - return msg.Content == "" && msg.Thinking == "" && len(msg.ToolCalls) == 0 -} diff --git a/agent/session_test.go b/agent/session_test.go deleted file mode 100644 index e95db328090..00000000000 --- a/agent/session_test.go +++ /dev/null @@ -1,2265 +0,0 @@ -package agent - -import ( - "context" - "encoding/json" - "errors" - "os" - "path/filepath" - "strings" - "testing" - - "github.com/ollama/ollama/api" -) - -type fakeClient struct { - calls int - responses [][]api.ChatResponse - requests []*api.ChatRequest - err error -} - -func (c *fakeClient) Chat(ctx context.Context, req *api.ChatRequest, fn api.ChatResponseFunc) error { - c.requests = append(c.requests, req) - if c.calls >= len(c.responses) { - return nil - } - responses := c.responses[c.calls] - c.calls++ - for _, response := range responses { - if err := fn(response); err != nil { - return err - } - } - return c.err -} - -type staticTool struct{} - -type approvalTestTool struct { - called *bool -} - -type namedApprovalTestTool struct { - name string -} - -type cwdTestTool struct{} - -type largeTool struct{} - -type preTruncatedTool struct{} - -type cancelingTool struct { - cancel context.CancelFunc -} - -type cancelAfterToolCallClient struct { - cancel context.CancelFunc -} - -type recordingCompactor struct { - requests []CompactionRequest -} - -type oversizedCompactor struct { - requests []CompactionRequest -} - -type recordingEventSink struct { - events []Event -} - -func (s *recordingEventSink) Emit(event Event) error { - s.events = append(s.events, event) - return nil -} - -func hasEventType(events []Event, eventType EventType) bool { - for _, event := range events { - if event.Type == eventType { - return true - } - } - return false -} - -func hasEventWithTokens(events []Event, eventType EventType, tokens int) bool { - for _, event := range events { - if event.Type == eventType && event.Tokens == tokens { - return true - } - } - return false -} - -func TestSessionEmitsToAllSinksAfterError(t *testing.T) { - errSink := EventSinkFunc(func(Event) error { - return errors.New("sink failed") - }) - events := &recordingEventSink{} - session := &Session{EventSinks: []EventSink{errSink, events}} - - err := session.emit(Event{Type: EventRunFinished}) - if err == nil { - t.Fatal("emit should return the first sink error") - } - if !hasEventType(events.events, EventRunFinished) { - t.Fatalf("later sink did not receive event after earlier error: %#v", events.events) - } -} - -func (c cancelAfterToolCallClient) Chat(ctx context.Context, req *api.ChatRequest, fn api.ChatResponseFunc) error { - args := api.NewToolCallFunctionArguments() - args.Set("value", "skip me") - if err := fn(api.ChatResponse{Message: api.Message{Role: "assistant", ToolCalls: []api.ToolCall{{ - ID: "call-1", - Function: api.ToolCallFunction{ - Name: "echo_tool", - Arguments: args, - }, - }}}}); err != nil { - return err - } - c.cancel() - return context.Canceled -} - -func (c *recordingCompactor) MaybeCompact(_ context.Context, req CompactionRequest) (CompactionResult, error) { - c.requests = append(c.requests, req) - result := CompactionResult{Messages: req.Messages, Due: true} - if len(req.Messages) > 0 && req.Messages[len(req.Messages)-1].Role == "tool" { - result.Messages = CompactionSummaryMessages("tool result summarized", false) - result.Compacted = true - result.Summary = "tool result summarized" - } - return result, nil -} - -func (c *recordingCompactor) ContextWindowTokens(options map[string]any) int { - return ResolveContextWindowTokens(options, 0) -} -func (c *recordingCompactor) Threshold() float64 { return 0 } -func (c *recordingCompactor) ShouldCompact(_ CompactionRequest) (string, bool) { - return "", false -} - -func (c *oversizedCompactor) MaybeCompact(_ context.Context, req CompactionRequest) (CompactionResult, error) { - c.requests = append(c.requests, req) - summary := strings.Repeat("oversized summary ", 300) - return CompactionResult{ - Messages: CompactionSummaryMessages(summary, req.ContinueTask), - Compacted: true, - Due: true, - Summary: summary, - }, nil -} - -func (c *oversizedCompactor) ContextWindowTokens(options map[string]any) int { - return ResolveContextWindowTokens(options, 0) -} -func (c *oversizedCompactor) Threshold() float64 { return 0 } -func (c *oversizedCompactor) ShouldCompact(_ CompactionRequest) (string, bool) { - return "", false -} - -type recordingApprovalPrompter struct { - requests []ApprovalRequest - results []Approval -} - -func approvalStateForTest(allowAll bool, scopes map[string]bool) *ApprovalState { - state := &ApprovalState{} - state.Set(allowAll, scopes) - return state -} - -func (staticTool) Name() string { - return "echo_tool" -} - -func (staticTool) Description() string { - return "echoes a value" -} - -func (staticTool) Schema() api.ToolFunction { - props := api.NewToolPropertiesMap() - props.Set("value", api.ToolProperty{Type: api.PropertyType{"string"}}) - return api.ToolFunction{ - Name: "echo_tool", - Description: "echoes a value", - Parameters: api.ToolFunctionParameters{ - Type: "object", - Properties: props, - }, - } -} - -func (staticTool) Execute(context.Context, ToolContext, map[string]any) (ToolResult, error) { - return ToolResult{Content: "tool says hello"}, nil -} - -func (largeTool) Name() string { - return "large_tool" -} - -func (largeTool) Description() string { - return "returns a large result" -} - -func (largeTool) Schema() api.ToolFunction { - return api.ToolFunction{ - Name: "large_tool", - Description: "returns a large result", - Parameters: api.ToolFunctionParameters{ - Type: "object", - }, - } -} - -func (largeTool) Execute(context.Context, ToolContext, map[string]any) (ToolResult, error) { - return ToolResult{Content: strings.Repeat("x", maxToolResultRunes+100)}, nil -} - -func (preTruncatedTool) Name() string { - return "pre_truncated_tool" -} - -func (preTruncatedTool) Description() string { - return "returns a large result that is already marked as truncated" -} - -func (preTruncatedTool) Schema() api.ToolFunction { - return api.ToolFunction{ - Name: "pre_truncated_tool", - Description: "returns a large result that is already marked as truncated", - Parameters: api.ToolFunctionParameters{ - Type: "object", - }, - } -} - -func (preTruncatedTool) Execute(context.Context, ToolContext, map[string]any) (ToolResult, error) { - content := strings.Repeat("x", smallContextToolResultRunes) + - "\n\n[tool output truncated: showing first ~1500 tokens; omitted ~999 tokens. Use a narrower command, line range, or search query if more detail is needed.]\n\n" + - strings.Repeat("y", smallContextToolResultRunes) - return ToolResult{Content: content}, nil -} - -func (t cancelingTool) Name() string { - return "cancel_tool" -} - -func (t cancelingTool) Description() string { - return "cancels while running" -} - -func (t cancelingTool) Schema() api.ToolFunction { - return api.ToolFunction{ - Name: t.Name(), - Description: t.Description(), - Parameters: api.ToolFunctionParameters{ - Type: "object", - }, - } -} - -func (t cancelingTool) Execute(ctx context.Context, _ ToolContext, _ map[string]any) (ToolResult, error) { - t.cancel() - <-ctx.Done() - return ToolResult{}, ctx.Err() -} - -func (t approvalTestTool) Name() string { - return "approval_tool" -} - -func (t approvalTestTool) Description() string { - return "requires approval" -} - -func (t approvalTestTool) Schema() api.ToolFunction { - return api.ToolFunction{ - Name: "approval_tool", - Description: "requires approval", - Parameters: api.ToolFunctionParameters{ - Type: "object", - }, - } -} - -func (t approvalTestTool) RequiresApproval(map[string]any) bool { - return true -} - -func (t approvalTestTool) Execute(context.Context, ToolContext, map[string]any) (ToolResult, error) { - if t.called != nil { - *t.called = true - } - return ToolResult{Content: "approved"}, nil -} - -func (t namedApprovalTestTool) Name() string { - return t.name -} - -func (t namedApprovalTestTool) Description() string { - return "requires approval" -} - -func (t namedApprovalTestTool) Schema() api.ToolFunction { - return api.ToolFunction{ - Name: t.name, - Description: "requires approval", - Parameters: api.ToolFunctionParameters{ - Type: "object", - }, - } -} - -func (t namedApprovalTestTool) RequiresApproval(map[string]any) bool { - return true -} - -func (t namedApprovalTestTool) Execute(context.Context, ToolContext, map[string]any) (ToolResult, error) { - return ToolResult{Content: "approved"}, nil -} - -// ApprovalScope mimics the Bash tool's command-scoping behavior so tests can -// exercise the shell approval flow without importing the tools package. -func (t namedApprovalTestTool) ApprovalScope(args map[string]any) string { - if t.name == "bash" || t.name == "powershell" { - if cmd, ok := args["command"].(string); ok { - cmd = strings.TrimSpace(cmd) - if cmd != "" { - return t.name + "\x00" + cmd - } - } - } - return t.name -} - -func (p *recordingApprovalPrompter) PromptApproval(_ context.Context, req ApprovalRequest) (Approval, error) { - p.requests = append(p.requests, req) - if len(p.results) == 0 { - return Approval{Allow: true}, nil - } - result := p.results[0] - p.results = p.results[1:] - return result, nil -} - -func (cwdTestTool) Name() string { - return "cwd_tool" -} - -func (cwdTestTool) Description() string { - return "tests cwd state" -} - -func (cwdTestTool) Schema() api.ToolFunction { - props := api.NewToolPropertiesMap() - props.Set("mode", api.ToolProperty{Type: api.PropertyType{"string"}}) - props.Set("path", api.ToolProperty{Type: api.PropertyType{"string"}}) - return api.ToolFunction{ - Name: "cwd_tool", - Description: "tests cwd state", - Parameters: api.ToolFunctionParameters{ - Type: "object", - Properties: props, - }, - } -} - -func (cwdTestTool) RequiresApproval(map[string]any) bool { - return true -} - -func (cwdTestTool) Execute(_ context.Context, toolCtx ToolContext, args map[string]any) (ToolResult, error) { - switch args["mode"] { - case "set": - path, _ := args["path"].(string) - return ToolResult{Content: "changed", WorkingDir: filepath.Join(toolCtx.WorkingDir, path)}, nil - case "escape": - return ToolResult{Content: "escaped", WorkingDir: filepath.Dir(toolCtx.WorkingDir)}, nil - default: - return ToolResult{Content: toolCtx.WorkingDir}, nil - } -} - -func TestSessionRunsToolLoop(t *testing.T) { - args := api.NewToolCallFunctionArguments() - args.Set("value", "hello") - - client := &fakeClient{ - responses: [][]api.ChatResponse{ - { - {Message: api.Message{Role: "assistant", ToolCalls: []api.ToolCall{{ - ID: "call-1", - Function: api.ToolCallFunction{ - Name: "echo_tool", - Arguments: args, - }, - }}}}, - }, - { - {Message: api.Message{Role: "assistant", Content: "done"}}, - }, - }, - } - - registry := &Registry{} - registry.Register(staticTool{}) - - session := &Session{ - Client: client, - Tools: registry, - ApprovalState: approvalStateForTest(true, nil), - } - - result, err := session.Run(context.Background(), RunOptions{ - ChatID: "chat-1", - Model: "model", - NewMessages: []api.Message{{Role: "user", Content: "use a tool"}}, - }) - if err != nil { - t.Fatal(err) - } - - if client.calls != 2 { - t.Fatalf("client calls = %d, want 2", client.calls) - } - if len(result.Messages) != 4 { - t.Fatalf("messages = %d, want 4", len(result.Messages)) - } - if result.Messages[2].Role != "tool" || result.Messages[2].Content != "tool says hello" { - t.Fatalf("tool message = %#v", result.Messages[2]) - } - if len(client.requests[0].Tools) != 1 { - t.Fatalf("first request tools = %d, want 1", len(client.requests[0].Tools)) - } - if len(client.requests[1].Messages) != 3 { - t.Fatalf("second request messages = %d, want 3", len(client.requests[1].Messages)) - } -} - -func TestSessionAddsSystemPromptOnlyToRequest(t *testing.T) { - client := &fakeClient{ - responses: [][]api.ChatResponse{ - {{Message: api.Message{Role: "assistant", Content: "done"}}}, - }, - } - session := &Session{Client: client} - - _, err := session.Run(context.Background(), RunOptions{ - ChatID: "chat-1", - Model: "model", - SystemPrompt: "available context: go-code", - NewMessages: []api.Message{{Role: "user", Content: "hello"}}, - }) - if err != nil { - t.Fatal(err) - } - if len(client.requests) != 1 { - t.Fatalf("requests = %d, want 1", len(client.requests)) - } - reqMessages := client.requests[0].Messages - if len(reqMessages) != 2 || reqMessages[0].Role != "system" || reqMessages[0].Content != "available context: go-code" { - t.Fatalf("request messages = %#v", reqMessages) - } -} - -func TestSessionChatRequestMatchesRunRequest(t *testing.T) { - client := &fakeClient{ - responses: [][]api.ChatResponse{ - {{Message: api.Message{Role: "assistant", Content: "done"}}}, - }, - } - registry := &Registry{} - registry.Register(staticTool{}) - session := &Session{Client: client, Tools: registry} - opts := RunOptions{ - ChatID: "chat-1", - Model: "model", - SystemPrompt: "available context: go-code", - NewMessages: []api.Message{{Role: "user", Content: "hello"}}, - Format: "json", - Options: map[string]any{"temperature": 0.5}, - } - - want := buildChatRequest(opts, opts.NewMessages, registry.Tools()) - _, err := session.Run(context.Background(), opts) - if err != nil { - t.Fatal(err) - } - if len(client.requests) != 1 { - t.Fatalf("requests = %d, want 1", len(client.requests)) - } - gotJSON, err := json.Marshal(client.requests[0]) - if err != nil { - t.Fatal(err) - } - wantJSON, err := json.Marshal(want) - if err != nil { - t.Fatal(err) - } - if string(gotJSON) != string(wantJSON) { - t.Fatalf("ChatRequest mismatch\ngot: %s\nwant: %s", gotJSON, wantJSON) - } -} - -func TestSessionAccumulatesStreamingAssistantMessage(t *testing.T) { - args := api.NewToolCallFunctionArguments() - args.Set("value", "hello") - responses := make([]api.ChatResponse, 0, 100) - var wantContent, wantThinking string - for range 99 { - wantContent += "x" - wantThinking += "t" - responses = append(responses, api.ChatResponse{ - Message: api.Message{Role: "assistant", Content: "x", Thinking: "t"}, - }) - } - toolCall := api.ToolCall{ - ID: "call-1", - Function: api.ToolCallFunction{ - Name: "echo_tool", - Arguments: args, - }, - } - responses = append(responses, api.ChatResponse{ - Message: api.Message{Role: "assistant", ToolCalls: []api.ToolCall{toolCall}}, - }) - - session := &Session{ - Client: &fakeClient{responses: [][]api.ChatResponse{responses}}, - } - - result, err := session.Run(context.Background(), RunOptions{ - ChatID: "chat-1", - Model: "model", - NewMessages: []api.Message{{Role: "user", Content: "stream"}}, - }) - if err != nil { - t.Fatal(err) - } - - if len(result.Messages) != 2 || result.Messages[1].Content != wantContent || result.Messages[1].Thinking != wantThinking || len(result.Messages[1].ToolCalls) != 1 { - t.Fatalf("result messages = %#v", result.Messages) - } -} - -func TestSessionRequestHistoryKeepsThinkingAndServerToolCallIDs(t *testing.T) { - args := api.NewToolCallFunctionArguments() - args.Set("value", "hello") - client := &fakeClient{ - responses: [][]api.ChatResponse{ - { - {Message: api.Message{Role: "assistant", Thinking: "private chain"}}, - {Message: api.Message{Role: "assistant", Content: "I'll check."}}, - {Message: api.Message{Role: "assistant", ToolCalls: []api.ToolCall{{ - ID: "volatile-random-id", - Function: api.ToolCallFunction{ - Name: "echo_tool", - Arguments: args, - }, - }}}}, - }, - {{Message: api.Message{Role: "assistant", Content: "done"}}}, - }, - } - registry := &Registry{} - registry.Register(staticTool{}) - session := &Session{ - Client: client, - Tools: registry, - ApprovalState: approvalStateForTest(true, nil), - } - - result, err := session.Run(context.Background(), RunOptions{ - ChatID: "chat-1", - Model: "model", - NewMessages: []api.Message{{Role: "user", Content: "use a tool"}}, - }) - if err != nil { - t.Fatal(err) - } - if len(client.requests) != 2 { - t.Fatalf("requests = %d, want 2", len(client.requests)) - } - - secondRequestMessages := client.requests[1].Messages - if len(secondRequestMessages) != 3 { - t.Fatalf("second request messages = %#v", secondRequestMessages) - } - assistant := secondRequestMessages[1] - if assistant.Role != "assistant" { - t.Fatalf("second request assistant = %#v", assistant) - } - if assistant.Thinking != "private chain" { - t.Fatalf("assistant thinking = %q, want preserved", assistant.Thinking) - } - if len(assistant.ToolCalls) != 1 || assistant.ToolCalls[0].ID != "volatile-random-id" { - t.Fatalf("assistant tool calls = %#v", assistant.ToolCalls) - } - tool := secondRequestMessages[2] - if tool.Role != "tool" || tool.ToolCallID != "volatile-random-id" { - t.Fatalf("tool result message = %#v", tool) - } - if len(result.Messages) < 3 || result.Messages[1].Thinking != "private chain" { - t.Fatalf("visible result messages lost thinking: %#v", result.Messages) - } -} - -func TestSessionKeepsPartialStreamOnCancellation(t *testing.T) { - session := &Session{ - Client: &fakeClient{ - responses: [][]api.ChatResponse{{ - {Message: api.Message{Role: "assistant", Content: "partial "}}, - {Message: api.Message{Role: "assistant", Content: "answer"}}, - }}, - err: context.Canceled, - }, - } - - result, err := session.Run(context.Background(), RunOptions{ - ChatID: "chat-1", - Model: "model", - NewMessages: []api.Message{{Role: "user", Content: "cancel"}}, - }) - if err != nil { - t.Fatal(err) - } - if len(result.Messages) != 2 || result.Messages[1].Content != "partial answer" { - t.Fatalf("result messages = %#v", result.Messages) - } -} - -func TestSessionCancellationKeepsPartialResultWhenUISinkCancels(t *testing.T) { - ctx, cancel := context.WithCancel(context.Background()) - cancel() - trace := &recordingEventSink{} - session := &Session{ - Client: &fakeClient{ - responses: [][]api.ChatResponse{{ - {Message: api.Message{Role: "assistant", Content: "partial"}}, - }}, - err: context.Canceled, - }, - EventSinks: []EventSink{ - EventSinkFunc(func(event Event) error { - if event.Type == EventRunFinished { - return context.Canceled - } - return nil - }), - trace, - }, - } - - result, err := session.Run(ctx, RunOptions{ - ChatID: "chat-1", - Model: "model", - NewMessages: []api.Message{{Role: "user", Content: "cancel"}}, - }) - if err != nil { - t.Fatal(err) - } - if result == nil || len(result.Messages) != 2 || result.Messages[1].Content != "partial" { - t.Fatalf("result messages = %#v, want partial assistant result", result) - } - if !hasEventType(trace.events, EventRunFinished) { - t.Fatalf("trace sink did not receive run finished event: %#v", trace.events) - } -} - -func TestSessionTreatsHTTPContextCanceledStringAsCancellation(t *testing.T) { - ctx, cancel := context.WithCancel(context.Background()) - cancel() - client := &fakeClient{err: errors.New(`Post "http://127.0.0.1:11434/api/chat": context canceled`)} - session := &Session{Client: client} - - result, err := session.Run(ctx, RunOptions{ - Model: "model", - NewMessages: []api.Message{{Role: "user", Content: "hello"}}, - }) - if err != nil { - t.Fatalf("Run returned error for canceled HTTP request: %v", err) - } - if result == nil { - t.Fatal("Run returned nil result") - } - if len(result.Messages) != 1 || result.Messages[0].Content != "hello" { - t.Fatalf("messages = %#v, want original user message only", result.Messages) - } -} - -func TestSessionDisabledToolsOmitToolsAndReturnsDisabledResults(t *testing.T) { - args := api.NewToolCallFunctionArguments() - args.Set("value", "hello") - client := &fakeClient{ - responses: [][]api.ChatResponse{ - {{ - Message: api.Message{Role: "assistant", ToolCalls: []api.ToolCall{{ - ID: "call-1", - Function: api.ToolCallFunction{ - Name: "echo_tool", - Arguments: args, - }, - }}}, - }}, - {{Message: api.Message{Role: "assistant", Content: "tools are off"}}}, - }, - } - registry := &Registry{} - registry.Register(staticTool{}) - events := &recordingEventSink{} - session := &Session{ - Client: client, - EventSinks: []EventSink{events}, - Tools: registry, - DisableTools: true, - } - - result, err := session.Run(context.Background(), RunOptions{ - ChatID: "chat-1", - Model: "model", - NewMessages: []api.Message{{Role: "user", Content: "use a tool"}}, - }) - if err != nil { - t.Fatal(err) - } - if len(client.requests) != 2 { - t.Fatalf("requests = %d, want 2", len(client.requests)) - } - if got := len(client.requests[0].Tools); got != 0 { - t.Fatalf("advertised tools = %d, want 0", got) - } - secondMessages := client.requests[1].Messages - if len(secondMessages) != 3 { - t.Fatalf("second request messages = %#v", secondMessages) - } - if secondMessages[2].Role != "tool" || secondMessages[2].ToolCallID != "call-1" || secondMessages[2].Content != toolExecutionDisabledMessage { - t.Fatalf("disabled tool message = %#v", secondMessages[2]) - } - if len(result.Messages) != 4 || result.Messages[2].Content != toolExecutionDisabledMessage { - t.Fatalf("result messages = %#v", result.Messages) - } - var sawDetected, sawDisabled bool - for _, event := range events.events { - if event.Type == EventToolCallDetected { - sawDetected = true - } - if event.Type == EventToolFinished && event.ToolStatus == ToolStatusDisabled && event.Content == toolExecutionDisabledMessage { - sawDisabled = true - } - } - if !sawDetected || !sawDisabled { - t.Fatalf("events missing detected/disabled: %#v", events.events) - } -} - -func TestSessionCancellationAfterToolCallAppendsSkippedToolMessage(t *testing.T) { - ctx, cancel := context.WithCancel(context.Background()) - registry := &Registry{} - registry.Register(staticTool{}) - session := &Session{ - Client: cancelAfterToolCallClient{cancel: cancel}, - Tools: registry, - ApprovalState: approvalStateForTest(true, nil), - } - - result, err := session.Run(ctx, RunOptions{ - ChatID: "chat-1", - Model: "model", - NewMessages: []api.Message{{Role: "user", Content: "cancel after tool"}}, - }) - if err != nil { - t.Fatal(err) - } - if len(result.Messages) != 3 { - t.Fatalf("messages = %#v", result.Messages) - } - if len(result.Messages[1].ToolCalls) != 1 { - t.Fatalf("assistant tool calls = %#v", result.Messages[1]) - } - if result.Messages[2].Role != "tool" || result.Messages[2].ToolCallID != "call-1" { - t.Fatalf("skipped tool message = %#v", result.Messages[2]) - } - if !strings.Contains(result.Messages[2].Content, "run was canceled") { - t.Fatalf("skipped content = %q", result.Messages[2].Content) - } -} - -func TestSessionCancellationDuringToolExecutionAppendsToolMessage(t *testing.T) { - ctx, cancel := context.WithCancel(context.Background()) - events := &recordingEventSink{} - registry := &Registry{} - registry.Register(cancelingTool{cancel: cancel}) - client := &fakeClient{responses: [][]api.ChatResponse{{ - {Message: api.Message{Role: "assistant", ToolCalls: []api.ToolCall{{ - ID: "call-1", - Function: api.ToolCallFunction{ - Name: "cancel_tool", - }, - }}}}, - }}} - session := &Session{ - Client: client, - Tools: registry, - ApprovalState: approvalStateForTest(true, nil), - EventSinks: []EventSink{events}, - } - - result, err := session.Run(ctx, RunOptions{ - ChatID: "chat-1", - Model: "model", - NewMessages: []api.Message{{Role: "user", Content: "cancel during tool"}}, - }) - if err != nil { - t.Fatal(err) - } - if len(result.Messages) != 3 { - t.Fatalf("messages = %#v", result.Messages) - } - if result.Messages[2].Role != "tool" || result.Messages[2].ToolCallID != "call-1" { - t.Fatalf("tool message = %#v", result.Messages[2]) - } - if !strings.Contains(result.Messages[2].Content, "context canceled") { - t.Fatalf("tool content = %q", result.Messages[2].Content) - } - var finished *Event - for i := range events.events { - if events.events[i].Type == EventRunFinished { - finished = &events.events[i] - } - } - if finished == nil { - t.Fatalf("run finished event missing: %#v", events.events) - } - if finished.Status != RunStatusCanceled { - t.Fatalf("run status = %q, want canceled", finished.Status) - } -} - -func TestSessionToolLoopAllowsRoundsUnderDefaultCap(t *testing.T) { - responses := make([][]api.ChatResponse, 0, 26) - for i := range 25 { - args := api.NewToolCallFunctionArguments() - args.Set("value", "hello") - responses = append(responses, []api.ChatResponse{{ - Message: api.Message{Role: "assistant", ToolCalls: []api.ToolCall{{ - ID: "call-" + string(rune('a'+i)), - Function: api.ToolCallFunction{ - Name: "echo_tool", - Arguments: args, - }, - }}}, - }}) - } - responses = append(responses, []api.ChatResponse{{ - Message: api.Message{Role: "assistant", Content: "done"}, - }}) - - client := &fakeClient{responses: responses} - registry := &Registry{} - registry.Register(staticTool{}) - - session := &Session{ - Client: client, - Tools: registry, - ApprovalState: approvalStateForTest(true, nil), - } - - if _, err := session.Run(context.Background(), RunOptions{ - Model: "model", - NewMessages: []api.Message{{Role: "user", Content: "keep going"}}, - }); err != nil { - t.Fatal(err) - } - - if client.calls != 26 { - t.Fatalf("client calls = %d, want 26", client.calls) - } -} - -func TestSessionLocalToolRoundLimitAppendsSkippedToolMessages(t *testing.T) { - firstArgs := api.NewToolCallFunctionArguments() - firstArgs.Set("value", "first") - secondArgs := api.NewToolCallFunctionArguments() - secondArgs.Set("value", "second") - thirdArgs := api.NewToolCallFunctionArguments() - thirdArgs.Set("value", "third") - - client := &fakeClient{ - responses: [][]api.ChatResponse{ - {{ - Message: api.Message{Role: "assistant", ToolCalls: []api.ToolCall{{ - ID: "call-1", - Function: api.ToolCallFunction{ - Name: "echo_tool", - Arguments: firstArgs, - }, - }}}, - }}, - {{ - Message: api.Message{Role: "assistant", ToolCalls: []api.ToolCall{ - { - ID: "call-2", - Function: api.ToolCallFunction{ - Name: "echo_tool", - Arguments: secondArgs, - }, - }, - { - ID: "call-3", - Function: api.ToolCallFunction{ - Name: "echo_tool", - Arguments: thirdArgs, - }, - }, - }}, - }}, - }, - } - registry := &Registry{} - registry.Register(staticTool{}) - session := &Session{ - Client: client, - Tools: registry, - ApprovalState: approvalStateForTest(true, nil), - } - - result, err := session.Run(context.Background(), RunOptions{ - ChatID: "chat-1", - Model: "test:local", - NewMessages: []api.Message{{Role: "user", Content: "hit cap"}}, - MaxToolRounds: 1, - }) - if err == nil || !strings.Contains(err.Error(), "tool round limit reached after 1 rounds") { - t.Fatalf("error = %v, want tool-round limit", err) - } - if result == nil { - t.Fatal("expected partial result with skipped tool messages") - } - if len(result.Messages) != 6 { - t.Fatalf("messages = %#v", result.Messages) - } - for i, wantID := range []string{"call-2", "call-3"} { - msg := result.Messages[4+i] - if msg.Role != "tool" || msg.ToolCallID != wantID { - t.Fatalf("skipped tool %d = %#v", i, msg) - } - if !strings.Contains(msg.Content, "max tool-round limit of 1") { - t.Fatalf("skipped content = %q", msg.Content) - } - } -} - -func TestSessionLocalToolLoopStopsAtDefaultRoundCap(t *testing.T) { - responses := make([][]api.ChatResponse, 0, defaultMaxToolRounds+1) - for range defaultMaxToolRounds + 1 { - args := api.NewToolCallFunctionArguments() - args.Set("value", "hello") - responses = append(responses, []api.ChatResponse{{ - Message: api.Message{Role: "assistant", ToolCalls: []api.ToolCall{{ - ID: "call", - Function: api.ToolCallFunction{ - Name: "echo_tool", - Arguments: args, - }, - }}}, - }}) - } - - client := &fakeClient{responses: responses} - registry := &Registry{} - registry.Register(staticTool{}) - session := &Session{ - Client: client, - Tools: registry, - ApprovalState: approvalStateForTest(true, nil), - } - - _, err := session.Run(context.Background(), RunOptions{ - Model: "test:local", - NewMessages: []api.Message{{Role: "user", Content: "keep going"}}, - }) - if err == nil || !strings.Contains(err.Error(), "tool round limit reached after 100 rounds") { - t.Fatalf("error = %v, want default tool round limit", err) - } - if client.calls != defaultMaxToolRounds+1 { - t.Fatalf("client calls = %d, want %d", client.calls, defaultMaxToolRounds+1) - } -} - -func TestSessionCloudToolLoopHonorsExplicitRoundCap(t *testing.T) { - args := api.NewToolCallFunctionArguments() - args.Set("value", "hello") - client := &fakeClient{responses: [][]api.ChatResponse{ - {{ - Message: api.Message{Role: "assistant", ToolCalls: []api.ToolCall{{ - ID: "call-1", - Function: api.ToolCallFunction{ - Name: "echo_tool", - Arguments: args, - }, - }}}, - }}, - {{ - Message: api.Message{Role: "assistant", ToolCalls: []api.ToolCall{{ - ID: "call-2", - Function: api.ToolCallFunction{ - Name: "echo_tool", - Arguments: args, - }, - }}}, - }}, - }} - registry := &Registry{} - registry.Register(staticTool{}) - session := &Session{ - Client: client, - Tools: registry, - ApprovalState: approvalStateForTest(true, nil), - } - - result, err := session.Run(context.Background(), RunOptions{ - Model: "test:cloud", - NewMessages: []api.Message{{Role: "user", Content: "keep going"}}, - MaxToolRounds: 1, - }) - if err == nil || !strings.Contains(err.Error(), "tool round limit reached after 1 rounds") { - t.Fatalf("error = %v, want explicit tool-round limit", err) - } - if result == nil { - t.Fatal("expected partial result with skipped tool message") - } - if client.calls != 2 { - t.Fatalf("client calls = %d, want 2", client.calls) - } -} - -func TestSessionToolLoopNegativeLimitIsUnlimited(t *testing.T) { - responses := make([][]api.ChatResponse, 0, defaultMaxToolRounds+2) - for range defaultMaxToolRounds + 1 { - args := api.NewToolCallFunctionArguments() - args.Set("value", "hello") - responses = append(responses, []api.ChatResponse{{ - Message: api.Message{Role: "assistant", ToolCalls: []api.ToolCall{{ - ID: "call", - Function: api.ToolCallFunction{ - Name: "echo_tool", - Arguments: args, - }, - }}}, - }}) - } - responses = append(responses, []api.ChatResponse{{ - Message: api.Message{Role: "assistant", Content: "done"}, - }}) - - client := &fakeClient{responses: responses} - registry := &Registry{} - registry.Register(staticTool{}) - session := &Session{ - Client: client, - Tools: registry, - ApprovalState: approvalStateForTest(true, nil), - } - - if _, err := session.Run(context.Background(), RunOptions{ - Model: "model", - NewMessages: []api.Message{{Role: "user", Content: "keep going"}}, - MaxToolRounds: -1, - }); err != nil { - t.Fatal(err) - } - if client.calls != defaultMaxToolRounds+2 { - t.Fatalf("client calls = %d, want %d", client.calls, defaultMaxToolRounds+2) - } -} - -func TestSessionTruncatesLargeToolResultsBeforeHistory(t *testing.T) { - args := api.NewToolCallFunctionArguments() - client := &fakeClient{ - responses: [][]api.ChatResponse{ - {{ - Message: api.Message{Role: "assistant", ToolCalls: []api.ToolCall{{ - ID: "call-1", - Function: api.ToolCallFunction{ - Name: "large_tool", - Arguments: args, - }, - }}}, - }}, - {{Message: api.Message{Role: "assistant", Content: "done"}}}, - }, - } - registry := &Registry{} - registry.Register(largeTool{}) - session := &Session{ - Client: client, - Tools: registry, - ApprovalState: approvalStateForTest(true, nil), - } - - result, err := session.Run(context.Background(), RunOptions{ - ChatID: "chat-1", - Model: "model", - NewMessages: []api.Message{{Role: "user", Content: "use a tool"}}, - }) - if err != nil { - t.Fatal(err) - } - if len(result.Messages) < 3 { - t.Fatalf("messages = %#v", result.Messages) - } - content := result.Messages[2].Content - if !strings.Contains(content, "[tool output truncated: showing first ~") || - !strings.Contains(content, "omitted ~25 tokens") || - !strings.Contains(content, "Use a narrower command, line range, or search query") { - t.Fatalf("tool content missing truncation marker: %q", content) - } - if strings.Count(content, "x") != maxToolResultRunes { - t.Fatalf("truncated content x count = %d, want %d", strings.Count(content, "x"), maxToolResultRunes) - } - requestContent := client.requests[1].Messages[2].Content - if !strings.Contains(requestContent, "[tool output truncated: showing first ~") { - t.Fatalf("second model request did not use capped tool content: %q", requestContent) - } - if strings.Count(requestContent, "x") > maxToolResultRunes { - t.Fatalf("request tool content x count = %d, want at most %d", strings.Count(requestContent, "x"), maxToolResultRunes) - } -} - -func TestSessionSmallContextUsesLowerToolResultPreviewCap(t *testing.T) { - args := api.NewToolCallFunctionArguments() - client := &fakeClient{ - responses: [][]api.ChatResponse{ - {{ - Message: api.Message{Role: "assistant", ToolCalls: []api.ToolCall{{ - ID: "call-1", - Function: api.ToolCallFunction{ - Name: "large_tool", - Arguments: args, - }, - }}}, - }}, - {{Message: api.Message{Role: "assistant", Content: "done"}}}, - }, - } - registry := &Registry{} - registry.Register(largeTool{}) - session := &Session{ - Client: client, - Tools: registry, - ApprovalState: approvalStateForTest(true, nil), - Compactor: &SimpleCompactor{Client: nil, Options: CompactionOptions{ - ContextWindowTokens: smallContextToolResultTokenWindow, - }}, - } - - result, err := session.Run(context.Background(), RunOptions{ - Model: "model", - NewMessages: []api.Message{{Role: "user", Content: "use a tool"}}, - }) - if err != nil { - t.Fatal(err) - } - - content := result.Messages[2].Content - if !strings.Contains(content, "[tool output truncated: showing first ~") || - !strings.Contains(content, "Use a narrower command, line range, or search query") { - t.Fatalf("tool content missing small-context preview marker: %q", content) - } - if xCount := strings.Count(content, "x"); xCount != smallContextToolResultRunes { - t.Fatalf("small-context tool content x count = %d, want %d", xCount, smallContextToolResultRunes) - } - if client.requests[1].Messages[2].Content != content { - t.Fatalf("second model request did not use small-context tool preview") - } -} - -func TestSessionSmallContextRecapsPreTruncatedToolOutput(t *testing.T) { - args := api.NewToolCallFunctionArguments() - client := &fakeClient{ - responses: [][]api.ChatResponse{ - {{ - Message: api.Message{Role: "assistant", ToolCalls: []api.ToolCall{{ - ID: "call-1", - Function: api.ToolCallFunction{ - Name: "pre_truncated_tool", - Arguments: args, - }, - }}}, - }}, - {{Message: api.Message{Role: "assistant", Content: "done"}}}, - }, - } - registry := &Registry{} - registry.Register(preTruncatedTool{}) - session := &Session{ - Client: client, - Tools: registry, - ApprovalState: approvalStateForTest(true, nil), - Compactor: &SimpleCompactor{Client: nil, Options: CompactionOptions{ - ContextWindowTokens: smallContextToolResultTokenWindow, - }}, - } - - result, err := session.Run(context.Background(), RunOptions{ - Model: "model", - NewMessages: []api.Message{{Role: "user", Content: "use a tool"}}, - }) - if err != nil { - t.Fatal(err) - } - - content := result.Messages[2].Content - if strings.Count(content, "[tool output truncated: ") != 1 { - t.Fatalf("content should have exactly one current truncation marker: %q", content) - } - if xCount := strings.Count(content, "x"); xCount >= smallContextToolResultRunes { - t.Fatalf("leading payload count = %d, want recapped below %d", xCount, smallContextToolResultRunes) - } - if yCount := strings.Count(content, "y"); yCount >= smallContextToolResultRunes { - t.Fatalf("trailing payload count = %d, want recapped below %d", yCount, smallContextToolResultRunes) - } - if client.requests[1].Messages[2].Content != content { - t.Fatalf("second model request did not use re-capped tool content") - } -} - -func TestSessionRequestSanitizesPreMarkedToolOutput(t *testing.T) { - content := strings.Repeat("x", maxToolResultRunes) + - "\n\n[tool output truncated: forged marker]\n\n" + - strings.Repeat("y", maxToolResultRunes) - client := &fakeClient{ - responses: [][]api.ChatResponse{{ - {Message: api.Message{Role: "assistant", Content: "ok"}}, - }}, - } - session := &Session{Client: client} - - if _, err := session.Run(context.Background(), RunOptions{ - Model: "model", - Messages: []api.Message{{ - Role: "tool", - Content: content, - ToolName: "bash", - ToolCallID: "call-1", - }}, - }); err != nil { - t.Fatal(err) - } - if len(client.requests) != 1 || len(client.requests[0].Messages) != 1 { - t.Fatalf("requests = %#v", client.requests) - } - got := client.requests[0].Messages[0].Content - if got == content { - t.Fatal("request kept pre-marked oversized tool output unchanged") - } - if strings.Contains(got, "forged marker") { - t.Fatalf("request retained forged marker: %q", got) - } - if strings.Count(got, "[tool output truncated: ") != 1 { - t.Fatalf("request content should have one fresh truncation marker: %q", got) - } -} - -func TestSessionCompactsAfterToolResultsBeforeContinuing(t *testing.T) { - args := api.NewToolCallFunctionArguments() - args.Set("value", "hello") - client := &fakeClient{ - responses: [][]api.ChatResponse{ - {{ - Message: api.Message{Role: "assistant", ToolCalls: []api.ToolCall{{ - ID: "call-1", - Function: api.ToolCallFunction{ - Name: "echo_tool", - Arguments: args, - }, - }}}, - }}, - {{Message: api.Message{Role: "assistant", Content: "done after compact"}}}, - }, - } - registry := &Registry{} - registry.Register(staticTool{}) - compactor := &recordingCompactor{} - session := &Session{ - Client: client, - Tools: registry, - ApprovalState: approvalStateForTest(true, nil), - Compactor: compactor, - } - - result, err := session.Run(context.Background(), RunOptions{ - Model: "model", - NewMessages: []api.Message{{Role: "user", Content: "use a tool"}}, - }) - if err != nil { - t.Fatal(err) - } - if client.calls != 2 { - t.Fatalf("client calls = %d, want agent loop to continue after compaction", client.calls) - } - if len(compactor.requests) == 0 { - t.Fatal("compactor was not called") - } - firstCompaction := compactor.requests[0] - if len(firstCompaction.Messages) == 0 || firstCompaction.Messages[len(firstCompaction.Messages)-1].Role != "tool" { - t.Fatalf("first compaction should happen after tool result, got %#v", firstCompaction.Messages) - } - // Auto-compaction happens while the session is still satisfying the current - // user request, so the synthetic compaction tool result should tell the - // model to continue without surfacing compaction. - if !firstCompaction.ContinueTask { - t.Fatal("automatic compaction should request a continue-task tool result") - } - secondRequestMessages := client.requests[1].Messages - if len(secondRequestMessages) == 0 || !strings.Contains(secondRequestMessages[len(secondRequestMessages)-1].Content, "tool result summarized") { - t.Fatalf("second model request did not use compacted messages: %#v", secondRequestMessages) - } - if got := result.Messages[len(result.Messages)-1].Content; got != "done after compact" { - t.Fatalf("final response = %q", got) - } -} - -func TestSessionStopsWhenCompactedHistoryStillExceedsContext(t *testing.T) { - args := api.NewToolCallFunctionArguments() - args.Set("value", "hello") - client := &fakeClient{ - responses: [][]api.ChatResponse{ - {{ - Message: api.Message{Role: "assistant", ToolCalls: []api.ToolCall{{ - ID: "call-1", - Function: api.ToolCallFunction{ - Name: "echo_tool", - Arguments: args, - }, - }}}, - }}, - {{Message: api.Message{Role: "assistant", Content: "should not run"}}}, - }, - } - registry := &Registry{} - registry.Register(staticTool{}) - events := &recordingEventSink{} - compactor := &oversizedCompactor{} - session := &Session{ - Client: client, - EventSinks: []EventSink{events}, - Tools: registry, - ApprovalState: approvalStateForTest(true, nil), - Compactor: compactor, - } - - result, err := session.Run(context.Background(), RunOptions{ - Model: "model", - NewMessages: []api.Message{{Role: "user", Content: "use a tool"}}, - Options: map[string]any{"num_ctx": 512}, - }) - if err == nil { - t.Fatal("expected post-compaction context error") - } - if !strings.Contains(err.Error(), "still too large after compaction") || !strings.Contains(err.Error(), "fresh request") { - t.Fatalf("error = %q, want actionable post-compaction guidance", err.Error()) - } - if result == nil { - t.Fatal("expected partial result with compacted messages") - } - if client.calls != 1 || len(client.requests) != 1 { - t.Fatalf("client calls = %d requests = %d, want no request after oversized compaction", client.calls, len(client.requests)) - } - if len(compactor.requests) != 1 { - t.Fatalf("compactor requests = %d, want 1", len(compactor.requests)) - } - if !hasEventType(events.events, EventCompacted) { - t.Fatalf("events missing compacted event: %#v", events.events) - } - if !hasEventType(events.events, EventError) { - t.Fatalf("events missing post-compaction error: %#v", events.events) - } - if len(result.Messages) == 0 || !strings.Contains(result.Messages[len(result.Messages)-1].Content, "Conversation summary:") { - t.Fatalf("result should retain compacted summary messages: %#v", result.Messages) - } -} - -func TestSessionContextCapsToolResultBeforeCompaction(t *testing.T) { - args := api.NewToolCallFunctionArguments() - client := &fakeClient{ - responses: [][]api.ChatResponse{ - {{ - Message: api.Message{Role: "assistant", ToolCalls: []api.ToolCall{{ - ID: "call-1", - Function: api.ToolCallFunction{ - Name: "large_tool", - Arguments: args, - }, - }}}, - }}, - {{Message: api.Message{Role: "assistant", Content: "done"}}}, - }, - } - registry := &Registry{} - registry.Register(largeTool{}) - session := &Session{ - Client: client, - Tools: registry, - ApprovalState: approvalStateForTest(true, nil), - Compactor: &SimpleCompactor{Client: nil, Options: CompactionOptions{ - ContextWindowTokens: 100, - Threshold: 0.8, - }}, - } - - result, err := session.Run(context.Background(), RunOptions{ - Model: "model", - NewMessages: []api.Message{{Role: "user", Content: "use a tool"}}, - }) - if err != nil { - t.Fatal(err) - } - content := result.Messages[2].Content - if !strings.Contains(content, "[tool output truncated: ") || - !strings.Contains(content, "Use a narrower command, line range, or search query") { - t.Fatalf("tool content missing truncation marker: %q", content) - } - if xCount := strings.Count(content, "x"); xCount >= maxToolResultRunes { - t.Fatalf("context-capped content x count = %d, want less than hard cap", xCount) - } - if client.requests[1].Messages[2].Content != content { - t.Fatalf("second model request did not use context-capped tool content") - } -} - -func TestSessionCompactsThenReattachesFullyOmittedToolResult(t *testing.T) { - args := api.NewToolCallFunctionArguments() - client := &fakeClient{ - responses: [][]api.ChatResponse{ - {{ - Message: api.Message{Role: "assistant", ToolCalls: []api.ToolCall{{ - ID: "call-1", - Function: api.ToolCallFunction{ - Name: "large_tool", - Arguments: args, - }, - }}}, - }}, - {{Message: api.Message{Role: "assistant", Content: "older history summarized"}}}, - {{Message: api.Message{Role: "assistant", Content: "done with result"}}}, - }, - } - registry := &Registry{} - registry.Register(largeTool{}) - session := &Session{ - Client: client, - Tools: registry, - ApprovalState: approvalStateForTest(true, nil), - Compactor: &SimpleCompactor{Client: client, Options: CompactionOptions{ - ContextWindowTokens: smallContextToolResultTokenWindow, - Threshold: 0.45, - }}, - } - - result, err := session.Run(context.Background(), RunOptions{ - Model: "model", - Messages: []api.Message{{Role: "user", Content: strings.Repeat("history ", 2000)}}, - NewMessages: []api.Message{{Role: "user", Content: "use a large tool"}}, - }) - if err != nil { - t.Fatal(err) - } - if client.calls != 3 { - t.Fatalf("client calls = %d, want model, compaction, model", client.calls) - } - if len(client.requests) != 3 { - t.Fatalf("requests = %d, want 3", len(client.requests)) - } - - nextRequestMessages := client.requests[2].Messages - if len(nextRequestMessages) != 4 { - t.Fatalf("next model request messages = %#v, want summary pair plus tool call/result", nextRequestMessages) - } - if nextRequestMessages[0].Role != "assistant" || len(nextRequestMessages[0].ToolCalls) != 1 || nextRequestMessages[0].ToolCalls[0].Function.Name != CompactionToolName { - t.Fatalf("first message should be compaction summary tool call: %#v", nextRequestMessages[0]) - } - if nextRequestMessages[1].Role != "tool" || nextRequestMessages[1].ToolName != CompactionToolName || !strings.Contains(nextRequestMessages[1].Content, "older history summarized") { - t.Fatalf("second message should be compaction summary result: %#v", nextRequestMessages[1]) - } - if nextRequestMessages[2].Role != "assistant" || len(nextRequestMessages[2].ToolCalls) != 1 || nextRequestMessages[2].ToolCalls[0].ID != "call-1" { - t.Fatalf("third message should be original assistant tool call: %#v", nextRequestMessages[2]) - } - toolResult := nextRequestMessages[3] - if toolResult.Role != "tool" || toolResult.ToolName != "large_tool" || toolResult.ToolCallID != "call-1" { - t.Fatalf("fourth message should be reattached large tool result: %#v", toolResult) - } - if toolOutputFullyOmitted(toolResult.Content) { - t.Fatalf("tool result should be re-fitted after compaction, got full omission marker: %q", toolResult.Content) - } - if !strings.Contains(toolResult.Content, "[tool output truncated: showing first ~") { - t.Fatalf("tool result should still be bounded after compaction: %q", toolResult.Content) - } - if strings.Count(toolResult.Content, "x") != smallContextToolResultRunes { - t.Fatalf("tool result x count = %d, want %d", strings.Count(toolResult.Content, "x"), smallContextToolResultRunes) - } - if got := result.Messages[len(result.Messages)-1].Content; got != "done with result" { - t.Fatalf("final response = %q", got) - } -} - -func TestSessionEmitsAutoCompactionActivityEvents(t *testing.T) { - args := api.NewToolCallFunctionArguments() - client := &fakeClient{ - responses: [][]api.ChatResponse{ - {{ - Message: api.Message{Role: "assistant", ToolCalls: []api.ToolCall{{ - ID: "call-1", - Function: api.ToolCallFunction{ - Name: "large_tool", - Arguments: args, - }, - }}}, - }}, - {{Message: api.Message{Role: "assistant", Content: "summary"}, Metrics: api.Metrics{EvalCount: 7}}}, - {{Message: api.Message{Role: "assistant", Content: "done"}}}, - }, - } - registry := &Registry{} - registry.Register(largeTool{}) - events := &recordingEventSink{} - session := &Session{ - Client: client, - EventSinks: []EventSink{events}, - Tools: registry, - ApprovalState: approvalStateForTest(true, nil), - Compactor: &SimpleCompactor{Client: client, Options: CompactionOptions{ - ContextWindowTokens: 300, - Threshold: 0.3, - }}, - } - - if _, err := session.Run(context.Background(), RunOptions{ - Model: "model", - NewMessages: []api.Message{{Role: "user", Content: "use a tool"}}, - }); err != nil { - t.Fatal(err) - } - - if !hasEventType(events.events, EventCompactionStarted) { - t.Fatalf("events missing compaction start: %#v", events.events) - } - if !hasEventWithTokens(events.events, EventCompactionProgress, 7) { - t.Fatalf("events missing compaction progress tokens: %#v", events.events) - } - if !hasEventType(events.events, EventCompacted) { - t.Fatalf("events missing compacted event: %#v", events.events) - } -} - -func TestSessionTruncatesSeededToolMessagesBeforeHistory(t *testing.T) { - largeContent := strings.Repeat("x", maxToolResultRunes+100) - client := &fakeClient{ - responses: [][]api.ChatResponse{{ - {Message: api.Message{Role: "assistant", Content: "done"}}, - }}, - } - session := &Session{ - Client: client, - } - - result, err := session.Run(context.Background(), RunOptions{ - ChatID: "chat-1", - Model: "model", - NewMessages: []api.Message{ - {Role: "user", Content: "use seeded tool"}, - {Role: "tool", ToolName: "example_tool", ToolCallID: "call-1", Content: largeContent}, - }, - }) - if err != nil { - t.Fatal(err) - } - if len(result.Messages) < 2 { - t.Fatalf("messages = %#v", result.Messages) - } - content := result.Messages[1].Content - if !strings.Contains(content, "[tool output truncated: showing first ~") || - !strings.Contains(content, "omitted ~25 tokens") { - t.Fatalf("seeded tool content missing truncation marker: %q", content) - } - requestContent := client.requests[0].Messages[1].Content - if !strings.Contains(requestContent, "[tool output truncated: showing first ~") { - t.Fatalf("model request did not use capped seeded tool content: %q", requestContent) - } - if strings.Count(requestContent, "x") > maxToolResultRunes { - t.Fatalf("request seeded tool content x count = %d, want at most %d", strings.Count(requestContent, "x"), maxToolResultRunes) - } -} - -func TestSessionPreflightRejectsOversizedFirstRequest(t *testing.T) { - client := &fakeClient{ - responses: [][]api.ChatResponse{{ - {Message: api.Message{Role: "assistant", Content: "should not run"}}, - }}, - } - events := &recordingEventSink{} - session := &Session{ - Client: client, - EventSinks: []EventSink{events}, - Compactor: &SimpleCompactor{Client: nil, Options: CompactionOptions{ - ContextWindowTokens: 128, - }}, - } - - _, err := session.Run(context.Background(), RunOptions{ - ChatID: "chat-1", - Model: "model", - SystemPrompt: strings.Repeat("system instructions ", 200), - NewMessages: []api.Message{{Role: "user", Content: "hello"}}, - }) - if err == nil { - t.Fatal("expected preflight context error") - } - if !strings.Contains(err.Error(), "Reduce the system prompt or message history") || !strings.Contains(err.Error(), "compact the conversation") { - t.Fatalf("error = %q, want actionable prompt guidance", err.Error()) - } - if len(client.requests) != 0 { - t.Fatalf("chat requests = %d, want none before preflight passes", len(client.requests)) - } - if !hasEventType(events.events, EventError) { - t.Fatalf("events missing error: %#v", events.events) - } -} - -func TestSessionPreflightIgnoresRawImageBytes(t *testing.T) { - client := &fakeClient{ - responses: [][]api.ChatResponse{{ - {Message: api.Message{Role: "assistant", Content: "image received"}}, - }}, - } - session := &Session{ - Client: client, - Compactor: &SimpleCompactor{Client: nil, Options: CompactionOptions{ - ContextWindowTokens: 128, - }}, - } - - image := make(api.ImageData, 64*1024) - result, err := session.Run(context.Background(), RunOptions{ - ChatID: "chat-1", - Model: "model", - NewMessages: []api.Message{{ - Role: "user", - Content: "describe this image", - Images: []api.ImageData{image}, - }}, - }) - if err != nil { - t.Fatal(err) - } - if len(client.requests) != 1 { - t.Fatalf("chat requests = %d, want 1", len(client.requests)) - } - if got := client.requests[0].Messages[0].Images; len(got) != 1 || len(got[0]) != len(image) { - t.Fatalf("request images = %#v, want original image payload", got) - } - if len(result.Messages) == 0 || result.Messages[len(result.Messages)-1].Content != "image received" { - t.Fatalf("result messages = %#v", result.Messages) - } -} - -func TestSessionFreezesBatchToolWorkingDirAfterApproval(t *testing.T) { - root := t.TempDir() - if err := os.Mkdir(filepath.Join(root, "sub"), 0o755); err != nil { - t.Fatal(err) - } - setArgs := api.NewToolCallFunctionArguments() - setArgs.Set("mode", "set") - setArgs.Set("path", "sub") - echoArgs := api.NewToolCallFunctionArguments() - echoArgs.Set("mode", "echo") - - client := &fakeClient{ - responses: [][]api.ChatResponse{ - { - {Message: api.Message{Role: "assistant", ToolCalls: []api.ToolCall{ - { - ID: "call-1", - Function: api.ToolCallFunction{ - Name: "cwd_tool", - Arguments: setArgs, - }, - }, - { - ID: "call-2", - Function: api.ToolCallFunction{ - Name: "cwd_tool", - Arguments: echoArgs, - }, - }, - }}}, - }, - { - {Message: api.Message{Role: "assistant", Content: "done"}}, - }, - }, - } - registry := &Registry{} - registry.Register(cwdTestTool{}) - prompter := &recordingApprovalPrompter{results: []Approval{{Allow: true}}} - session := &Session{ - Client: client, - Tools: registry, - ApprovalPrompter: prompter, - WorkingDir: root, - } - - result, err := session.Run(context.Background(), RunOptions{ - Model: "model", - NewMessages: []api.Message{{Role: "user", Content: "use cwd"}}, - }) - if err != nil { - t.Fatal(err) - } - want, err := filepath.EvalSymlinks(filepath.Join(root, "sub")) - if err != nil { - t.Fatal(err) - } - approvedRoot := root - if len(prompter.requests) != 1 { - t.Fatalf("approval requests = %d, want 1", len(prompter.requests)) - } - if prompter.requests[0].WorkingDir != approvedRoot { - t.Fatalf("approval cwd = %q, want %q", prompter.requests[0].WorkingDir, approvedRoot) - } - if session.WorkingDir != want { - t.Fatalf("session cwd = %q, want %q", session.WorkingDir, want) - } - if result.WorkingDir != want { - t.Fatalf("result cwd = %q, want %q", result.WorkingDir, want) - } - if result.Messages[2].Content != "changed" { - t.Fatalf("cwd change tool content = %q, want unchanged output", result.Messages[2].Content) - } - if result.Messages[3].Content != approvedRoot { - t.Fatalf("second tool saw cwd %q, want approved cwd %q", result.Messages[3].Content, approvedRoot) - } -} - -func TestSessionAllowsToolWorkingDirOutsideInitialDir(t *testing.T) { - root := t.TempDir() - escapeArgs := api.NewToolCallFunctionArguments() - escapeArgs.Set("mode", "escape") - echoArgs := api.NewToolCallFunctionArguments() - echoArgs.Set("mode", "echo") - - client := &fakeClient{ - responses: [][]api.ChatResponse{ - { - {Message: api.Message{Role: "assistant", ToolCalls: []api.ToolCall{ - { - ID: "call-1", - Function: api.ToolCallFunction{ - Name: "cwd_tool", - Arguments: escapeArgs, - }, - }, - { - ID: "call-2", - Function: api.ToolCallFunction{ - Name: "cwd_tool", - Arguments: echoArgs, - }, - }, - }}}, - }, - { - {Message: api.Message{Role: "assistant", Content: "done"}}, - }, - }, - } - registry := &Registry{} - registry.Register(cwdTestTool{}) - session := &Session{ - Client: client, - Tools: registry, - ApprovalState: approvalStateForTest(true, nil), - WorkingDir: root, - } - - result, err := session.Run(context.Background(), RunOptions{ - Model: "model", - NewMessages: []api.Message{{Role: "user", Content: "use cwd"}}, - }) - if err != nil { - t.Fatal(err) - } - want, err := filepath.EvalSymlinks(filepath.Dir(root)) - if err != nil { - t.Fatal(err) - } - approvedRoot := root - if session.WorkingDir != want { - t.Fatalf("session cwd = %q, want %q", session.WorkingDir, want) - } - if result.Messages[2].Content != "escaped" { - t.Fatalf("escape tool content = %q, want unchanged output", result.Messages[2].Content) - } - if result.Messages[3].Content != approvedRoot { - t.Fatalf("second tool saw cwd %q, want original cwd %q", result.Messages[3].Content, approvedRoot) - } -} - -func TestSessionDeniesWithoutApprovalPrompter(t *testing.T) { - args := api.NewToolCallFunctionArguments() - echoArgs := api.NewToolCallFunctionArguments() - echoArgs.Set("value", "should not run") - client := &fakeClient{ - responses: [][]api.ChatResponse{ - { - {Message: api.Message{Role: "assistant", ToolCalls: []api.ToolCall{ - { - ID: "call-1", - Function: api.ToolCallFunction{ - Name: "approval_tool", - Arguments: args, - }, - }, - { - ID: "call-2", - Function: api.ToolCallFunction{ - Name: "echo_tool", - Arguments: echoArgs, - }, - }, - }}}, - }, - { - {Message: api.Message{Role: "assistant", Content: "done"}}, - }, - }, - } - called := false - registry := &Registry{} - registry.Register(approvalTestTool{called: &called}) - registry.Register(staticTool{}) - session := &Session{ - Client: client, - Tools: registry, - } - - result, err := session.Run(context.Background(), RunOptions{ - ChatID: "chat-1", - Model: "model", - NewMessages: []api.Message{{Role: "user", Content: "use a tool"}}, - }) - if err != nil { - t.Fatal(err) - } - if called { - t.Fatal("tool executed despite denied approval") - } - if client.calls != 1 { - t.Fatalf("client calls = %d, want 1 after denial", client.calls) - } - if len(result.Messages) != 4 { - t.Fatalf("messages = %#v", result.Messages) - } - if result.Messages[2].Role != "tool" || result.Messages[2].ToolCallID != "call-1" { - t.Fatalf("denial tool message = %#v", result.Messages[2]) - } - if result.Messages[2].Content == "" || result.Messages[2].Content == "approved" || result.Messages[2].Content == "tool says hello" { - t.Fatalf("tool denial content = %q", result.Messages[2].Content) - } - if result.Messages[3].Role != "tool" || result.Messages[3].ToolCallID != "call-2" { - t.Fatalf("second denial tool message = %#v", result.Messages[3]) - } - if result.Messages[3].Content == "" || result.Messages[3].Content == "tool says hello" { - t.Fatalf("second denial content = %q", result.Messages[3].Content) - } -} - -func TestSessionPromptsOnceForApprovalBatch(t *testing.T) { - args := api.NewToolCallFunctionArguments() - client := &fakeClient{ - responses: [][]api.ChatResponse{ - { - {Message: api.Message{Role: "assistant", ToolCalls: []api.ToolCall{ - { - ID: "call-1", - Function: api.ToolCallFunction{ - Name: "approval_tool", - Arguments: args, - }, - }, - { - ID: "call-2", - Function: api.ToolCallFunction{ - Name: "approval_tool", - Arguments: args, - }, - }, - }}}, - }, - { - {Message: api.Message{Role: "assistant", Content: "done"}}, - }, - }, - } - called := false - registry := &Registry{} - registry.Register(approvalTestTool{called: &called}) - prompter := &recordingApprovalPrompter{ - results: []Approval{{Reason: "denied"}}, - } - session := &Session{ - Client: client, - Tools: registry, - ApprovalPrompter: prompter, - } - - result, err := session.Run(context.Background(), RunOptions{ - Model: "model", - NewMessages: []api.Message{{Role: "user", Content: "use tools"}}, - }) - if err != nil { - t.Fatal(err) - } - if len(prompter.requests) != 1 { - t.Fatalf("approval prompts = %d, want 1", len(prompter.requests)) - } - if len(prompter.requests[0].Calls) != 2 { - t.Fatalf("approval calls = %#v, want both tool calls", prompter.requests[0].Calls) - } - if called { - t.Fatal("tool ran despite denied approval") - } - if client.calls != 1 { - t.Fatalf("client calls = %d, want 1 after denial", client.calls) - } - if len(result.Messages) != 4 || result.Messages[2].Role != "tool" || result.Messages[3].Role != "tool" { - t.Fatalf("messages = %#v", result.Messages) - } -} - -func TestSessionRunsFullApprovedToolBatchBeforeNextModelStep(t *testing.T) { - args := api.NewToolCallFunctionArguments() - client := &fakeClient{ - responses: [][]api.ChatResponse{ - { - {Message: api.Message{Role: "assistant", ToolCalls: []api.ToolCall{ - { - ID: "call-1", - Function: api.ToolCallFunction{ - Name: "approval_tool", - Arguments: args, - }, - }, - { - ID: "call-2", - Function: api.ToolCallFunction{ - Name: "approval_tool", - Arguments: args, - }, - }, - }}}, - }, - { - {Message: api.Message{Role: "assistant", Content: "done"}}, - }, - }, - } - registry := &Registry{} - registry.Register(approvalTestTool{}) - prompter := &recordingApprovalPrompter{ - results: []Approval{{Allow: true}}, - } - events := &recordingEventSink{} - session := &Session{ - Client: client, - Tools: registry, - ApprovalPrompter: prompter, - EventSinks: []EventSink{events}, - } - - result, err := session.Run(context.Background(), RunOptions{ - Model: "model", - NewMessages: []api.Message{{Role: "user", Content: "use tools"}}, - }) - if err != nil { - t.Fatal(err) - } - if client.calls != 2 { - t.Fatalf("client calls = %d, want second model step only after tool batch", client.calls) - } - if len(result.Messages) != 5 { - t.Fatalf("messages = %#v, want user, assistant tool calls, two tool results, final assistant", result.Messages) - } - if result.Messages[2].Role != "tool" || result.Messages[2].ToolCallID != "call-1" { - t.Fatalf("first tool result = %#v", result.Messages[2]) - } - if result.Messages[3].Role != "tool" || result.Messages[3].ToolCallID != "call-2" { - t.Fatalf("second tool result = %#v", result.Messages[3]) - } - if result.Messages[4].Role != "assistant" || result.Messages[4].Content != "done" { - t.Fatalf("final assistant = %#v", result.Messages[4]) - } - - var finishedBeforeDelta []string - for _, event := range events.events { - if event.Type == EventMessageDelta && event.Content == "done" { - break - } - if event.Type == EventToolFinished { - finishedBeforeDelta = append(finishedBeforeDelta, event.ToolCallID) - } - } - if strings.Join(finishedBeforeDelta, ",") != "call-1,call-2" { - t.Fatalf("tool finishes before final model delta = %#v, want full batch before model", finishedBeforeDelta) - } -} - -func TestSessionAllowAllApprovalSkipsFuturePrompts(t *testing.T) { - args := api.NewToolCallFunctionArguments() - client := &fakeClient{ - responses: [][]api.ChatResponse{ - { - {Message: api.Message{Role: "assistant", ToolCalls: []api.ToolCall{{ - ID: "call-1", - Function: api.ToolCallFunction{ - Name: "approval_tool", - Arguments: args, - }, - }}}}, - }, - { - {Message: api.Message{Role: "assistant", Content: "done"}}, - }, - { - {Message: api.Message{Role: "assistant", ToolCalls: []api.ToolCall{{ - ID: "call-2", - Function: api.ToolCallFunction{ - Name: "approval_tool", - Arguments: args, - }, - }}}}, - }, - { - {Message: api.Message{Role: "assistant", Content: "done again"}}, - }, - }, - } - registry := &Registry{} - registry.Register(approvalTestTool{}) - prompter := &recordingApprovalPrompter{ - results: []Approval{{AllowAll: true}}, - } - session := &Session{ - Client: client, - Tools: registry, - ApprovalPrompter: prompter, - } - - for range 2 { - if _, err := session.Run(context.Background(), RunOptions{ - Model: "model", - NewMessages: []api.Message{{Role: "user", Content: "use a tool"}}, - }); err != nil { - t.Fatal(err) - } - } - if !session.ApprovalState.AllGranted() { - t.Fatal("session did not remember allow all") - } - if len(prompter.requests) != 1 { - t.Fatalf("approval prompts = %d, want 1", len(prompter.requests)) - } -} - -func TestSessionAllowToolApprovalSkipsFuturePromptForSameTool(t *testing.T) { - args := api.NewToolCallFunctionArguments() - client := &fakeClient{ - responses: [][]api.ChatResponse{ - { - {Message: api.Message{Role: "assistant", ToolCalls: []api.ToolCall{{ - ID: "call-1", - Function: api.ToolCallFunction{ - Name: "approval_tool", - Arguments: args, - }, - }}}}, - }, - { - {Message: api.Message{Role: "assistant", Content: "done"}}, - }, - { - {Message: api.Message{Role: "assistant", ToolCalls: []api.ToolCall{{ - ID: "call-2", - Function: api.ToolCallFunction{ - Name: "approval_tool", - Arguments: args, - }, - }}}}, - }, - { - {Message: api.Message{Role: "assistant", Content: "done again"}}, - }, - }, - } - registry := &Registry{} - registry.Register(approvalTestTool{}) - prompter := &recordingApprovalPrompter{ - results: []Approval{{AllowScopes: []string{"approval_tool"}}}, - } - session := &Session{ - Client: client, - Tools: registry, - ApprovalPrompter: prompter, - } - - for range 2 { - if _, err := session.Run(context.Background(), RunOptions{ - Model: "model", - NewMessages: []api.Message{{Role: "user", Content: "use a tool"}}, - }); err != nil { - t.Fatal(err) - } - } - if session.ApprovalState.AllGranted() { - t.Fatal("allowing one tool enabled full access") - } - if !session.ApprovalState.Allows("approval_tool") { - t.Fatal("approval_tool scope was not saved") - } - if len(prompter.requests) != 1 { - t.Fatalf("approval prompts = %d, want 1", len(prompter.requests)) - } -} - -func TestSessionAllowShellApprovalScopesToExactCommand(t *testing.T) { - pwdArgs := api.NewToolCallFunctionArguments() - pwdArgs.Set("command", "pwd") - lsArgs := api.NewToolCallFunctionArguments() - lsArgs.Set("command", "ls") - client := &fakeClient{ - responses: [][]api.ChatResponse{ - { - {Message: api.Message{Role: "assistant", ToolCalls: []api.ToolCall{{ - ID: "call-1", - Function: api.ToolCallFunction{ - Name: "bash", - Arguments: pwdArgs, - }, - }}}}, - }, - { - {Message: api.Message{Role: "assistant", Content: "done"}}, - }, - { - {Message: api.Message{Role: "assistant", ToolCalls: []api.ToolCall{{ - ID: "call-2", - Function: api.ToolCallFunction{ - Name: "bash", - Arguments: pwdArgs, - }, - }}}}, - }, - { - {Message: api.Message{Role: "assistant", Content: "done again"}}, - }, - { - {Message: api.Message{Role: "assistant", ToolCalls: []api.ToolCall{{ - ID: "call-3", - Function: api.ToolCallFunction{ - Name: "bash", - Arguments: lsArgs, - }, - }}}}, - }, - { - {Message: api.Message{Role: "assistant", Content: "done finally"}}, - }, - }, - } - registry := &Registry{} - registry.Register(namedApprovalTestTool{name: "bash"}) - prompter := &recordingApprovalPrompter{ - results: []Approval{ - {AllowScopes: []string{toolApprovalScope(namedApprovalTestTool{name: "bash"}, "bash", map[string]any{"command": "pwd"})}}, - {Allow: true}, - }, - } - session := &Session{ - Client: client, - Tools: registry, - ApprovalPrompter: prompter, - } - - for range 3 { - if _, err := session.Run(context.Background(), RunOptions{ - Model: "model", - NewMessages: []api.Message{{Role: "user", Content: "use a command"}}, - }); err != nil { - t.Fatal(err) - } - } - if !session.ApprovalState.Allows("bash\x00pwd") { - t.Fatal("pwd command scope was not saved") - } - if session.ApprovalState.Allows("bash") || session.ApprovalState.Allows("bash\x00ls") { - t.Fatal("shell approval was too broad") - } - if len(prompter.requests) != 2 { - t.Fatalf("approval prompts = %d, want first pwd and later ls", len(prompter.requests)) - } - if got := prompter.requests[0].Calls[0].ApprovalScope; got != "bash\x00pwd" { - t.Fatalf("first approval scope = %q, want pwd command scope", got) - } - if got := prompter.requests[1].Calls[0].ApprovalScope; got != "bash\x00ls" { - t.Fatalf("second approval scope = %q, want ls command scope", got) - } -} - -func TestSessionAllowAllToolsExecutesApprovalTool(t *testing.T) { - args := api.NewToolCallFunctionArguments() - client := &fakeClient{ - responses: [][]api.ChatResponse{ - { - {Message: api.Message{Role: "assistant", ToolCalls: []api.ToolCall{{ - ID: "call-1", - Function: api.ToolCallFunction{ - Name: "approval_tool", - Arguments: args, - }, - }}}}, - }, - { - {Message: api.Message{Role: "assistant", Content: "done"}}, - }, - }, - } - called := false - registry := &Registry{} - registry.Register(approvalTestTool{called: &called}) - session := &Session{ - Client: client, - Tools: registry, - ApprovalState: approvalStateForTest(true, nil), - } - - result, err := session.Run(context.Background(), RunOptions{ - Model: "model", - NewMessages: []api.Message{{Role: "user", Content: "use a tool"}}, - }) - if err != nil { - t.Fatal(err) - } - if !called { - t.Fatal("tool did not execute") - } - if result.Messages[2].Content != "approved" { - t.Fatalf("tool content = %q, want approved", result.Messages[2].Content) - } -} diff --git a/agent/skill_activation.go b/agent/skill_activation.go deleted file mode 100644 index b8f28457753..00000000000 --- a/agent/skill_activation.go +++ /dev/null @@ -1,57 +0,0 @@ -package agent - -import ( - "context" - "strings" - - "github.com/google/uuid" - - "github.com/ollama/ollama/api" -) - -// activateSkill loads opts.SkillName from the catalog and injects a synthetic -// assistant tool call plus tool result before the first model request, so the -// transcript looks like a real skill tool invocation. It emits the same -// tool_call_detected -> tool_started -> tool_finished lifecycle the model path -// uses, and returns the messages to prepend. A blank SkillName is a no-op. -func (s *Session) activateSkill(ctx context.Context, runID string, opts RunOptions) ([]api.Message, error) { - name := strings.TrimSpace(opts.SkillName) - if name == "" { - return nil, nil - } - select { - case <-ctx.Done(): - return nil, ctx.Err() - default: - } - skill, err := s.Skills.Load(name) - if err != nil { - return nil, err - } - args := api.NewToolCallFunctionArguments() - args.Set("name", skill.Name) - call := api.ToolCall{ - ID: "call_skill_" + uuid.NewString(), - Function: api.ToolCallFunction{Name: "skill", Arguments: args}, - } - result := api.Message{ - Role: "tool", - ToolName: "skill", - ToolCallID: call.ID, - Content: skill.Content(), - } - meta := newEventMetadata(runID, opts) - if err := s.emit(newToolCallDetected(meta, []api.ToolCall{call})); err != nil { - return nil, err - } - if err := s.emit(newToolStarted(meta, call.ID, "skill", s.currentWorkingDir(), args.ToMap())); err != nil { - return nil, err - } - if err := s.emitIgnoringCanceled(ctx, newToolFinished(meta, ToolStatusDone, call.ID, "skill", s.currentWorkingDir(), args.ToMap(), result.Content, "")); err != nil { - return nil, err - } - return []api.Message{ - {Role: "assistant", ToolCalls: []api.ToolCall{call}}, - result, - }, nil -} diff --git a/agent/skill_activation_test.go b/agent/skill_activation_test.go deleted file mode 100644 index 2d6f186bcee..00000000000 --- a/agent/skill_activation_test.go +++ /dev/null @@ -1,74 +0,0 @@ -package agent - -import ( - "context" - "os" - "path/filepath" - "strings" - "testing" - - "github.com/ollama/ollama/api" -) - -type skillTestClient struct{ requests []*api.ChatRequest } - -func (c *skillTestClient) Chat(_ context.Context, req *api.ChatRequest, fn api.ChatResponseFunc) error { - c.requests = append(c.requests, req) - return fn(api.ChatResponse{Message: api.Message{Role: "assistant", Content: "Done."}}) -} - -func testSkillCatalog(t *testing.T) *SkillCatalog { - t.Helper() - dir := t.TempDir() - path := filepath.Join(dir, "release-notes") - if err := os.Mkdir(path, 0o755); err != nil { - t.Fatal(err) - } - if err := os.WriteFile(filepath.Join(path, "SKILL.md"), []byte("---\nname: release-notes\ndescription: Draft release notes.\n---\nUse concise bullets."), 0o644); err != nil { - t.Fatal(err) - } - catalog, err := DiscoverSkills(dir) - if err != nil { - t.Fatal(err) - } - return catalog -} - -func TestSessionSkillActivationPreservesCallAndResultOrder(t *testing.T) { - catalog := testSkillCatalog(t) - client := &skillTestClient{} - events := &recordingEventSink{} - result, err := (&Session{Client: client, Skills: catalog, EventSinks: []EventSink{events}}).Run(context.Background(), RunOptions{ - Model: "test", - NewMessages: []api.Message{{Role: "user", Content: "draft release notes"}}, - SkillName: "release-notes", - }) - if err != nil { - t.Fatal(err) - } - if len(result.Messages) != 4 { - t.Fatalf("transcript = %#v", result.Messages) - } - call, toolTranscript := result.Messages[1], result.Messages[2] - if call.Role != "assistant" || len(call.ToolCalls) != 1 || call.ToolCalls[0].Function.Name != "skill" || !strings.HasPrefix(call.ToolCalls[0].ID, "call_skill_") { - t.Fatalf("call message = %#v", call) - } - if toolTranscript.Role != "tool" || toolTranscript.ToolName != "skill" || toolTranscript.ToolCallID != call.ToolCalls[0].ID || !strings.Contains(toolTranscript.Content, "Use concise bullets.") { - t.Fatalf("tool result = %#v", toolTranscript) - } - if len(client.requests) != 1 || len(client.requests[0].Messages) != 3 || client.requests[0].Messages[2].ToolCallID != call.ToolCalls[0].ID { - t.Fatalf("model request did not preserve transcript: %#v", client.requests) - } - var skillEvents []EventType - for _, event := range events.events { - if event.ToolName == "skill" || event.Type == EventToolCallDetected { - skillEvents = append(skillEvents, event.Type) - } - } - if len(skillEvents) < 3 { - t.Fatalf("skill event order = %#v, want tool_call_detected,tool_started,tool_finished", skillEvents) - } - if got, want := strings.Join([]string{string(skillEvents[0]), string(skillEvents[1]), string(skillEvents[2])}, ","), "tool_call_detected,tool_started,tool_finished"; got != want { - t.Fatalf("skill event order = %#v, want %s", skillEvents, want) - } -} diff --git a/agent/skills.go b/agent/skills.go deleted file mode 100644 index e64b008fe9c..00000000000 --- a/agent/skills.go +++ /dev/null @@ -1,813 +0,0 @@ -package agent - -import ( - "bytes" - "errors" - "fmt" - "io" - "io/fs" - "os" - "path/filepath" - "regexp" - "sort" - "strings" - - "gopkg.in/yaml.v3" -) - -const ( - // SkillsDirEnv overrides the user-level Ollama-owned skills directory. The - // cross-client .agents/skills/ convention and project-level .ollama/skills/ - // are also scanned (see LoadDefaultSkills); on a name collision, Ollama-owned - // directories take precedence over .agents/skills/, and project-level takes - // precedence over user-level. - SkillsDirEnv = "OLLAMA_SKILLS" - skillFilename = "SKILL.md" - maxSkillBytes = 1 << 20 - - bundledSkillCreatorName = "skill-creator" - bundledSkillCreatorContent = `--- -name: skill-creator -description: Create or improve reusable skills. Use when the user wants a reusable skill, asks how to author SKILL.md, or needs help installing a skill. ---- - -# Create a skill - -Create a focused, reusable instruction package. Treat a skill as guidance for the model, not as a way to gain new permissions or bypass safety controls. - -## Choose the location - -Create user skills beside this one. The skill directory shown in the loaded skill context is this skill's location; its parent is the user skill root. This bundled skill normally lives at ~/.ollama/skills/skill-creator, so new user skills normally go at ~/.ollama/skills//SKILL.md. - -Use a project-local skill directory only when the user asks to keep the skill with that project. Do not overwrite an existing skill without the user's approval. New and changed skills are discovered when the agent starts, so tell the user to begin a new agent session afterward. - -## Follow the required shape - -Use the directory name as the skill name. Use lowercase letters, numbers, and single hyphens only. Keep the name short and no longer than 64 characters. - -Every skill needs a SKILL.md with YAML frontmatter followed by Markdown instructions: - -~~~md ---- -name: release-notes -description: Draft concise release notes from completed changes. Use when the user asks for a changelog, release notes, or GitHub release copy. ---- - -# Draft release notes - -Write the workflow here. -~~~ - -Require a non-empty description that says both what the skill does and when to use it. Keep the body procedural and concise. Put detailed schemas, long examples, and variant-specific guidance in references/ only when the skill needs them. - -Use scripts/ for repeatable or fragile operations that benefit from deterministic execution. Use assets/ for files that belong in generated output. Do not add README files, changelogs, or setup notes that do not help the model perform the task. - -## Create safely - -1. Identify the repeated task, expected inputs, and useful output. -2. Choose the smallest name and description that reliably trigger the skill. -3. Create the folder and SKILL.md; add resources only when they remove real repeated work. -4. Re-read the completed file and verify its frontmatter, directory-name match, and relative resource paths. -5. Tell the user where it was created and that a new agent session will discover it. - -Skills provide instructions only. They do not grant filesystem, network, shell, or approval privileges, and they do not make a tool available. Use only the tools that are actually available, follow their normal approval rules, and ask before actions that need user authorization. -` -) - -var skillName = regexp.MustCompile(`^[a-z0-9]+(?:-[a-z0-9]+)*$`) - -// SkillsDir returns the canonical runtime-owned skill directory. -func SkillsDir() (string, error) { - if path := strings.TrimSpace(os.Getenv(SkillsDirEnv)); path != "" { - return filepath.Abs(path) - } - if xdg := strings.TrimSpace(os.Getenv("XDG_CONFIG_HOME")); xdg != "" { - return filepath.Join(xdg, "ollama", "skills"), nil - } - home, err := os.UserHomeDir() - if err != nil { - return "", err - } - return filepath.Join(home, ".ollama", "skills"), nil -} - -// Skill is a validated, loadable instruction set. It never grants tool -// permissions; it is supplied to the model as ordinary tool-result content. -type Skill struct { - Name string - Description string - Instructions string - Path string -} - -func (s Skill) Content() string { - var b strings.Builder - fmt.Fprintf(&b, "\n%s\n", s.Name, strings.TrimSpace(s.Instructions)) - if s.Path != "" { - dir := filepath.Dir(s.Path) - fmt.Fprintf(&b, "Skill directory: %s\n", dir) - b.WriteString("Relative paths in this skill are relative to the skill directory.\n") - } - if resources := s.resources(); len(resources) > 0 { - b.WriteString("\n") - for _, r := range resources { - fmt.Fprintf(&b, " %s\n", r) - } - b.WriteString("\n") - } - b.WriteString("") - return b.String() -} - -// resources lists bundled files one level deep under scripts/, references/, -// and assets/ without reading them, so the model can load them on demand. -func (s Skill) resources() []string { - if s.Path == "" { - return nil - } - dir := filepath.Dir(s.Path) - var resources []string - for _, sub := range []string{"scripts", "references", "assets"} { - entries, err := os.ReadDir(filepath.Join(dir, sub)) - if err != nil { - continue - } - for _, e := range entries { - if e.IsDir() { - continue - } - resources = append(resources, sub+"/"+e.Name()) - } - } - sort.Strings(resources) - return resources -} - -// SkillCatalog contains valid skills and diagnostics for ignored invalid -// entries, so one malformed skill cannot hide the rest. -type SkillCatalog struct { - dir string - skills map[string]Skill - diagnostics []error -} - -func DiscoverSkills(dir string) (*SkillCatalog, error) { - dir, err := filepath.Abs(strings.TrimSpace(dir)) - if err != nil { - return nil, err - } - catalog := &SkillCatalog{dir: dir, skills: make(map[string]Skill)} - entries, err := os.ReadDir(dir) - if errors.Is(err, fs.ErrNotExist) { - return catalog, nil - } - if err != nil { - return nil, fmt.Errorf("read skills directory: %w", err) - } - for _, entry := range entries { - name := entry.Name() - // Follow symlinks so users can point at shared skill repositories. - // The link name (not the target) is the canonical skill name. - info, err := os.Stat(filepath.Join(dir, name)) - if err != nil { - if errors.Is(err, fs.ErrNotExist) { - continue - } - catalog.diagnostics = append(catalog.diagnostics, fmt.Errorf("skill %q: %w", name, err)) - continue - } - if !info.IsDir() { - continue - } - if !skillName.MatchString(name) { - catalog.diagnostics = append(catalog.diagnostics, fmt.Errorf("invalid skill directory %q", name)) - continue - } - skill, err := parseSkill(filepath.Join(dir, name, skillFilename), name) - if errors.Is(err, fs.ErrNotExist) { - continue - } - if err != nil { - catalog.diagnostics = append(catalog.diagnostics, err) - continue - } - catalog.skills[skill.Name] = skill - } - return catalog, nil -} - -// LoadDefaultSkills discovers skills from the spec's scopes, merged with -// deterministic precedence. Roots are scanned lowest-precedence first so later -// roots override earlier ones on name collisions (recording a diagnostic): -// -// 1. ~/.agents/skills/ (user, cross-client) -// 2. user Ollama skills dir (user, Ollama-owned; SkillsDir) -// 3. /.agents/skills/ (project, cross-client) -// 4. /.ollama/skills/ (project, Ollama-owned) -// -// Project-level overrides user-level, and within a scope Ollama-owned -// directories override .agents/skills/. projectDir is the agent's working -// directory at startup (discovery is a session-start snapshot per the spec). -func LoadDefaultSkills(projectDir string) (*SkillCatalog, error) { - roots, err := defaultSkillRoots(projectDir) - if err != nil { - return nil, err - } - catalog := &SkillCatalog{skills: make(map[string]Skill)} - bundled, err := bundledSkillCreator() - if err != nil { - return nil, err - } - catalog.skills[bundled.Name] = bundled - if err := installBundledSkillCreator(); err != nil { - catalog.diagnostics = append(catalog.diagnostics, err) - } - for _, root := range roots { - sub, err := DiscoverSkills(root.path) - if err != nil { - catalog.diagnostics = append(catalog.diagnostics, fmt.Errorf("discover skills in %s: %w", root.path, err)) - continue - } - catalog.diagnostics = append(catalog.diagnostics, sub.diagnostics...) - for _, skill := range sub.skills { - // Name collisions across roots are expected precedence resolution, - // not errors: later (higher-precedence) roots legitimately override - // earlier ones. The skill is still loaded; no diagnostic needed. - catalog.skills[skill.Name] = skill - } - } - return catalog, nil -} - -func bundledSkillCreator() (Skill, error) { - skill, err := parseSkillContent("", bundledSkillCreatorName, bundledSkillCreatorContent) - if err != nil { - return Skill{}, fmt.Errorf("load bundled %s skill: %w", bundledSkillCreatorName, err) - } - return skill, nil -} - -func installBundledSkillCreator() error { - dir, err := SkillsDir() - if err != nil { - return fmt.Errorf("resolve bundled skill directory: %w", err) - } - path := filepath.Join(dir, bundledSkillCreatorName, skillFilename) - if err := os.MkdirAll(filepath.Dir(path), 0o755); err != nil { - return fmt.Errorf("create bundled skill directory: %w", err) - } - contents, err := os.ReadFile(path) - if err == nil && string(contents) == bundledSkillCreatorContent { - return nil - } - if err != nil && !errors.Is(err, fs.ErrNotExist) { - return fmt.Errorf("read bundled skill: %w", err) - } - if err := os.WriteFile(path, []byte(bundledSkillCreatorContent), 0o644); err != nil { - return fmt.Errorf("write bundled skill: %w", err) - } - return nil -} - -type skillRoot struct { - path string -} - -// SkillImportResult describes one import attempt. Failed skills do not prevent -// other valid skills in the same source root from being imported. -type SkillImportResult struct { - Source string - SourceDir string - Destination string - Imported []string - Existing []string - Failures []SkillImportFailure -} - -// SkillImportFailure identifies a source skill that was deliberately skipped. -// The destination is never changed for a failed skill. -type SkillImportFailure struct { - Name string - Err error -} - -// ImportSkills imports skills from a conventional coding-agent source into the -// canonical Ollama skills directory. Supported sources are codex, claude, and -// pi. Existing skills are left untouched: an identical directory is reported -// as existing, and a differing one is reported as a conflict. -func ImportSkills(source string) (SkillImportResult, error) { - home, err := os.UserHomeDir() - if err != nil { - return SkillImportResult{}, fmt.Errorf("resolve home directory: %w", err) - } - - destination, err := SkillsDir() - if err != nil { - return SkillImportResult{}, fmt.Errorf("resolve Ollama skills directory: %w", err) - } - return importSkillsFromRoots(source, conventionalSkillImportRoots(home), destination) -} - -func conventionalSkillImportRoots(home string) map[string]string { - return map[string]string{ - "codex": filepath.Join(home, ".codex", "skills"), - "claude": filepath.Join(home, ".claude", "skills"), - "pi": filepath.Join(home, ".pi", "agent", "skills"), - } -} - -func importSkillsFromRoots(source string, roots map[string]string, destination string) (SkillImportResult, error) { - source = strings.ToLower(strings.TrimSpace(source)) - sourceDir, ok := roots[source] - if !ok { - return SkillImportResult{}, fmt.Errorf("unknown skill source %q", source) - } - return importSkillsFromDir(source, sourceDir, destination) -} - -func importSkillsFromDir(source, sourceDir, destination string) (SkillImportResult, error) { - result := SkillImportResult{Source: source, SourceDir: sourceDir, Destination: destination} - info, err := os.Lstat(sourceDir) - if errors.Is(err, fs.ErrNotExist) { - return result, nil - } - if err != nil { - return result, fmt.Errorf("inspect %s skills directory: %w", source, err) - } - if info.Mode()&os.ModeSymlink != 0 { - return result, fmt.Errorf("inspect %s skills directory: symlinks are not supported", source) - } - if !info.IsDir() { - return result, fmt.Errorf("inspect %s skills directory: not a directory", source) - } - - entries, err := os.ReadDir(sourceDir) - if err != nil { - return result, fmt.Errorf("read %s skills directory: %w", source, err) - } - for _, entry := range entries { - name := entry.Name() - path := filepath.Join(sourceDir, name) - if entry.Type()&os.ModeSymlink != 0 { - result.Failures = append(result.Failures, SkillImportFailure{Name: name, Err: errors.New("symlinked skill directories are not supported")}) - continue - } - info, err := entry.Info() - if err != nil { - result.Failures = append(result.Failures, SkillImportFailure{Name: name, Err: fmt.Errorf("inspect source: %w", err)}) - continue - } - if !info.IsDir() { - continue - } - if !skillName.MatchString(name) { - result.Failures = append(result.Failures, SkillImportFailure{Name: name, Err: errors.New("invalid skill directory name")}) - continue - } - if err := validateImportSkill(path, name); err != nil { - result.Failures = append(result.Failures, SkillImportFailure{Name: name, Err: err}) - continue - } - - state, err := importSkillDirectory(path, filepath.Join(destination, name)) - if err != nil { - result.Failures = append(result.Failures, SkillImportFailure{Name: name, Err: err}) - continue - } - if state == skillImportExisting { - result.Existing = append(result.Existing, name) - } else { - result.Imported = append(result.Imported, name) - } - } - return result, nil -} - -func validateImportSkill(dir, name string) error { - manifest := filepath.Join(dir, skillFilename) - info, err := os.Lstat(manifest) - if err != nil { - return fmt.Errorf("inspect %s: %w", skillFilename, err) - } - if info.Mode()&os.ModeSymlink != 0 || !info.Mode().IsRegular() { - return fmt.Errorf("%s must be a regular, non-symlinked file", skillFilename) - } - if _, err := parseSkill(manifest, name); err != nil { - return err - } - return walkImportTree(dir, func(path string, entry fs.DirEntry, info fs.FileInfo) error { - if info.IsDir() || path == dir { - return nil - } - if !info.Mode().IsRegular() { - return fmt.Errorf("only regular files may be imported: %s", path) - } - file, err := os.Open(path) - if err != nil { - return fmt.Errorf("read %s: %w", path, err) - } - return file.Close() - }) -} - -func walkImportTree(root string, visit func(string, fs.DirEntry, fs.FileInfo) error) error { - return filepath.WalkDir(root, func(path string, entry fs.DirEntry, err error) error { - if err != nil { - return err - } - rel, err := filepath.Rel(root, path) - if err != nil || rel == ".." || strings.HasPrefix(rel, ".."+string(filepath.Separator)) { - return fmt.Errorf("unsafe skill path %q", path) - } - if entry.Type()&os.ModeSymlink != 0 { - return fmt.Errorf("symlinks may not be imported: %s", path) - } - info, err := entry.Info() - if err != nil { - return err - } - return visit(path, entry, info) - }) -} - -type skillImportState int - -const ( - skillImportCopied skillImportState = iota - skillImportExisting -) - -func importSkillDirectory(source, destination string) (skillImportState, error) { - if info, err := os.Lstat(destination); err == nil { - if info.Mode()&os.ModeSymlink != 0 || !info.IsDir() { - return 0, errors.New("destination exists but is not a regular directory") - } - same, err := sameImportTree(source, destination) - if err != nil { - return 0, fmt.Errorf("inspect existing destination: %w", err) - } - if same { - return skillImportExisting, nil - } - return 0, errors.New("destination skill already exists with different contents") - } else if !errors.Is(err, fs.ErrNotExist) { - return 0, fmt.Errorf("inspect destination: %w", err) - } - - if err := ensureImportDestination(filepath.Dir(destination)); err != nil { - return 0, err - } - stage, err := os.MkdirTemp(filepath.Dir(destination), "."+filepath.Base(destination)+".import-") - if err != nil { - return 0, fmt.Errorf("create import staging directory: %w", err) - } - defer os.RemoveAll(stage) - if err := copyImportTree(source, stage); err != nil { - return 0, err - } - if _, err := os.Lstat(destination); err == nil { - return 0, errors.New("destination skill was created during import") - } else if !errors.Is(err, fs.ErrNotExist) { - return 0, fmt.Errorf("inspect destination before install: %w", err) - } - if err := os.Rename(stage, destination); err != nil { - return 0, fmt.Errorf("install imported skill: %w", err) - } - return skillImportCopied, nil -} - -func ensureImportDestination(dir string) error { - if err := os.MkdirAll(dir, 0o755); err != nil { - return fmt.Errorf("create Ollama skills directory: %w", err) - } - info, err := os.Lstat(dir) - if err != nil { - return fmt.Errorf("inspect Ollama skills directory: %w", err) - } - if info.Mode()&os.ModeSymlink != 0 || !info.IsDir() { - return errors.New("Ollama skills directory must be a regular, non-symlinked directory") - } - return nil -} - -func copyImportTree(source, destination string) error { - return walkImportTree(source, func(path string, entry fs.DirEntry, info fs.FileInfo) error { - rel, err := filepath.Rel(source, path) - if err != nil { - return err - } - target := destination - if rel != "." { - target = filepath.Join(destination, rel) - } - if info.IsDir() { - if rel == "." { - return nil - } - return os.Mkdir(target, info.Mode().Perm()) - } - if !info.Mode().IsRegular() { - return fmt.Errorf("only regular files may be imported: %s", path) - } - return copyImportFile(path, target, info.Mode().Perm()) - }) -} - -func copyImportFile(source, destination string, mode fs.FileMode) error { - in, err := os.Open(source) - if err != nil { - return fmt.Errorf("read %s: %w", source, err) - } - defer in.Close() - out, err := os.OpenFile(destination, os.O_WRONLY|os.O_CREATE|os.O_EXCL, mode) - if err != nil { - return fmt.Errorf("create %s: %w", destination, err) - } - _, copyErr := io.Copy(out, in) - closeErr := out.Close() - if copyErr != nil { - return fmt.Errorf("copy %s: %w", source, copyErr) - } - if closeErr != nil { - return fmt.Errorf("write %s: %w", destination, closeErr) - } - return nil -} - -func sameImportTree(source, destination string) (bool, error) { - seen := make(map[string]struct{}) - same := true - err := walkImportTree(source, func(path string, entry fs.DirEntry, info fs.FileInfo) error { - rel, err := filepath.Rel(source, path) - if err != nil { - return err - } - seen[rel] = struct{}{} - other := destination - if rel != "." { - other = filepath.Join(destination, rel) - } - otherInfo, err := os.Lstat(other) - if errors.Is(err, fs.ErrNotExist) { - same = false - return nil - } - if err != nil { - return err - } - if otherInfo.Mode()&os.ModeSymlink != 0 || otherInfo.IsDir() != info.IsDir() || (!info.IsDir() && !otherInfo.Mode().IsRegular()) { - same = false - return nil - } - if info.Mode().IsRegular() { - equal, err := sameImportFile(path, other) - if err != nil { - return err - } - if !equal { - same = false - } - } - return nil - }) - if err != nil || !same { - return same, err - } - err = walkImportTree(destination, func(path string, entry fs.DirEntry, info fs.FileInfo) error { - rel, err := filepath.Rel(destination, path) - if err != nil { - return err - } - if _, ok := seen[rel]; !ok { - same = false - } - return nil - }) - return same, err -} - -func sameImportFile(first, second string) (bool, error) { - a, err := os.Open(first) - if err != nil { - return false, err - } - defer a.Close() - b, err := os.Open(second) - if err != nil { - return false, err - } - defer b.Close() - - left := make([]byte, 32*1024) - right := make([]byte, len(left)) - for { - n, errA := a.Read(left) - m, errB := b.Read(right) - if n != m || !bytes.Equal(left[:n], right[:m]) { - return false, nil - } - if errA == io.EOF && errB == io.EOF { - return true, nil - } - if errA != nil && errA != io.EOF { - return false, errA - } - if errB != nil && errB != io.EOF { - return false, errB - } - if errA == io.EOF || errB == io.EOF { - return false, nil - } - } -} - -// defaultSkillRoots returns skill directories ordered lowest- to -// highest-precedence. Non-existent directories are scanned harmlessly -// (DiscoverSkills skips them). -func defaultSkillRoots(projectDir string) ([]skillRoot, error) { - var roots []skillRoot - - if home, err := os.UserHomeDir(); err == nil && home != "" { - roots = append(roots, skillRoot{path: filepath.Join(home, ".agents", "skills")}) - } - - userOllama, err := SkillsDir() - if err != nil { - return nil, err - } - roots = append(roots, skillRoot{path: userOllama}) - - projectDir = strings.TrimSpace(projectDir) - if projectDir != "" { - if abs, err := filepath.Abs(projectDir); err == nil { - roots = append(roots, - skillRoot{path: filepath.Join(abs, ".agents", "skills")}, - skillRoot{path: filepath.Join(abs, ".ollama", "skills")}, - ) - } - } - return roots, nil -} - -func (c *SkillCatalog) Dir() string { - if c == nil { - return "" - } - return c.dir -} - -func (c *SkillCatalog) List() []Skill { - if c == nil { - return nil - } - list := make([]Skill, 0, len(c.skills)) - for _, skill := range c.skills { - list = append(list, skill) - } - sort.Slice(list, func(i, j int) bool { return list[i].Name < list[j].Name }) - return list -} - -func (c *SkillCatalog) Diagnostics() []error { - if c == nil { - return nil - } - return append([]error(nil), c.diagnostics...) -} - -// ExcludeNames removes skills whose names are reserved by a caller. It returns -// the excluded names in sorted order. -func (c *SkillCatalog) ExcludeNames(names []string) []string { - if c == nil { - return nil - } - reserved := make(map[string]struct{}, len(names)) - for _, name := range names { - name = strings.TrimPrefix(strings.ToLower(strings.TrimSpace(name)), "/") - if name != "" { - reserved[name] = struct{}{} - } - } - var excluded []string - for name := range c.skills { - if _, ok := reserved[name]; !ok { - continue - } - delete(c.skills, name) - excluded = append(excluded, name) - } - sort.Strings(excluded) - return excluded -} - -func (c *SkillCatalog) Load(name string) (Skill, error) { - name = strings.TrimSpace(name) - if !skillName.MatchString(name) { - return Skill{}, fmt.Errorf("invalid skill name %q", name) - } - if c == nil { - return Skill{}, errors.New("skills are unavailable") - } - skill, ok := c.skills[name] - if !ok { - return Skill{}, fmt.Errorf("skill %q not found in %s", name, c.dir) - } - return skill, nil -} - -// SystemContext advertises the catalog without expanding full instructions in -// every request. The skill call is the explicit loading boundary. -func (c *SkillCatalog) SystemContext() string { - list := c.List() - if len(list) == 0 { - return "" - } - lines := []string{""} - for _, skill := range list { - description := skill.Description - if description == "" { - description = "No description provided." - } - lines = append(lines, fmt.Sprintf("- %s: %s", skill.Name, description)) - } - lines = append(lines, "", "Load a matching skill with the skill tool before following its instructions. Skills only provide instructions; use ordinary tools for filesystem or network access, with their normal approval rules.") - return strings.Join(lines, "\n") -} - -func parseSkill(path, directoryName string) (Skill, error) { - // Stat (not Lstat) so a symlinked SKILL.md resolves to its target file. - info, err := os.Stat(path) - if err != nil { - return Skill{}, err - } - if !info.Mode().IsRegular() { - return Skill{}, fmt.Errorf("skill %q: %s is not a regular file", directoryName, skillFilename) - } - if info.Size() > maxSkillBytes { - return Skill{}, fmt.Errorf("skill %q: %s exceeds %d bytes", directoryName, skillFilename, maxSkillBytes) - } - data, err := os.ReadFile(path) - if err != nil { - return Skill{}, fmt.Errorf("read skill %q: %w", directoryName, err) - } - return parseSkillContent(path, directoryName, string(data)) -} - -func parseSkillContent(path, directoryName, input string) (Skill, error) { - instructions := strings.TrimSpace(input) - if instructions == "" { - return Skill{}, fmt.Errorf("skill %q: %s is empty", directoryName, skillFilename) - } - if !strings.HasPrefix(instructions, "---\n") && !strings.HasPrefix(instructions, "---\r\n") { - return Skill{}, fmt.Errorf("skill %q: missing YAML front matter", directoryName) - } - metadata, body, err := skillFrontMatter(instructions) - if err != nil { - return Skill{}, fmt.Errorf("skill %q: %w", directoryName, err) - } - if metadata.Name == "" { - return Skill{}, fmt.Errorf("skill %q: front matter requires name", directoryName) - } - if metadata.Description == "" { - return Skill{}, fmt.Errorf("skill %q: front matter requires description", directoryName) - } - if !skillName.MatchString(metadata.Name) { - return Skill{}, fmt.Errorf("skill %q: invalid front matter name %q", directoryName, metadata.Name) - } - if metadata.Name != directoryName { - return Skill{}, fmt.Errorf("skill %q: front matter name %q must match directory name", directoryName, metadata.Name) - } - skill := Skill{Name: metadata.Name, Description: metadata.Description, Path: path} - instructions = body - if strings.TrimSpace(instructions) == "" { - return Skill{}, fmt.Errorf("skill %q: instructions are empty", directoryName) - } - skill.Instructions = strings.TrimSpace(instructions) - return skill, nil -} - -type skillFrontMatterMetadata struct { - Name string `yaml:"name"` - Description string `yaml:"description"` - Metadata map[string]any `yaml:"metadata"` -} - -func skillFrontMatter(input string) (skillFrontMatterMetadata, string, error) { - input = strings.ReplaceAll(input, "\r\n", "\n") - lines := strings.Split(input, "\n") - if len(lines) < 3 || lines[0] != "---" { - return skillFrontMatterMetadata{}, "", errors.New("invalid front matter") - } - for i := 1; i < len(lines); i++ { - if lines[i] == "---" { - var metadata skillFrontMatterMetadata - if err := yaml.Unmarshal([]byte(strings.Join(lines[1:i], "\n")), &metadata); err != nil { - return skillFrontMatterMetadata{}, "", fmt.Errorf("parse YAML front matter: %w", err) - } - metadata.Name = strings.TrimSpace(metadata.Name) - metadata.Description = strings.TrimSpace(metadata.Description) - return metadata, strings.Join(lines[i+1:], "\n"), nil - } - } - return skillFrontMatterMetadata{}, "", errors.New("front matter is not closed") -} diff --git a/agent/skills_test.go b/agent/skills_test.go deleted file mode 100644 index 99e93feb6e9..00000000000 --- a/agent/skills_test.go +++ /dev/null @@ -1,516 +0,0 @@ -package agent - -import ( - "os" - "path/filepath" - "strings" - "testing" -) - -func writeCatalogSkill(t *testing.T, dir, name, content string) { - t.Helper() - path := filepath.Join(dir, name) - if err := os.MkdirAll(path, 0o755); err != nil { - t.Fatal(err) - } - if !strings.HasPrefix(content, "---") { - content = "---\nname: " + name + "\ndescription: Test skill.\n---\n" + content - } - if err := os.WriteFile(filepath.Join(path, skillFilename), []byte(content), 0o644); err != nil { - t.Fatal(err) - } -} - -func writeImportFixtureSkill(t *testing.T, dir string) { - t.Helper() - contents, err := os.ReadFile(filepath.Join("testdata", "import", "release-notes", skillFilename)) - if err != nil { - t.Fatal(err) - } - path := filepath.Join(dir, "release-notes", skillFilename) - if err := os.MkdirAll(filepath.Dir(path), 0o755); err != nil { - t.Fatal(err) - } - if err := os.WriteFile(path, contents, 0o644); err != nil { - t.Fatal(err) - } -} - -func TestDiscoverAndLoadSkills(t *testing.T) { - dir := t.TempDir() - writeCatalogSkill(t, dir, "release-notes", "---\nname: release-notes\ndescription: Draft concise release notes.\nmetadata:\n author: Ollama\n labels:\n - release\n - docs\n---\n# Release notes\n\nUse short bullets.") - catalog, err := DiscoverSkills(dir) - if err != nil { - t.Fatal(err) - } - list := catalog.List() - if len(list) != 1 || list[0].Name != "release-notes" || list[0].Description != "Draft concise release notes." { - t.Fatalf("skills = %#v", list) - } - skill, err := catalog.Load("release-notes") - if err != nil { - t.Fatal(err) - } - if !strings.Contains(skill.Content(), ``) || !strings.Contains(skill.Content(), "Use short bullets.") { - t.Fatalf("skill content = %q", skill.Content()) - } - if context := catalog.SystemContext(); !strings.Contains(context, "release-notes: Draft concise release notes.") || !strings.Contains(context, "normal approval rules") { - t.Fatalf("system context = %q", context) - } -} - -func TestDiscoverSkillsSkipsMalformedEntries(t *testing.T) { - dir := t.TempDir() - writeCatalogSkill(t, dir, "valid", "do the useful thing") - writeCatalogSkill(t, dir, "mismatched", "---\nname: whatever\ndescription: wrong name\n---\nbody") - // Genuinely malformed front matter (a line without a key:value pair) is still rejected. - writeCatalogSkill(t, dir, "broken", "---\nname: broken\ndescription\n---\nnope") - writeCatalogSkill(t, dir, "missing-name", "---\ndescription: missing name\n---\nbody") - writeCatalogSkill(t, dir, "missing-description", "---\nname: missing-description\n---\nbody") - writeCatalogSkill(t, dir, "bad-name", "---\nname: bad_name\ndescription: invalid name\n---\nbody") - writeCatalogSkill(t, dir, "under_score", "---\nname: under_score\ndescription: invalid directory\n---\nbody") - if err := os.MkdirAll(filepath.Join(dir, "no-front-matter"), 0o755); err != nil { - t.Fatal(err) - } - if err := os.WriteFile(filepath.Join(dir, "no-front-matter", skillFilename), []byte("body"), 0o644); err != nil { - t.Fatal(err) - } - catalog, err := DiscoverSkills(dir) - if err != nil { - t.Fatal(err) - } - if got, want := len(catalog.List()), 1; got != want { - t.Fatalf("valid skills = %d, want %d", got, want) - } - if got, want := len(catalog.Diagnostics()), 7; got != want { - t.Fatalf("diagnostics = %d, want %d: %#v", got, want, catalog.Diagnostics()) - } - if _, err := catalog.Load("broken"); err == nil || !strings.Contains(err.Error(), "not found") { - t.Fatalf("load broken error = %v", err) - } - if _, err := catalog.Load("../valid"); err == nil || !strings.Contains(err.Error(), "invalid skill name") { - t.Fatalf("unsafe name error = %v", err) - } -} - -func TestDiscoverSkillsFollowsSymlinks(t *testing.T) { - dir := t.TempDir() - target := t.TempDir() - writeCatalogSkill(t, target, "shared", "---\nname: shared\ndescription: From a linked repo.\n---\nshared instructions") - if err := os.Symlink(filepath.Join(target, "shared"), filepath.Join(dir, "shared")); err != nil { - t.Skipf("symlink not supported: %v", err) - } - catalog, err := DiscoverSkills(dir) - if err != nil { - t.Fatal(err) - } - list := catalog.List() - if len(list) != 1 || list[0].Name != "shared" || list[0].Description != "From a linked repo." { - t.Fatalf("symlinked skills = %#v", list) - } - if !strings.Contains(list[0].Content(), "shared instructions") { - t.Fatalf("symlinked skill content = %q", list[0].Content()) - } -} - -func TestLoadDefaultSkillsContinuesAfterBadRoot(t *testing.T) { - home := t.TempDir() - t.Setenv("HOME", home) - t.Setenv("USERPROFILE", home) - project := t.TempDir() - writeCatalogSkill(t, filepath.Join(project, ".ollama", "skills"), "release-notes", "project instructions") - - badRoot := filepath.Join(t.TempDir(), "not-a-directory") - if err := os.WriteFile(badRoot, []byte("not a directory"), 0o644); err != nil { - t.Fatal(err) - } - t.Setenv(SkillsDirEnv, badRoot) - - catalog, err := LoadDefaultSkills(project) - if err != nil { - t.Fatal(err) - } - if _, err := catalog.Load("release-notes"); err != nil { - t.Fatalf("valid skill was hidden by bad root: %v", err) - } - if _, err := catalog.Load(bundledSkillCreatorName); err != nil { - t.Fatalf("bundled skill was hidden by bad root: %v", err) - } - var foundDiagnostic bool - for _, diagnostic := range catalog.Diagnostics() { - if strings.Contains(diagnostic.Error(), badRoot) { - foundDiagnostic = true - break - } - } - if !foundDiagnostic { - t.Fatalf("diagnostics = %#v, want bad root %q", catalog.Diagnostics(), badRoot) - } -} - -func TestLoadDefaultSkillsInstallsBundledSkillCreator(t *testing.T) { - dir := t.TempDir() - t.Setenv(SkillsDirEnv, dir) - - catalog, err := LoadDefaultSkills("") - if err != nil { - t.Fatal(err) - } - skill, err := catalog.Load(bundledSkillCreatorName) - if err != nil { - t.Fatal(err) - } - path := filepath.Join(dir, bundledSkillCreatorName, skillFilename) - contents, err := os.ReadFile(path) - if err != nil { - t.Fatal(err) - } - if string(contents) != bundledSkillCreatorContent { - t.Fatalf("installed skill = %q, want bundled contents", contents) - } - if skill.Path != path { - t.Fatalf("skill path = %q, want %q", skill.Path, path) - } - if !strings.Contains(skill.Content(), "Skill directory: "+filepath.Dir(path)) { - t.Fatalf("skill content does not identify its directory: %q", skill.Content()) - } -} - -func TestLoadDefaultSkillsUpdatesExistingSkillCreator(t *testing.T) { - dir := t.TempDir() - t.Setenv(SkillsDirEnv, dir) - writeCatalogSkill(t, dir, bundledSkillCreatorName, "custom instructions") - - if _, err := LoadDefaultSkills(""); err != nil { - t.Fatal(err) - } - contents, err := os.ReadFile(filepath.Join(dir, bundledSkillCreatorName, skillFilename)) - if err != nil { - t.Fatal(err) - } - if string(contents) != bundledSkillCreatorContent { - t.Fatalf("installed skill = %q, want bundled contents", contents) - } -} - -func TestSkillsDirUsesOverrideAndXDG(t *testing.T) { - base := t.TempDir() - - override := filepath.Join(base, "skills-override") - t.Setenv(SkillsDirEnv, override) - got, err := SkillsDir() - if err != nil { - t.Fatal(err) - } - want, err := filepath.Abs(override) - if err != nil { - t.Fatal(err) - } - if got != want { - t.Fatalf("SkillsDir override = %q, want %q", got, want) - } - - t.Setenv(SkillsDirEnv, "") - xdg := filepath.Join(base, "xdg") - t.Setenv("XDG_CONFIG_HOME", xdg) - if got, err := SkillsDir(); err != nil || got != filepath.Join(xdg, "ollama", "skills") { - t.Fatalf("SkillsDir xdg = %q, want %q, %v", got, filepath.Join(xdg, "ollama", "skills"), err) - } - - t.Setenv("XDG_CONFIG_HOME", "") - home := filepath.Join(base, "home") - t.Setenv("HOME", home) - t.Setenv("USERPROFILE", home) - if got, err := SkillsDir(); err != nil || got != filepath.Join(home, ".ollama", "skills") { - t.Fatalf("SkillsDir default = %q, want %q, %v", got, filepath.Join(home, ".ollama", "skills"), err) - } -} - -func TestLoadDefaultSkillsPrecedenceAndCollisions(t *testing.T) { - home := t.TempDir() - t.Setenv("HOME", home) - t.Setenv("USERPROFILE", home) // Windows: os.UserHomeDir uses %USERPROFILE% - - userOllama := t.TempDir() - t.Setenv(SkillsDirEnv, userOllama) - - userAgents := filepath.Join(home, ".agents", "skills") - project := t.TempDir() - projectAgents := filepath.Join(project, ".agents", "skills") - projectOllama := filepath.Join(project, ".ollama", "skills") - - // release-notes exists in all four roots; project ollama must win. - writeCatalogSkill(t, userAgents, "release-notes", "from user agents") - writeCatalogSkill(t, userOllama, "release-notes", "from user ollama") - writeCatalogSkill(t, projectOllama, "release-notes", "from project ollama") - // code-review exists in both project roots; project ollama beats project agents. - writeCatalogSkill(t, projectAgents, "code-review", "from project agents") - writeCatalogSkill(t, projectOllama, "code-review", "from project ollama") - // unique appears only in user ollama (via env override). - writeCatalogSkill(t, userOllama, "unique", "only here") - - catalog, err := LoadDefaultSkills(project) - if err != nil { - t.Fatal(err) - } - rn, err := catalog.Load("release-notes") - if err != nil || !strings.Contains(rn.Instructions, "from project ollama") || !strings.Contains(rn.Path, ".ollama") { - t.Fatalf("release-notes = %#v, want project ollama to win", rn) - } - cr, err := catalog.Load("code-review") - if err != nil || !strings.Contains(cr.Instructions, "from project ollama") { - t.Fatalf("code-review = %#v, want project ollama to win over project agents", cr) - } - if _, err := catalog.Load("unique"); err != nil { - t.Fatalf("unique should load from user ollama: %v", err) - } - // Collisions are resolved silently by precedence — no diagnostics. - for _, d := range catalog.Diagnostics() { - if strings.Contains(d.Error(), "shadows") { - t.Fatalf("unexpected shadow diagnostic: %v", d) - } - } -} - -func TestSkillCatalogExcludeNames(t *testing.T) { - dir := t.TempDir() - for _, name := range []string{"release-notes", "system", "exit"} { - writeCatalogSkill(t, dir, name, "instructions") - } - catalog, err := DiscoverSkills(dir) - if err != nil { - t.Fatal(err) - } - - if got, want := strings.Join(catalog.ExcludeNames([]string{"/system", "EXIT"}), ","), "exit,system"; got != want { - t.Fatalf("excluded skills = %q, want %q", got, want) - } - if _, err := catalog.Load("system"); err == nil { - t.Fatal("excluded system skill should not load") - } - if _, err := catalog.Load("exit"); err == nil { - t.Fatal("excluded exit skill should not load") - } - if _, err := catalog.Load("release-notes"); err != nil { - t.Fatalf("non-conflicting skill should remain available: %v", err) - } -} - -func TestSkillContentListsDirectoryAndResources(t *testing.T) { - root := t.TempDir() - skillDir := filepath.Join(root, "pdf-processing") - if err := os.MkdirAll(filepath.Join(skillDir, "scripts"), 0o755); err != nil { - t.Fatal(err) - } - if err := os.MkdirAll(filepath.Join(skillDir, "references"), 0o755); err != nil { - t.Fatal(err) - } - if err := os.WriteFile(filepath.Join(skillDir, "SKILL.md"), []byte("---\nname: pdf-processing\ndescription: Handle PDFs.\n---\nHandle PDFs."), 0o644); err != nil { - t.Fatal(err) - } - if err := os.WriteFile(filepath.Join(skillDir, "scripts", "extract.py"), []byte("#!/usr/bin/env python3"), 0o755); err != nil { - t.Fatal(err) - } - if err := os.WriteFile(filepath.Join(skillDir, "references", "ref.md"), []byte("ref"), 0o644); err != nil { - t.Fatal(err) - } - catalog, err := DiscoverSkills(root) - if err != nil { - t.Fatal(err) - } - skill, err := catalog.Load("pdf-processing") - if err != nil { - t.Fatal(err) - } - content := skill.Content() - if !strings.Contains(content, "Skill directory:") || !strings.Contains(content, skillDir) { - t.Fatalf("content missing skill directory: %q", content) - } - if !strings.Contains(content, "scripts/extract.py") || !strings.Contains(content, "references/ref.md") { - t.Fatalf("content missing resource listing: %q", content) - } -} - -func TestImportSkillsCopiesFixtureAndIsIdempotent(t *testing.T) { - source := t.TempDir() - destination := t.TempDir() - writeImportFixtureSkill(t, source) - writeCatalogSkill(t, source, "broken", "---\nname: another-skill\ndescription: Deliberately invalid.\n---\nIgnore this.") - if err := os.MkdirAll(filepath.Join(source, "release-notes", "references"), 0o755); err != nil { - t.Fatal(err) - } - if err := os.WriteFile(filepath.Join(source, "release-notes", "references", "style.txt"), []byte("Keep it short.\n"), 0o644); err != nil { - t.Fatal(err) - } - if err := os.MkdirAll(filepath.Join(source, "release-notes", "scripts"), 0o755); err != nil { - t.Fatal(err) - } - if err := os.WriteFile(filepath.Join(source, "release-notes", "scripts", "prepare.sh"), []byte("#!/bin/sh\n"), 0o755); err != nil { - t.Fatal(err) - } - if err := os.WriteFile(filepath.Join(source, "ignored.md"), []byte("Ignored root file.\n"), 0o644); err != nil { - t.Fatal(err) - } - - result, err := importSkillsFromDir("codex", source, destination) - if err != nil { - t.Fatal(err) - } - if got, want := strings.Join(result.Imported, ","), "release-notes"; got != want { - t.Fatalf("imported = %q, want %q", got, want) - } - catalog, err := DiscoverSkills(destination) - if err != nil { - t.Fatal(err) - } - skill, err := catalog.Load("release-notes") - if err != nil || skill.Description != "Draft concise release notes." { - t.Fatalf("imported skill = %#v, %v", skill, err) - } - if got := len(result.Failures); got != 1 || result.Failures[0].Name != "broken" { - t.Fatalf("failures = %#v, want broken fixture failure", result.Failures) - } - for _, file := range []string{skillFilename, filepath.Join("references", "style.txt"), filepath.Join("scripts", "prepare.sh")} { - if _, err := os.Stat(filepath.Join(destination, "release-notes", file)); err != nil { - t.Fatalf("imported fixture file %q: %v", file, err) - } - } - - result, err = importSkillsFromDir("codex", source, destination) - if err != nil { - t.Fatal(err) - } - if got, want := strings.Join(result.Existing, ","), "release-notes"; got != want { - t.Fatalf("existing = %q, want %q", got, want) - } - if len(result.Imported) != 0 { - t.Fatalf("repeated import copied skills: %#v", result.Imported) - } -} - -func TestImportSkillsLeavesConflictsAndUnsafeSourcesUntouched(t *testing.T) { - source := t.TempDir() - destination := t.TempDir() - writeCatalogSkill(t, source, "release-notes", "source instructions") - writeCatalogSkill(t, destination, "release-notes", "existing instructions") - writeCatalogSkill(t, source, "nested-link", "safe manifest") - if err := os.Symlink(filepath.Join(source, "release-notes", skillFilename), filepath.Join(source, "nested-link", "reference")); err != nil { - t.Skipf("symlink not supported: %v", err) - } - if err := os.Symlink(filepath.Join(source, "release-notes"), filepath.Join(source, "linked-skill")); err != nil { - t.Skipf("symlink not supported: %v", err) - } - - result, err := importSkillsFromDir("codex", source, destination) - if err != nil { - t.Fatal(err) - } - if len(result.Imported) != 0 || len(result.Existing) != 0 { - t.Fatalf("unexpected successful import: %#v", result) - } - if got, err := os.ReadFile(filepath.Join(destination, "release-notes", skillFilename)); err != nil || !strings.Contains(string(got), "existing instructions") { - t.Fatalf("conflicting destination changed: %q, %v", got, err) - } - failed := make(map[string]bool) - for _, failure := range result.Failures { - failed[failure.Name] = true - } - for _, name := range []string{"release-notes", "nested-link", "linked-skill"} { - if !failed[name] { - t.Fatalf("missing failure for %q: %#v", name, result.Failures) - } - } -} - -func TestImportSkillsRejectsSymlinkedRoot(t *testing.T) { - root := t.TempDir() - source := filepath.Join(t.TempDir(), "codex-skills") - if err := os.Symlink(root, source); err != nil { - t.Skipf("symlink not supported: %v", err) - } - result, err := importSkillsFromDir("codex", source, t.TempDir()) - if err == nil || !strings.Contains(err.Error(), "symlinks are not supported") { - t.Fatalf("symlinked root error = %v", err) - } - if len(result.Imported) != 0 || len(result.Existing) != 0 || len(result.Failures) != 0 { - t.Fatalf("symlinked root result = %#v", result) - } -} - -func TestImportSkillsMissingRootAndConfiguredRoots(t *testing.T) { - result, err := importSkillsFromDir("codex", filepath.Join(t.TempDir(), "missing"), t.TempDir()) - if err != nil { - t.Fatal(err) - } - if len(result.Imported) != 0 || len(result.Existing) != 0 || len(result.Failures) != 0 { - t.Fatalf("missing root result = %#v", result) - } - - destination := t.TempDir() - rootBase := t.TempDir() - roots := map[string]string{ - "codex": filepath.Join(rootBase, "codex"), - "claude": filepath.Join(rootBase, "claude"), - "pi": filepath.Join(rootBase, "pi"), - } - for _, test := range []struct { - source string - root string - name string - }{ - {source: "codex", root: roots["codex"], name: "from-codex"}, - {source: "claude", root: roots["claude"], name: "from-claude"}, - {source: "pi", root: roots["pi"], name: "from-pi"}, - } { - t.Run(test.source, func(t *testing.T) { - writeCatalogSkill(t, test.root, test.name, "from "+test.source) - result, err = importSkillsFromRoots(test.source, roots, destination) - if err != nil { - t.Fatal(err) - } - if result.SourceDir != test.root { - t.Fatalf("source dir = %q, want %q", result.SourceDir, test.root) - } - if _, err := os.Stat(filepath.Join(destination, test.name, skillFilename)); err != nil { - t.Fatalf("conventional source was not imported: %v", err) - } - }) - } - if _, err := importSkillsFromRoots("unknown", roots, destination); err == nil || !strings.Contains(err.Error(), "unknown skill source") { - t.Fatalf("unknown source error = %v", err) - } -} - -func TestConventionalSkillImportRoots(t *testing.T) { - home := t.TempDir() - roots := conventionalSkillImportRoots(home) - for source, want := range map[string]string{ - "codex": filepath.Join(home, ".codex", "skills"), - "claude": filepath.Join(home, ".claude", "skills"), - "pi": filepath.Join(home, ".pi", "agent", "skills"), - } { - if got := roots[source]; got != want { - t.Fatalf("%s root = %q, want %q", source, got, want) - } - } -} - -func TestImportSkillsRejectsUnreadableManifest(t *testing.T) { - source := t.TempDir() - writeCatalogSkill(t, source, "private", "do not read") - manifest := filepath.Join(source, "private", skillFilename) - if err := os.Chmod(manifest, 0); err != nil { - t.Fatal(err) - } - t.Cleanup(func() { _ = os.Chmod(manifest, 0o644) }) - if _, err := os.ReadFile(manifest); err == nil { - t.Skip("test user can read a mode-000 file") - } - result, err := importSkillsFromDir("codex", source, t.TempDir()) - if err != nil { - t.Fatal(err) - } - if len(result.Failures) != 1 || result.Failures[0].Name != "private" { - t.Fatalf("failures = %#v", result.Failures) - } -} diff --git a/agent/testdata/import/release-notes/SKILL.md b/agent/testdata/import/release-notes/SKILL.md deleted file mode 100644 index dae33c01526..00000000000 --- a/agent/testdata/import/release-notes/SKILL.md +++ /dev/null @@ -1,8 +0,0 @@ ---- -name: release-notes -description: Draft concise release notes. ---- - -# Release notes - -Use short bullets. diff --git a/agent/tools/bash.go b/agent/tools/bash.go deleted file mode 100644 index 733629b515f..00000000000 --- a/agent/tools/bash.go +++ /dev/null @@ -1,450 +0,0 @@ -package tools - -import ( - "context" - "errors" - "fmt" - "os" - "os/exec" - "path/filepath" - "runtime" - "strings" - "time" - "unicode/utf8" - - "github.com/ollama/ollama/agent" - "github.com/ollama/ollama/api" -) - -const ( - bashTimeout = 3 * time.Minute - bashWaitDelay = 1 * time.Second - maxBashOutputBytes = 60_000 -) - -type Bash struct{} - -func (b *Bash) Name() string { - return shellToolName() -} - -func (b *Bash) Description() string { - return shellToolDescription() -} - -func (b *Bash) Schema() api.ToolFunction { - props := api.NewToolPropertiesMap() - props.Set("command", api.ToolProperty{ - Type: api.PropertyType{"string"}, - Description: shellCommandDescription(), - }) - return api.ToolFunction{ - Name: b.Name(), - Description: b.Description(), - Parameters: api.ToolFunctionParameters{ - Type: "object", - Properties: props, - Required: []string{"command"}, - }, - } -} - -func (b *Bash) RequiresApproval(map[string]any) bool { - return true -} - -// ApprovalScope scopes shell approval to the exact, trimmed command string -// using a NUL separator: "\x00". "Always allow this command" -// matches ONLY that precise string — any whitespace, quoting, or casing -// variant re-prompts. The NUL separator is safe because a shell command -// string cannot contain a literal NUL. -func (b *Bash) ApprovalScope(args map[string]any) string { - name := b.Name() - if command, ok := args["command"].(string); ok { - command = strings.TrimSpace(command) - if command != "" { - return name + "\x00" + command - } - } - return name -} - -func (b *Bash) Execute(ctx context.Context, toolCtx agent.ToolContext, args map[string]any) (agent.ToolResult, error) { - // TODO: use shared agent.RequiredStringArg for the "command" parameter (see agent package cleanup plan). - command, ok := args["command"].(string) - if !ok || strings.TrimSpace(command) == "" { - return agent.ToolResult{}, fmt.Errorf("command parameter is required") - } - if err := rejectUnsafeShellCommand(command); err != nil { - return agent.ToolResult{}, err - } - - ctx, cancel := context.WithTimeout(ctx, bashTimeout) - defer cancel() - - cwdFile, err := os.CreateTemp("", "ollama-agent-cwd-*") - if err != nil { - return agent.ToolResult{}, err - } - cwdPath := cwdFile.Name() - _ = cwdFile.Close() - defer os.Remove(cwdPath) - - cmd := newBashCommand(ctx, command, cwdPath) - cmd.WaitDelay = bashWaitDelay - cmd.Cancel = func() error { - return killBashCommand(cmd) - } - if toolCtx.WorkingDir != "" { - cmd.Dir = toolCtx.WorkingDir - } - - var stdout, stderr boundedOutput - stdout.Limit = maxBashOutputBytes - stderr.Limit = maxBashOutputBytes - cmd.Stdout = &stdout - cmd.Stderr = &stderr - - err = runBashCommand(cmd) - finalWorkingDir := readFinalWorkingDir(cwdPath) - - var sb strings.Builder - if stdout.Len() > 0 { - sb.WriteString(stdout.String("stdout")) - } - if stderr.Len() > 0 { - if sb.Len() > 0 { - sb.WriteString("\n") - } - sb.WriteString("stderr:\n") - sb.WriteString(stderr.String("stderr")) - } - - if err != nil { - if ctx.Err() == context.DeadlineExceeded { - return agent.ToolResult{Content: bashContentWithError(sb.String(), "Error: command timed out after "+bashTimeout.String()), WorkingDir: finalWorkingDir}, nil - } - if ctx.Err() == context.Canceled { - return agent.ToolResult{Content: bashContentWithError(sb.String(), "Error: command was canceled"), WorkingDir: finalWorkingDir}, nil - } - if errors.Is(err, exec.ErrWaitDelay) { - _ = killBashCommand(cmd) - return agent.ToolResult{Content: bashContentWithError(sb.String(), "Error: command output pipes did not close after "+bashWaitDelay.String()), WorkingDir: finalWorkingDir}, nil - } - if exitErr, ok := err.(*exec.ExitError); ok { - return agent.ToolResult{Content: bashContentWithError(sb.String(), fmt.Sprintf("Exit code: %d", exitErr.ExitCode())), WorkingDir: finalWorkingDir}, nil - } - return agent.ToolResult{Content: sb.String(), WorkingDir: finalWorkingDir}, fmt.Errorf("executing command: %w", err) - } - - if sb.Len() == 0 { - return agent.ToolResult{Content: "(no output)", WorkingDir: finalWorkingDir}, nil - } - return agent.ToolResult{Content: sb.String(), WorkingDir: finalWorkingDir}, nil -} - -func bashContentWithError(content, msg string) string { - if content == "" { - return msg - } - return content + "\n\n" + msg -} - -// rejectUnsafeShellCommand applies a best-effort blocklist for obviously -// destructive or credential-exfiltrating commands. It is defense-in-depth -// ONLY: the interactive approval prompt is the real security control, and -// this check must not be relied upon as a sandbox. Sophisticated or novel -// dangerous commands (e.g. find / -delete, dd, fork bombs, custom binaries) -// are NOT caught here and will simply be routed through approval like any -// other command. Keep the approval prompt as the gate. -func rejectUnsafeShellCommand(command string) error { - switch { - case hasUnsafeRecursiveDelete(command): - return fmt.Errorf("refusing to run unsafe command: recursive delete target is too broad") - case readsCredentialPath(command): - return fmt.Errorf("refusing to run unsafe command: credential file reads are not allowed") - default: - return nil - } -} - -func hasUnsafeRecursiveDelete(command string) bool { - // Check each command segment independently. shellSafetyText flattens - // separators (; & | newlines) to spaces, which would otherwise let the - // rm target scan bleed across command boundaries — e.g. - // "rm -rf build && echo ~/.ssh/config" flattened to one token stream - // would treat the unrelated ~/.ssh/config (a ~/-prefixed "unsafe - // target") as an rm argument. Splitting on separators first restores - // command boundaries while still catching multi-target single commands - // like "rm -rf build /etc". - for _, segment := range shellSegments(command) { - fields := shellSafetyFields(segment) - for i, field := range fields { - if isRMCommand(field) && rmCommandDeletesUnsafeTarget(fields[i+1:]) { - return true - } - if isPowerShellDeleteCommand(field) && powerShellDeleteCommandDeletesUnsafeTarget(fields[i+1:]) { - return true - } - } - } - return false -} - -// shellSegments splits a command on shell control operators (;, &, |, &&, -// ||) and newlines, returning the individual command segments. It operates on -// the lowercased raw command before quote/separator normalization so that -// command boundaries are preserved for per-segment checks. Subshell parens are -// intentionally NOT treated as separators: splitting on them would fragment -// command substitutions like "rm -rf $(echo /)" into "rm -rf $" and "echo /", -// hiding the destructive "/" target from the per-segment scan. Empty segments -// are dropped. -func shellSegments(command string) []string { - command = strings.ToLower(command) - var segments []string - for _, segment := range strings.FieldsFunc(command, func(r rune) bool { - switch r { - case ';', '&', '|', '\n', '\r': - return true - } - return false - }) { - if segment = strings.TrimSpace(segment); segment != "" { - segments = append(segments, segment) - } - } - return segments -} - -func rmCommandDeletesUnsafeTarget(fields []string) bool { - var flags string - for _, field := range fields { - if field == "--" { - continue - } - if strings.HasPrefix(field, "-") { - flags += field - continue - } - if strings.Contains(flags, "r") && strings.Contains(flags, "f") && isUnsafeDeleteTarget(field) { - return true - } - } - return false -} - -func powerShellDeleteCommandDeletesUnsafeTarget(fields []string) bool { - var recurse, force bool - var targets []string - for _, field := range fields { - switch field { - case "-r", "-recurse", "-recursive": - recurse = true - case "-f", "-force": - force = true - default: - if !strings.HasPrefix(field, "-") { - targets = append(targets, field) - } - } - } - if !recurse || !force { - return false - } - for _, target := range targets { - if isUnsafeDeleteTarget(target) { - return true - } - } - return false -} - -func readsCredentialPath(command string) bool { - fields := shellSafetyFields(command) - if !hasCredentialReadVerb(fields) { - return false - } - normalized := shellSafetyText(command) - for _, fragment := range []string{ - "/.ssh/id_rsa", - "/.ssh/id_dsa", - "/.ssh/id_ecdsa", - "/.ssh/id_ed25519", - "/.ssh/config", - "/.ssh/known_hosts", - "/.aws/credentials", - "/.aws/config", - "/.config/gcloud/application_default_credentials.json", - "/.kube/config", - "/.netrc", - "/.npmrc", - "/.docker/config.json", - "/.config/gh/hosts.yml", - "/.gnupg/", - "/etc/shadow", - } { - if strings.Contains(normalized, fragment) { - return true - } - } - return false -} - -func hasCredentialReadVerb(fields []string) bool { - for _, field := range fields { - switch field { - case "cat", "less", "more", "head", "tail", "type", "get-content", "gc", "select-string", "grep", "rg", "sed", "awk": - return true - case "env", "printenv": - return true - } - } - return false -} - -func isRMCommand(field string) bool { - return field == "rm" || strings.HasSuffix(field, "/rm") -} - -func isPowerShellDeleteCommand(field string) bool { - switch field { - case "remove-item", "del", "erase", "rd", "rmdir": - return true - default: - return false - } -} - -func isUnsafeDeleteTarget(target string) bool { - if target == "." || target == "./" || target == "*" { - return true - } - if target == "/*" { - return true - } - target = strings.TrimSuffix(target, "/*") - for _, prefix := range []string{"~/", "$home/", "${home}/", "$env:home/", "$env:userprofile/", "%userprofile%/"} { - if strings.HasPrefix(target, prefix) { - return true - } - } - for _, prefix := range []string{"/etc/", "/bin/", "/sbin/", "/usr/", "/var/", "/lib/", "/library/", "/system/", "/applications/", "c:/windows/", "c:/program files/"} { - if strings.HasPrefix(target, prefix) { - return true - } - } - for _, exact := range []string{"/", "~", "$home", "${home}", "$env:home", "$env:userprofile", "%userprofile%", "c:", "c:/", "/etc", "/bin", "/sbin", "/usr", "/var", "/lib", "/library", "/system", "/applications", "c:/windows", "c:/program files"} { - if target == exact { - return true - } - } - return false -} - -func shellSafetyFields(command string) []string { - return strings.Fields(shellSafetyText(command)) -} - -func shellSafetyText(command string) string { - command = strings.ToLower(command) - return strings.NewReplacer( - "\\", "/", - "\n", " ", - "\t", " ", - ";", " ", - "&", " ", - "|", " ", - "(", " ", - ")", " ", - "\"", "", - "'", "", - "`", "", - ).Replace(command) -} - -func readFinalWorkingDir(path string) string { - content, err := os.ReadFile(path) - if err != nil { - return "" - } - workingDir := strings.TrimPrefix(string(content), "\ufeff") - workingDir = strings.TrimSpace(workingDir) - if workingDir == "" { - return "" - } - workingDir = normalizeBashWorkingDir(workingDir) - info, err := os.Stat(workingDir) - if err != nil || !info.IsDir() { - return "" - } - return workingDir -} - -func normalizeBashWorkingDir(workingDir string) string { - if runtime.GOOS == "windows" && len(workingDir) >= 3 && workingDir[0] == '/' && workingDir[2] == '/' && isASCIIAlpha(workingDir[1]) { - workingDir = strings.ToUpper(string(workingDir[1])) + ":" + workingDir[2:] - } - workingDir = filepath.Clean(filepath.FromSlash(workingDir)) - if runtime.GOOS == "windows" && len(workingDir) >= 2 && workingDir[1] == ':' && isASCIIAlpha(workingDir[0]) { - workingDir = strings.ToUpper(string(workingDir[0])) + workingDir[1:] - } - return workingDir -} - -func isASCIIAlpha(b byte) bool { - return (b >= 'a' && b <= 'z') || (b >= 'A' && b <= 'Z') -} - -type boundedOutput struct { - Limit int - buf []byte - omitted int -} - -func (b *boundedOutput) Write(p []byte) (int, error) { - if b.Limit <= 0 { - b.omitted += len(p) - return len(p), nil - } - remaining := b.Limit - len(b.buf) - if remaining <= 0 { - b.omitted += len(p) - return len(p), nil - } - if len(p) <= remaining { - b.buf = append(b.buf, p...) - return len(p), nil - } - writeLen := utf8SafePrefixLen(p[:remaining]) - b.buf = append(b.buf, p[:writeLen]...) - b.omitted += len(p) - writeLen - return len(p), nil -} - -func (b *boundedOutput) Len() int { - return len(b.buf) + b.omitted -} - -func (b *boundedOutput) String(label string) string { - safeLen := utf8SafePrefixLen(b.buf) - content := string(b.buf[:safeLen]) - omitted := b.omitted + len(b.buf) - safeLen - if omitted == 0 { - return content - } - return content + agent.TruncMarker(label, safeLen, 0, omitted, false, "") -} - -func utf8SafePrefixLen(p []byte) int { - if len(p) == 0 { - return 0 - } - for i := 0; i < len(p); { - r, size := utf8.DecodeRune(p[i:]) - if r == utf8.RuneError && size == 1 { - return i - } - i += size - } - return len(p) -} diff --git a/agent/tools/bash_test.go b/agent/tools/bash_test.go deleted file mode 100644 index 537ac98a961..00000000000 --- a/agent/tools/bash_test.go +++ /dev/null @@ -1,258 +0,0 @@ -package tools - -import ( - "context" - "os" - "path/filepath" - "runtime" - "strings" - "testing" - "unicode/utf8" - - "github.com/ollama/ollama/agent" -) - -func TestBashReportsFinalWorkingDir(t *testing.T) { - root := t.TempDir() - subdir := filepath.Join(root, "sub") - if err := os.Mkdir(subdir, 0o755); err != nil { - t.Fatal(err) - } - - result, err := (&Bash{}).Execute(context.Background(), agent.ToolContext{WorkingDir: root}, map[string]any{ - "command": shellTestCommand("cd sub && pwd", "Set-Location sub; Get-Location"), - }) - if err != nil { - t.Fatal(err) - } - wantDir, err := filepath.EvalSymlinks(subdir) - if err != nil { - t.Fatal(err) - } - if result.WorkingDir != wantDir { - t.Fatalf("working dir = %q, want %q", result.WorkingDir, wantDir) - } - if !strings.Contains(result.Content, "sub") { - t.Fatalf("content = %q, want pwd output", result.Content) - } -} - -func TestBashBoundsOutputWhileRunning(t *testing.T) { - result, err := (&Bash{}).Execute(context.Background(), agent.ToolContext{WorkingDir: t.TempDir()}, map[string]any{ - "command": shellTestCommand("yes x | head -c 70000", "[Console]::Out.Write(('x' * 70000))"), - }) - if err != nil { - t.Fatal(err) - } - if !strings.Contains(result.Content, "[stdout truncated: showing first ~") || !strings.Contains(result.Content, "omitted ~") || !strings.Contains(result.Content, " tokens.]") { - t.Fatalf("content = %q, want stdout truncation marker", result.Content) - } - if count, want := strings.Count(result.Content, "x"), shellTestCapturedXCount(); count != want { - t.Fatalf("captured x count = %d, want %d", count, want) - } - if len(result.Content) > maxBashOutputBytes+200 { - t.Fatalf("content length = %d, want bounded output", len(result.Content)) - } -} - -func TestBoundedOutputTruncatesAtUTF8Boundary(t *testing.T) { - var out boundedOutput - out.Limit = len([]byte("abc")) + 1 - - if _, err := out.Write([]byte("abcédef")); err != nil { - t.Fatal(err) - } - content := out.String("stdout") - if !utf8.ValidString(content) { - t.Fatalf("content is not valid UTF-8: %q", content) - } - if strings.ContainsRune(content, utf8.RuneError) { - t.Fatalf("content contains replacement rune: %q", content) - } - if !strings.HasPrefix(content, "abc\n\n[stdout truncated:") { - t.Fatalf("content = %q, want complete ASCII prefix and truncation marker", content) - } -} - -func TestBoundedOutputKeepsCompleteUTF8AtBoundary(t *testing.T) { - var out boundedOutput - out.Limit = len([]byte("abcé")) - - if _, err := out.Write([]byte("abcédef")); err != nil { - t.Fatal(err) - } - if content := out.String("stdout"); !strings.HasPrefix(content, "abcé\n\n[stdout truncated:") { - t.Fatalf("content = %q, want complete UTF-8 prefix", content) - } -} - -func TestBoundedOutputTrimsTrailingPartialUTF8(t *testing.T) { - var out boundedOutput - out.Limit = 4 - - if _, err := out.Write([]byte{'a', 'b', 'c', 0xc3}); err != nil { - t.Fatal(err) - } - if _, err := out.Write([]byte{0xa9}); err != nil { - t.Fatal(err) - } - if content := out.String("stdout"); !utf8.ValidString(content) || !strings.HasPrefix(content, "abc\n\n[stdout truncated:") { - t.Fatalf("content = %q, want valid UTF-8 with partial suffix trimmed", content) - } -} - -func TestUTF8SafePrefixRejectsMalformedLeadByte(t *testing.T) { - input := []byte{'a', 0xc0, 0x80, 'b'} - if got := utf8SafePrefixLen(input); got != 1 { - t.Fatalf("safe prefix length = %d, want 1", got) - } -} - -func TestBoundedOutputDropsMalformedUTF8(t *testing.T) { - var out boundedOutput - out.Limit = 4 - - if _, err := out.Write([]byte{'a', 0xc0, 0x80, 'b'}); err != nil { - t.Fatal(err) - } - content := out.String("stdout") - if !utf8.ValidString(content) { - t.Fatalf("content is not valid UTF-8: %q", content) - } - if strings.ContainsRune(content, utf8.RuneError) { - t.Fatalf("content contains replacement rune: %q", content) - } - if !strings.HasPrefix(content, "a\n\n[stdout truncated:") { - t.Fatalf("content = %q, want valid prefix and truncation marker", content) - } -} - -func TestBashReportsCanceledCommand(t *testing.T) { - ctx, cancel := context.WithCancel(context.Background()) - cancel() - - result, err := (&Bash{}).Execute(ctx, agent.ToolContext{WorkingDir: t.TempDir()}, map[string]any{ - "command": shellTestCommand("sleep 10", "Start-Sleep -Seconds 10"), - }) - if err != nil { - t.Fatal(err) - } - if !strings.Contains(result.Content, "Error: command was canceled") { - t.Fatalf("content = %q, want canceled message", result.Content) - } - if strings.Contains(result.Content, "Exit code: -1") { - t.Fatalf("content = %q, should not mask cancellation as exit code", result.Content) - } -} - -func TestRejectUnsafeShellCommand(t *testing.T) { - tests := []struct { - name string - command string - wantErr bool - }{ - {name: "rm root", command: "rm -rf /", wantErr: true}, - {name: "sudo rm root", command: "sudo rm -rf -- /", wantErr: true}, - {name: "rm home", command: "rm -fr $HOME", wantErr: true}, - {name: "rm root wildcard", command: "rm -rf /*", wantErr: true}, - {name: "rm system subdir", command: "rm -rf /etc/ssh", wantErr: true}, - {name: "rm cwd", command: "rm -rf .", wantErr: true}, - {name: "powershell remove root", command: `Remove-Item -Recurse -Force C:\`, wantErr: true}, - {name: "powershell remove system subdir", command: `Remove-Item -Recurse -Force C:\Windows\Temp`, wantErr: true}, - {name: "ssh private key", command: "cat ~/.ssh/id_rsa", wantErr: true}, - {name: "aws credentials", command: "Get-Content $HOME/.aws/credentials", wantErr: true}, - {name: "shadow", command: "head /etc/shadow", wantErr: true}, - {name: "netrc", command: "cat ~/.netrc", wantErr: true}, - {name: "docker config", command: "cat ~/.docker/config.json", wantErr: true}, - {name: "gnupg dir", command: "cat ~/.gnupg/private-keys-v1.d/key", wantErr: true}, - {name: "gh hosts", command: "cat ~/.config/gh/hosts.yml", wantErr: true}, - {name: "ssh config", command: "cat ~/.ssh/config", wantErr: true}, - {name: "printenv dump", command: "printenv", wantErr: false}, - {name: "delete build dir", command: "rm -rf build", wantErr: false}, - {name: "read project file", command: "cat README.md", wantErr: false}, - {name: "mention key text", command: "rg id_rsa docs", wantErr: false}, - {name: "env example", command: "cat .env.example", wantErr: false}, - {name: "rm build then unrelated tilde path", command: "rm -rf build && echo ~/.ssh/config", wantErr: false}, - {name: "rm build then unrelated slash path", command: "rm -rf build; cat /etc/passwd", wantErr: false}, - {name: "rm build then unrelated star glob", command: "rm -rf build && ls *.go", wantErr: false}, - {name: "rm multiple targets one unsafe", command: "rm -rf build /etc", wantErr: true}, - {name: "rm unsafe then safe piped", command: "rm -rf / | tee log", wantErr: true}, - {name: "rm unsafe via command substitution", command: "rm -rf $(echo /)", wantErr: true}, - } - - for _, tt := range tests { - t.Run(tt.name, func(t *testing.T) { - err := rejectUnsafeShellCommand(tt.command) - if tt.wantErr && err == nil { - t.Fatal("expected unsafe command to be rejected") - } - if !tt.wantErr && err != nil { - t.Fatalf("command rejected: %v", err) - } - }) - } -} - -func TestBashRejectsUnsafeCommandBeforeExecution(t *testing.T) { - _, err := (&Bash{}).Execute(context.Background(), agent.ToolContext{WorkingDir: t.TempDir()}, map[string]any{ - "command": "rm -rf /", - }) - if err == nil || !strings.Contains(err.Error(), "refusing to run unsafe command") { - t.Fatalf("err = %v, want unsafe command rejection", err) - } -} - -func shellTestCommand(unix, windows string) string { - if runtime.GOOS == "windows" { - return windows - } - return unix -} - -func shellTestCapturedXCount() int { - if runtime.GOOS == "windows" { - return maxBashOutputBytes - } - return maxBashOutputBytes / 2 -} - -func TestReadFinalWorkingDirRejectsInvalidPaths(t *testing.T) { - dir := t.TempDir() - cwdFile := filepath.Join(dir, "cwd") - notDir := filepath.Join(dir, "file.txt") - if err := os.WriteFile(notDir, []byte("not a dir"), 0o644); err != nil { - t.Fatal(err) - } - - if err := os.WriteFile(cwdFile, []byte(notDir+"\n"), 0o644); err != nil { - t.Fatal(err) - } - if got := readFinalWorkingDir(cwdFile); got != "" { - t.Fatalf("regular file cwd = %q, want empty", got) - } - - if err := os.WriteFile(cwdFile, []byte(filepath.Join(dir, "missing")+"\n"), 0o644); err != nil { - t.Fatal(err) - } - if got := readFinalWorkingDir(cwdFile); got != "" { - t.Fatalf("missing cwd = %q, want empty", got) - } - - if err := os.WriteFile(cwdFile, []byte(dir+"\n"), 0o644); err != nil { - t.Fatal(err) - } - if got := readFinalWorkingDir(cwdFile); got != dir { - t.Fatalf("directory cwd = %q, want %q", got, dir) - } -} - -func TestNormalizeBashWorkingDirWindowsDriveLetter(t *testing.T) { - if runtime.GOOS != "windows" { - t.Skip("windows path normalization") - } - got := normalizeBashWorkingDir("/c/Users/jdoe/project") - want := filepath.Clean(`C:\Users\jdoe\project`) - if got != want { - t.Fatalf("working dir = %q, want %q", got, want) - } -} diff --git a/agent/tools/bash_unix.go b/agent/tools/bash_unix.go deleted file mode 100644 index 10683e8d00e..00000000000 --- a/agent/tools/bash_unix.go +++ /dev/null @@ -1,49 +0,0 @@ -//go:build !windows - -package tools - -import ( - "context" - "os/exec" - "strings" - "syscall" -) - -func shellToolName() string { - return "bash" -} - -func shellToolDescription() string { - return "Execute a bash command on the system. Use this to inspect files, run tests, and perform development tasks." -} - -func shellCommandDescription() string { - return "The bash command to execute." -} - -func newBashCommand(ctx context.Context, command, cwdPath string) *exec.Cmd { - script := command + "\n__ollama_status=$?\npwd -P > " + shellQuote(cwdPath) + "\nexit $__ollama_status" - cmd := exec.CommandContext(ctx, "bash", "-c", script) - configureBashCommand(cmd) - return cmd -} - -func shellQuote(value string) string { - return "'" + strings.ReplaceAll(value, "'", "'\\''") + "'" -} - -func configureBashCommand(cmd *exec.Cmd) { - cmd.SysProcAttr = &syscall.SysProcAttr{Setpgid: true} -} - -func runBashCommand(cmd *exec.Cmd) error { - return cmd.Run() -} - -func killBashCommand(cmd *exec.Cmd) error { - if cmd == nil || cmd.Process == nil { - return nil - } - _ = syscall.Kill(-cmd.Process.Pid, syscall.SIGKILL) - return nil -} diff --git a/agent/tools/bash_unix_test.go b/agent/tools/bash_unix_test.go deleted file mode 100644 index 5c7992adb85..00000000000 --- a/agent/tools/bash_unix_test.go +++ /dev/null @@ -1,40 +0,0 @@ -//go:build !windows - -package tools - -import ( - "context" - "os/exec" - "strings" - "testing" - "time" - - "github.com/ollama/ollama/agent" -) - -func TestConfigureBashCommandSetsProcessGroup(t *testing.T) { - cmd := exec.Command("bash", "-c", "true") - configureBashCommand(cmd) - if cmd.SysProcAttr == nil || !cmd.SysProcAttr.Setpgid { - t.Fatalf("configureBashCommand should start bash in a new process group") - } -} - -func TestBashWaitDelayBoundsBackgroundOutputPipe(t *testing.T) { - start := time.Now() - result, err := (&Bash{}).Execute(context.Background(), agent.ToolContext{WorkingDir: t.TempDir()}, map[string]any{ - "command": "sleep 5 & echo done", - }) - if err != nil { - t.Fatal(err) - } - if elapsed := time.Since(start); elapsed > bashWaitDelay+2*time.Second { - t.Fatalf("command elapsed = %s, want bounded near %s", elapsed, bashWaitDelay) - } - if !strings.Contains(result.Content, "done") { - t.Fatalf("content = %q, want command output", result.Content) - } - if !strings.Contains(result.Content, "output pipes did not close") { - t.Fatalf("content = %q, want wait delay message", result.Content) - } -} diff --git a/agent/tools/bash_windows.go b/agent/tools/bash_windows.go deleted file mode 100644 index 840ff24fca2..00000000000 --- a/agent/tools/bash_windows.go +++ /dev/null @@ -1,134 +0,0 @@ -//go:build windows - -package tools - -import ( - "context" - "os/exec" - "strings" - "sync" - "unsafe" - - "golang.org/x/sys/windows" -) - -var bashJobHandles sync.Map - -func shellToolName() string { - return "powershell" -} - -func shellToolDescription() string { - return "Execute a PowerShell command on the system. Use this to inspect files, run tests, and perform development tasks." -} - -func shellCommandDescription() string { - return "The PowerShell command to execute." -} - -func newBashCommand(ctx context.Context, command, cwdPath string) *exec.Cmd { - return exec.CommandContext( - ctx, - "powershell.exe", - "-NoLogo", - "-NoProfile", - "-NonInteractive", - "-ExecutionPolicy", - "Bypass", - "-Command", - powerShellCommandScript(command, cwdPath), - ) -} - -func powerShellCommandScript(command, cwdPath string) string { - cwdPath = powerShellSingleQuote(cwdPath) - return strings.Join([]string{ - "$__ollama_status = 0", - ". {", - "try {", - command, - " $__ollama_success = $?", - " $__ollama_last_exit = $global:LASTEXITCODE", - " if ($__ollama_success) {", - " $__ollama_status = 0", - " } elseif ($__ollama_last_exit -is [int] -and $__ollama_last_exit -ne 0) {", - " $__ollama_status = $__ollama_last_exit", - " } else {", - " $__ollama_status = 1", - " }", - "} catch {", - " Write-Error $_", - " $__ollama_status = 1", - "} finally {", - " try { [System.IO.File]::WriteAllText(" + cwdPath + ", (Get-Location).ProviderPath, [System.Text.Encoding]::UTF8) } catch {}", - "}", - "} | Out-String -Stream -Width 4096", - "exit $__ollama_status", - }, "\n") -} - -func powerShellSingleQuote(value string) string { - return "'" + strings.ReplaceAll(value, "'", "''") + "'" -} - -func runBashCommand(cmd *exec.Cmd) error { - if err := cmd.Start(); err != nil { - return err - } - if job, err := createBashJob(cmd.Process.Pid); err == nil { - bashJobHandles.Store(cmd.Process.Pid, job) - defer releaseBashJob(cmd.Process.Pid) - } - return cmd.Wait() -} - -func killBashCommand(cmd *exec.Cmd) error { - if cmd == nil || cmd.Process == nil { - return nil - } - releaseBashJob(cmd.Process.Pid) - _ = cmd.Process.Kill() - return nil -} - -func createBashJob(pid int) (windows.Handle, error) { - job, err := windows.CreateJobObject(nil, nil) - if err != nil { - return 0, err - } - - info := windows.JOBOBJECT_EXTENDED_LIMIT_INFORMATION{} - info.BasicLimitInformation.LimitFlags = windows.JOB_OBJECT_LIMIT_KILL_ON_JOB_CLOSE - if _, err := windows.SetInformationJobObject( - job, - windows.JobObjectExtendedLimitInformation, - uintptr(unsafe.Pointer(&info)), - uint32(unsafe.Sizeof(info)), - ); err != nil { - _ = windows.CloseHandle(job) - return 0, err - } - - process, err := windows.OpenProcess(windows.PROCESS_SET_QUOTA|windows.PROCESS_TERMINATE, false, uint32(pid)) - if err != nil { - _ = windows.CloseHandle(job) - return 0, err - } - defer windows.CloseHandle(process) - - if err := windows.AssignProcessToJobObject(job, process); err != nil { - _ = windows.CloseHandle(job) - return 0, err - } - return job, nil -} - -func releaseBashJob(pid int) { - value, ok := bashJobHandles.LoadAndDelete(pid) - if !ok { - return - } - if job, ok := value.(windows.Handle); ok { - _ = windows.CloseHandle(job) - } -} diff --git a/agent/tools/bash_windows_test.go b/agent/tools/bash_windows_test.go deleted file mode 100644 index 9a9be4c23e3..00000000000 --- a/agent/tools/bash_windows_test.go +++ /dev/null @@ -1,15 +0,0 @@ -//go:build windows - -package tools - -import ( - "strings" - "testing" -) - -func TestPowerShellCommandScriptUsesWideOutString(t *testing.T) { - script := powerShellCommandScript("Get-ChildItem", `C:\cwd.txt`) - if !strings.Contains(script, "Out-String -Stream -Width 4096") { - t.Fatalf("script = %q, want explicit Out-String width", script) - } -} diff --git a/agent/tools/file.go b/agent/tools/file.go deleted file mode 100644 index b12a1c0e00a..00000000000 --- a/agent/tools/file.go +++ /dev/null @@ -1,711 +0,0 @@ -package tools - -import ( - "bufio" - "cmp" - "context" - "encoding/json" - "fmt" - "io" - "os" - "path/filepath" - "slices" - "strconv" - "strings" - - "github.com/ollama/ollama/agent" - "github.com/ollama/ollama/api" -) - -const ( - maxReadBytes = 200000 -) - -type Read struct{} - -func (r *Read) Name() string { - return "read" -} - -func (r *Read) Description() string { - return "Read a text file from the current working directory." -} - -func (r *Read) Schema() api.ToolFunction { - props := api.NewToolPropertiesMap() - props.Set("path", api.ToolProperty{ - Type: api.PropertyType{"string"}, - Description: "Path to the file to read, relative to the working directory.", - }) - props.Set("start", api.ToolProperty{ - Type: api.PropertyType{"integer"}, - Description: "Optional 1-based line to start reading from.", - }) - props.Set("end", api.ToolProperty{ - Type: api.PropertyType{"integer"}, - Description: "Optional 1-based inclusive line to stop reading at.", - }) - return api.ToolFunction{ - Name: r.Name(), - Description: r.Description(), - Parameters: api.ToolFunctionParameters{ - Type: "object", - Properties: props, - Required: []string{"path"}, - }, - } -} - -func (r *Read) RequiresApproval(map[string]any) bool { - return true -} - -func (r *Read) Execute(ctx context.Context, toolCtx agent.ToolContext, args map[string]any) (agent.ToolResult, error) { - // TODO: use shared agent.RequiredStringArg / agent.OptionalIntArg for args (see agent package cleanup plan). - path, ok := args["path"].(string) - if !ok || strings.TrimSpace(path) == "" { - return agent.ToolResult{}, fmt.Errorf("path parameter is required") - } - - file, info, err := openRegularFile(toolCtx.WorkingDir, path, true) - if err != nil { - return agent.ToolResult{}, err - } - defer file.Close() - - selection, err := readSelectionFromArgs(args) - if err != nil { - return agent.ToolResult{}, err - } - if !selection.enabled && info.Size() > maxReadBytes { - return agent.ToolResult{}, fmt.Errorf("%s is too large to read (%d bytes)", path, info.Size()) - } - - select { - case <-ctx.Done(): - return agent.ToolResult{}, ctx.Err() - default: - } - - var content string - if selection.enabled { - content, err = readLineSelection(file, selection) - } else { - var contentBytes []byte - contentBytes, err = readAllWithinLimit(file, maxReadBytes) - content = string(contentBytes) - } - if err != nil { - return agent.ToolResult{}, err - } - return agent.ToolResult{Content: content}, nil -} - -type Edit struct{} - -func (e *Edit) Name() string { - return "edit" -} - -func (e *Edit) Description() string { - return "Edit a text file in the current working directory by replacing exact text. Pass multiple edits to change separate parts of the file in one call." -} - -func (e *Edit) Schema() api.ToolFunction { - editProps := api.NewToolPropertiesMap() - editProps.Set("old_text", api.ToolProperty{ - Type: api.PropertyType{"string"}, - Description: "Exact text for one targeted replacement. Must match the original file exactly once and must not overlap with any other edit's old_text.", - }) - editProps.Set("new_text", api.ToolProperty{ - Type: api.PropertyType{"string"}, - Description: "Replacement text for this targeted edit.", - }) - - props := api.NewToolPropertiesMap() - props.Set("path", api.ToolProperty{ - Type: api.PropertyType{"string"}, - Description: "Path to the file to edit, relative to the working directory.", - }) - props.Set("edits", api.ToolProperty{ - Type: api.PropertyType{"array"}, - Items: api.ToolProperty{ - Type: api.PropertyType{"object"}, - Properties: editProps, - Required: []string{"old_text", "new_text"}, - }, - Description: "One or more exact-text replacements. Each is matched against the original file, not against the output of earlier edits. Keep old_text as small as possible while still unique in the file; merge changes to the same or adjacent lines into a single edit.", - }) - props.Set("replace_all", api.ToolProperty{ - Type: api.PropertyType{"boolean"}, - Description: "Replace every occurrence. Defaults to false; only applies when a single edit is provided.", - }) - return api.ToolFunction{ - Name: e.Name(), - Description: e.Description(), - Parameters: api.ToolFunctionParameters{ - Type: "object", - Properties: props, - Required: []string{"path", "edits"}, - }, - } -} - -func (e *Edit) RequiresApproval(map[string]any) bool { - return true -} - -func (e *Edit) Execute(ctx context.Context, toolCtx agent.ToolContext, args map[string]any) (agent.ToolResult, error) { - // TODO: use shared agent.RequiredStringArg / agent.OptionalBoolArg for args (see agent package cleanup plan). - path, ok := args["path"].(string) - if !ok || strings.TrimSpace(path) == "" { - return agent.ToolResult{}, fmt.Errorf("path parameter is required") - } - - edits, replaceAll, err := parseEditArgs(args) - if err != nil { - return agent.ToolResult{}, err - } - - if err := rejectFinalSymlink(toolCtx.WorkingDir, path); err != nil { - return agent.ToolResult{}, err - } - - file, info, err := openRegularFile(toolCtx.WorkingDir, path, false) - if err != nil { - return agent.ToolResult{}, err - } - if info.Size() > maxReadBytes { - file.Close() - return agent.ToolResult{}, fmt.Errorf("%s is too large to edit (%d bytes)", path, info.Size()) - } - - select { - case <-ctx.Done(): - file.Close() - return agent.ToolResult{}, ctx.Err() - default: - } - - contentBytes, err := readAllWithinLimit(file, maxReadBytes) - if closeErr := file.Close(); err == nil && closeErr != nil { - err = closeErr - } - if err != nil { - return agent.ToolResult{}, err - } - content := string(contentBytes) - - var updated string - replacements := 0 - if replaceAll { - matches := strings.Count(content, edits[0].OldText) - if matches == 0 { - return agent.ToolResult{}, fmt.Errorf("old_text was not found in %s", path) - } - updated = strings.ReplaceAll(content, edits[0].OldText, edits[0].NewText) - replacements = matches - } else { - // Every edit is matched against the original file content rather - // than the output of earlier edits, so each edit must match exactly - // once and edits must target disjoint regions. - matched := make([]editMatch, 0, len(edits)) - for i, edit := range edits { - count := strings.Count(content, edit.OldText) - if count == 0 { - return agent.ToolResult{}, editNotFoundError(path, i, len(edits)) - } - if count > 1 { - return agent.ToolResult{}, editAmbiguousError(path, i, len(edits), count) - } - matched = append(matched, editMatch{ - editIndex: i, - offset: strings.Index(content, edit.OldText), - length: len(edit.OldText), - newText: edit.NewText, - }) - replacements++ - } - - slices.SortFunc(matched, func(a, b editMatch) int { return cmp.Compare(a.offset, b.offset) }) - for i := 1; i < len(matched); i++ { - prev, cur := matched[i-1], matched[i] - if prev.offset+prev.length > cur.offset { - return agent.ToolResult{}, fmt.Errorf("edits[%d] and edits[%d] overlap in %s; merge them into one edit or target disjoint text", prev.editIndex, cur.editIndex, path) - } - } - - // Apply from the end of the file backwards so earlier offsets stay valid. - updated = content - for i := len(matched) - 1; i >= 0; i-- { - m := matched[i] - updated = updated[:m.offset] + m.newText + updated[m.offset+m.length:] - } - } - - if updated == content { - return agent.ToolResult{}, fmt.Errorf("edit produced no changes in %s; replacement text is identical to the original", path) - } - if len(updated) > maxReadBytes { - return agent.ToolResult{}, fmt.Errorf("edited content is too large (%d bytes)", len(updated)) - } - - if err := writeFileAtomic(toolCtx.WorkingDir, path, []byte(updated), info.Mode().Perm()); err != nil { - return agent.ToolResult{}, err - } - - return agent.ToolResult{Content: fmt.Sprintf("Updated %s (%d edit%s, %d replacement%s).", path, len(edits), plural(len(edits)), replacements, plural(replacements))}, nil -} - -// editReplacement is one targeted replacement within an edit call. -type editReplacement struct { - OldText string - NewText string -} - -// editMatch locates one editReplacement within the original file content. -type editMatch struct { - editIndex int - offset int - length int - newText string -} - -// parseEditArgs normalizes edit arguments from a tool call into a list of -// replacements. It accepts the `edits` array form and tolerates legacy -// top-level old_text/new_text args as well as stringified JSON, mirroring -// the pi coding agent's argument handling. -func parseEditArgs(args map[string]any) ([]editReplacement, bool, error) { - replaceAll, _ := args["replace_all"].(bool) - - var edits []editReplacement - if raw, ok := args["edits"]; ok { - parsed, err := parseEditArray(raw) - if err != nil { - return nil, false, err - } - edits = parsed - } - - // Fold a legacy top-level old_text/new_text pair into edits. - if oldText, ok := args["old_text"].(string); ok { - newText, ok := args["new_text"].(string) - if !ok { - return nil, false, fmt.Errorf("new_text parameter is required") - } - edits = append(edits, editReplacement{OldText: oldText, NewText: newText}) - } - - if len(edits) == 0 { - return nil, false, fmt.Errorf("edits parameter is required") - } - for i, edit := range edits { - if edit.OldText == "" { - if len(edits) == 1 { - return nil, false, fmt.Errorf("old_text parameter is required") - } - return nil, false, fmt.Errorf("edits[%d].old_text must not be empty", i) - } - } - if replaceAll && len(edits) != 1 { - return nil, false, fmt.Errorf("replace_all only applies to a single edit") - } - return edits, replaceAll, nil -} - -func parseEditArray(raw any) ([]editReplacement, error) { - if s, ok := raw.(string); ok { - // Some models serialize array arguments as a JSON string. - if err := json.Unmarshal([]byte(s), &raw); err != nil { - return nil, fmt.Errorf("edits must be an array of {old_text, new_text} objects") - } - } - items, ok := raw.([]any) - if !ok { - return nil, fmt.Errorf("edits must be an array of {old_text, new_text} objects") - } - - edits := make([]editReplacement, 0, len(items)) - for i, item := range items { - entry, ok := item.(map[string]any) - if !ok { - return nil, fmt.Errorf("edits[%d] must be an object with old_text and new_text", i) - } - oldText, oldOK := editTextArg(entry, "old_text", "oldText") - newText, newOK := editTextArg(entry, "new_text", "newText") - if !oldOK || !newOK { - return nil, fmt.Errorf("edits[%d] must be an object with old_text and new_text", i) - } - edits = append(edits, editReplacement{OldText: oldText, NewText: newText}) - } - return edits, nil -} - -// editTextArg reads the first present string key, tolerating both snake_case -// and camelCase spellings that models emit. -func editTextArg(entry map[string]any, keys ...string) (string, bool) { - for _, key := range keys { - if value, ok := entry[key].(string); ok { - return value, true - } - } - return "", false -} - -func editNotFoundError(path string, editIndex, totalEdits int) error { - if totalEdits == 1 { - return fmt.Errorf("old_text was not found in %s", path) - } - return fmt.Errorf("edits[%d].old_text was not found in %s", editIndex, path) -} - -func editAmbiguousError(path string, editIndex, totalEdits, occurrences int) error { - if totalEdits == 1 { - return fmt.Errorf("old_text matched %d times in %s; set replace_all to true to replace every match", occurrences, path) - } - return fmt.Errorf("edits[%d].old_text matched %d times in %s; each edit must match exactly once, so provide more surrounding context", editIndex, occurrences, path) -} - -func cleanRelativePath(path string) (string, error) { - path = strings.TrimSpace(path) - if path == "" { - return "", fmt.Errorf("path parameter is required") - } - if filepath.IsAbs(path) { - return "", fmt.Errorf("absolute paths are not allowed") - } - cleaned := filepath.Clean(path) - if cleaned == "." || cleaned == ".." || strings.HasPrefix(cleaned, ".."+string(os.PathSeparator)) { - return "", fmt.Errorf("path escapes working directory") - } - return cleaned, nil -} - -func openRegularFile(workingDir, path string, allowAbsolute bool) (*os.File, os.FileInfo, error) { - path = strings.TrimSpace(path) - if path == "" { - return nil, nil, fmt.Errorf("path parameter is required") - } - if allowAbsolute && filepath.IsAbs(path) { - cleaned := filepath.Clean(path) - info, err := os.Lstat(cleaned) - if err != nil { - return nil, nil, err - } - if info.Mode()&os.ModeSymlink != 0 { - return nil, nil, fmt.Errorf("%s is a symlink; read the target file directly", path) - } - if err := rejectNonRegularFile(path, info); err != nil { - return nil, nil, err - } - file, err := os.Open(cleaned) - if err != nil { - return nil, nil, err - } - info, err = file.Stat() - if err != nil { - file.Close() - return nil, nil, err - } - if err := rejectNonRegularFile(path, info); err != nil { - file.Close() - return nil, nil, err - } - return file, info, nil - } - - rel, err := cleanRelativePath(path) - if err != nil { - return nil, nil, err - } - root, err := openWorkingRoot(workingDir) - if err != nil { - return nil, nil, err - } - defer root.Close() - - if _, err := regularRootFileInfo(root, rel, path); err != nil { - return nil, nil, err - } - file, err := root.Open(rel) - if err != nil { - return nil, nil, rootPathError(err) - } - info, err := file.Stat() - if err != nil { - file.Close() - return nil, nil, err - } - if err := rejectNonRegularFile(path, info); err != nil { - file.Close() - return nil, nil, err - } - return file, info, nil -} - -func regularRootFileInfo(root *os.Root, rel, path string) (os.FileInfo, error) { - info, err := root.Lstat(rel) - if err != nil { - return nil, rootPathError(err) - } - // Reject symlinks outright. os.Root.Open follows symlinks via openat - // without O_NOFOLLOW, so a symlink inside the working root that points - // outside it (e.g. ./notes -> ~/.ssh/id_rsa) would otherwise be read - // transparently, bypassing the working-directory confinement that the - // bash denylist enforces for direct credential reads. The caller must - // operate on the real target file instead. - if info.Mode()&os.ModeSymlink != 0 { - return nil, fmt.Errorf("%s is a symlink; read the target file directly", path) - } - if err := rejectNonRegularFile(path, info); err != nil { - return nil, err - } - return info, nil -} - -func rejectNonRegularFile(path string, info os.FileInfo) error { - if info.IsDir() { - return fmt.Errorf("%s is a directory", path) - } - if !info.Mode().IsRegular() { - return fmt.Errorf("%s is not a regular file", path) - } - return nil -} - -func writeFileAtomic(workingDir, path string, data []byte, perm os.FileMode) error { - rel, err := cleanRelativePath(path) - if err != nil { - return err - } - root, err := openWorkingRoot(workingDir) - if err != nil { - return err - } - defer root.Close() - if err := rejectRootFinalSymlink(root, rel, path); err != nil { - return err - } - - parent, name := filepath.Split(rel) - tmpBase := fmt.Sprintf(".%s.ollama-tmp-%d", name, os.Getpid()) - for i := 0; ; i++ { - candidateName := tmpBase - if i > 0 { - candidateName = fmt.Sprintf("%s-%d", tmpBase, i) - } - candidate := filepath.Join(parent, candidateName) - file, err := root.OpenFile(candidate, os.O_WRONLY|os.O_CREATE|os.O_EXCL, perm) - if os.IsExist(err) { - continue - } - if err != nil { - return rootPathError(err) - } - if err := file.Chmod(perm); err != nil { - closeErr := file.Close() - _ = root.Remove(candidate) - if closeErr != nil { - return closeErr - } - return err - } - writeErr := writeAllAndSync(file, data) - closeErr := file.Close() - if writeErr != nil || closeErr != nil { - _ = root.Remove(candidate) - if writeErr != nil { - return writeErr - } - return closeErr - } - if err := root.Rename(candidate, rel); err != nil { - _ = root.Remove(candidate) - return rootPathError(err) - } - return nil - } -} - -func rejectFinalSymlink(workingDir, path string) error { - rel, err := cleanRelativePath(path) - if err != nil { - return err - } - root, err := openWorkingRoot(workingDir) - if err != nil { - return err - } - defer root.Close() - return rejectRootFinalSymlink(root, rel, path) -} - -func rejectRootFinalSymlink(root *os.Root, rel, path string) error { - info, err := root.Lstat(rel) - if err != nil { - return rootPathError(err) - } - if info.Mode()&os.ModeSymlink != 0 { - return fmt.Errorf("%s is a symlink; edit the target file directly", path) - } - return nil -} - -func rootPathError(err error) error { - if err != nil && strings.Contains(err.Error(), "path escapes") { - return fmt.Errorf("path escapes working directory") - } - return err -} - -func openWorkingRoot(workingDir string) (*os.Root, error) { - base, err := workingDirAbs(workingDir) - if err != nil { - return nil, err - } - return os.OpenRoot(base) -} - -func writeAllAndSync(file *os.File, data []byte) error { - if _, err := file.Write(data); err != nil { - return err - } - return file.Sync() -} - -func readAllWithinLimit(reader io.Reader, limit int) ([]byte, error) { - if limit < 0 { - limit = 0 - } - content, err := io.ReadAll(io.LimitReader(reader, int64(limit)+1)) - if err != nil { - return nil, err - } - if len(content) > limit { - return nil, fmt.Errorf("content is too large (%d byte limit)", limit) - } - return content, nil -} - -func workingDirAbs(workingDir string) (string, error) { - base := workingDir - if base == "" { - var err error - base, err = os.Getwd() - if err != nil { - return "", err - } - } - return canonicalPath(base) -} - -func canonicalPath(path string) (string, error) { - abs, err := filepath.Abs(path) - if err != nil { - return "", err - } - resolved, err := filepath.EvalSymlinks(abs) - if err == nil { - return resolved, nil - } - return abs, nil -} - -type readSelection struct { - enabled bool - start int - end int -} - -func readSelectionFromArgs(args map[string]any) (readSelection, error) { - selection := readSelection{start: 1} - - if start, ok, err := intReadArg(args, "start"); err != nil { - return readSelection{}, err - } else if ok { - selection.enabled = true - selection.start = start - } - if end, ok, err := intReadArg(args, "end"); err != nil { - return readSelection{}, err - } else if ok { - selection.enabled = true - selection.end = end - } - - if !selection.enabled { - return selection, nil - } - if selection.start < 1 { - return readSelection{}, fmt.Errorf("start must be greater than 0") - } - if selection.end > 0 && selection.end < selection.start { - return readSelection{}, fmt.Errorf("end must be greater than or equal to start") - } - return selection, nil -} - -func readLineSelection(file *os.File, selection readSelection) (string, error) { - reader := bufio.NewReader(file) - var b strings.Builder - for lineNo := 1; ; { - line, err := reader.ReadSlice('\n') - if lineNo >= selection.start && (selection.end == 0 || lineNo <= selection.end) { - if b.Len()+len(line) > maxReadBytes { - return "", fmt.Errorf("selected content is too large (%d byte limit)", maxReadBytes) - } - b.Write(line) - } - if err != nil { - if err == bufio.ErrBufferFull { - continue - } - if err == io.EOF { - break - } - return "", err - } - if selection.end > 0 && lineNo >= selection.end { - break - } - lineNo++ - } - return b.String(), nil -} - -func intReadArg(args map[string]any, key string) (int, bool, error) { - value, ok := args[key] - if !ok { - return 0, false, nil - } - switch v := value.(type) { - case int: - return v, true, nil - case int64: - return int(v), true, nil - case float64: - if v != float64(int(v)) { - return 0, true, fmt.Errorf("%s must be a whole number", key) - } - return int(v), true, nil - case string: - v = strings.TrimSpace(v) - if v == "" { - return 0, false, nil - } - n, err := strconv.Atoi(v) - if err != nil { - return 0, true, fmt.Errorf("%s must be a whole number", key) - } - return n, true, nil - default: - return 0, true, fmt.Errorf("%s must be a whole number", key) - } -} - -func plural(n int) string { - if n == 1 { - return "" - } - return "s" -} diff --git a/agent/tools/file_test.go b/agent/tools/file_test.go deleted file mode 100644 index 15f28ab2792..00000000000 --- a/agent/tools/file_test.go +++ /dev/null @@ -1,571 +0,0 @@ -package tools - -import ( - "context" - "io" - "os" - "path/filepath" - "strings" - "testing" - - "github.com/ollama/ollama/agent" -) - -func TestEditReplacesUniqueText(t *testing.T) { - dir := t.TempDir() - path := filepath.Join(dir, "note.txt") - if err := os.WriteFile(path, []byte("hello world\n"), 0o644); err != nil { - t.Fatal(err) - } - - result, err := (&Edit{}).Execute(context.Background(), agent.ToolContext{WorkingDir: dir}, map[string]any{ - "path": "note.txt", - "old_text": "hello", - "new_text": "hi", - }) - if err != nil { - t.Fatal(err) - } - if !strings.Contains(result.Content, "Updated note.txt") { - t.Fatalf("result = %q", result.Content) - } - - content, err := os.ReadFile(path) - if err != nil { - t.Fatal(err) - } - if string(content) != "hi world\n" { - t.Fatalf("content = %q", content) - } -} - -func TestEditRequiresUniqueMatchByDefault(t *testing.T) { - dir := t.TempDir() - path := filepath.Join(dir, "note.txt") - if err := os.WriteFile(path, []byte("same same\n"), 0o644); err != nil { - t.Fatal(err) - } - - _, err := (&Edit{}).Execute(context.Background(), agent.ToolContext{WorkingDir: dir}, map[string]any{ - "path": "note.txt", - "old_text": "same", - "new_text": "other", - }) - if err == nil { - t.Fatal("expected ambiguous edit to fail") - } - if !strings.Contains(err.Error(), "matched 2 times") { - t.Fatalf("err = %v", err) - } -} - -func TestEditAppliesMultipleEdits(t *testing.T) { - dir := t.TempDir() - path := filepath.Join(dir, "note.txt") - if err := os.WriteFile(path, []byte("alpha beta gamma delta\n"), 0o644); err != nil { - t.Fatal(err) - } - - result, err := (&Edit{}).Execute(context.Background(), agent.ToolContext{WorkingDir: dir}, map[string]any{ - "path": "note.txt", - "edits": []any{ - map[string]any{"old_text": "beta", "new_text": "BETA"}, - map[string]any{"old_text": "delta", "new_text": "DELTA"}, - }, - }) - if err != nil { - t.Fatal(err) - } - if !strings.Contains(result.Content, "2 edits, 2 replacements") { - t.Fatalf("result = %q", result.Content) - } - - content, err := os.ReadFile(path) - if err != nil { - t.Fatal(err) - } - if string(content) != "alpha BETA gamma DELTA\n" { - t.Fatalf("content = %q", content) - } -} - -func TestEditMatchesEditsAgainstOriginalContent(t *testing.T) { - dir := t.TempDir() - path := filepath.Join(dir, "note.txt") - if err := os.WriteFile(path, []byte("abc def\n"), 0o644); err != nil { - t.Fatal(err) - } - - // edits[1] must target the original "def", not the one introduced by edits[0]. - _, err := (&Edit{}).Execute(context.Background(), agent.ToolContext{WorkingDir: dir}, map[string]any{ - "path": "note.txt", - "edits": []any{ - map[string]any{"old_text": "abc", "new_text": "def"}, - map[string]any{"old_text": "def", "new_text": "ghi"}, - }, - }) - if err != nil { - t.Fatal(err) - } - - content, err := os.ReadFile(path) - if err != nil { - t.Fatal(err) - } - if string(content) != "def ghi\n" { - t.Fatalf("content = %q", content) - } -} - -func TestEditRejectsOverlappingEdits(t *testing.T) { - dir := t.TempDir() - path := filepath.Join(dir, "note.txt") - if err := os.WriteFile(path, []byte("abc\n"), 0o644); err != nil { - t.Fatal(err) - } - - _, err := (&Edit{}).Execute(context.Background(), agent.ToolContext{WorkingDir: dir}, map[string]any{ - "path": "note.txt", - "edits": []any{ - map[string]any{"old_text": "ab", "new_text": "x"}, - map[string]any{"old_text": "bc", "new_text": "y"}, - }, - }) - if err == nil { - t.Fatal("expected overlapping edits to fail") - } - if !strings.Contains(err.Error(), "overlap") { - t.Fatalf("err = %v", err) - } -} - -func TestEditMultipleEditsNotFoundIndexed(t *testing.T) { - dir := t.TempDir() - path := filepath.Join(dir, "note.txt") - if err := os.WriteFile(path, []byte("hello world\n"), 0o644); err != nil { - t.Fatal(err) - } - - _, err := (&Edit{}).Execute(context.Background(), agent.ToolContext{WorkingDir: dir}, map[string]any{ - "path": "note.txt", - "edits": []any{ - map[string]any{"old_text": "hello", "new_text": "hi"}, - map[string]any{"old_text": "missing", "new_text": "x"}, - }, - }) - if err == nil { - t.Fatal("expected missing edit to fail") - } - if !strings.Contains(err.Error(), "edits[1]") { - t.Fatalf("err = %v", err) - } -} - -func TestEditMultipleEditsAmbiguousIndexed(t *testing.T) { - dir := t.TempDir() - path := filepath.Join(dir, "note.txt") - if err := os.WriteFile(path, []byte("hello same same\n"), 0o644); err != nil { - t.Fatal(err) - } - - _, err := (&Edit{}).Execute(context.Background(), agent.ToolContext{WorkingDir: dir}, map[string]any{ - "path": "note.txt", - "edits": []any{ - map[string]any{"old_text": "hello", "new_text": "hi"}, - map[string]any{"old_text": "same", "new_text": "x"}, - }, - }) - if err == nil { - t.Fatal("expected ambiguous edit to fail") - } - if !strings.Contains(err.Error(), "edits[1]") || !strings.Contains(err.Error(), "matched 2 times") { - t.Fatalf("err = %v", err) - } -} - -func TestEditRejectsEmptyEdits(t *testing.T) { - dir := t.TempDir() - path := filepath.Join(dir, "note.txt") - if err := os.WriteFile(path, []byte("hello\n"), 0o644); err != nil { - t.Fatal(err) - } - - for name, args := range map[string]map[string]any{ - "missing edits": {"path": "note.txt"}, - "empty edits": {"path": "note.txt", "edits": []any{}}, - } { - if _, err := (&Edit{}).Execute(context.Background(), agent.ToolContext{WorkingDir: dir}, args); err == nil { - t.Fatalf("%s: expected error", name) - } else if !strings.Contains(err.Error(), "edits parameter is required") { - t.Fatalf("%s: err = %v", name, err) - } - } -} - -func TestEditRejectsEmptyOldTextInArray(t *testing.T) { - dir := t.TempDir() - path := filepath.Join(dir, "note.txt") - if err := os.WriteFile(path, []byte("hello\n"), 0o644); err != nil { - t.Fatal(err) - } - - _, err := (&Edit{}).Execute(context.Background(), agent.ToolContext{WorkingDir: dir}, map[string]any{ - "path": "note.txt", - "edits": []any{ - map[string]any{"old_text": "hello", "new_text": "hi"}, - map[string]any{"old_text": "", "new_text": "x"}, - }, - }) - if err == nil { - t.Fatal("expected empty old_text to fail") - } - if !strings.Contains(err.Error(), "edits[1].old_text must not be empty") { - t.Fatalf("err = %v", err) - } -} - -func TestEditAcceptsJSONStringEdits(t *testing.T) { - dir := t.TempDir() - path := filepath.Join(dir, "note.txt") - if err := os.WriteFile(path, []byte("hello world\n"), 0o644); err != nil { - t.Fatal(err) - } - - // Some models serialize array arguments as a JSON string. - _, err := (&Edit{}).Execute(context.Background(), agent.ToolContext{WorkingDir: dir}, map[string]any{ - "path": "note.txt", - "edits": `[{"oldText": "hello", "newText": "hi"}, {"oldText": "world", "newText": "earth"}]`, - }) - if err != nil { - t.Fatal(err) - } - - content, err := os.ReadFile(path) - if err != nil { - t.Fatal(err) - } - if string(content) != "hi earth\n" { - t.Fatalf("content = %q", content) - } -} - -func TestEditRejectsReplaceAllWithMultipleEdits(t *testing.T) { - dir := t.TempDir() - path := filepath.Join(dir, "note.txt") - if err := os.WriteFile(path, []byte("a b c\n"), 0o644); err != nil { - t.Fatal(err) - } - - _, err := (&Edit{}).Execute(context.Background(), agent.ToolContext{WorkingDir: dir}, map[string]any{ - "path": "note.txt", - "replace_all": true, - "edits": []any{ - map[string]any{"old_text": "a", "new_text": "x"}, - map[string]any{"old_text": "b", "new_text": "y"}, - }, - }) - if err == nil { - t.Fatal("expected replace_all with multiple edits to fail") - } - if !strings.Contains(err.Error(), "replace_all") { - t.Fatalf("err = %v", err) - } -} - -func TestEditRejectsNoChange(t *testing.T) { - dir := t.TempDir() - path := filepath.Join(dir, "note.txt") - if err := os.WriteFile(path, []byte("hello\n"), 0o644); err != nil { - t.Fatal(err) - } - - _, err := (&Edit{}).Execute(context.Background(), agent.ToolContext{WorkingDir: dir}, map[string]any{ - "path": "note.txt", - "old_text": "hello", - "new_text": "hello", - }) - if err == nil { - t.Fatal("expected no-change edit to fail") - } - if !strings.Contains(err.Error(), "no changes") { - t.Fatalf("err = %v", err) - } -} - -func TestEditRejectsEscapingPath(t *testing.T) { - dir := t.TempDir() - _, err := (&Edit{}).Execute(context.Background(), agent.ToolContext{WorkingDir: dir}, map[string]any{ - "path": "../outside.txt", - "old_text": "old", - "new_text": "new", - }) - if err == nil { - t.Fatal("expected escaping path to fail") - } - if !strings.Contains(err.Error(), "path escapes working directory") { - t.Fatalf("err = %v", err) - } -} - -func TestEditRejectsSymlinkEscape(t *testing.T) { - dir := t.TempDir() - outside := t.TempDir() - if err := os.WriteFile(filepath.Join(outside, "note.txt"), []byte("old\n"), 0o644); err != nil { - t.Fatal(err) - } - if err := os.Symlink(outside, filepath.Join(dir, "link")); err != nil { - t.Skipf("symlinks unavailable: %v", err) - } - - _, err := (&Edit{}).Execute(context.Background(), agent.ToolContext{WorkingDir: dir}, map[string]any{ - "path": filepath.Join("link", "note.txt"), - "old_text": "old", - "new_text": "new", - }) - if err == nil { - t.Fatal("expected symlink escape to fail") - } - if !strings.Contains(err.Error(), "path escapes working directory") { - t.Fatalf("err = %v", err) - } - - content, err := os.ReadFile(filepath.Join(outside, "note.txt")) - if err != nil { - t.Fatal(err) - } - if string(content) != "old\n" { - t.Fatalf("outside content changed to %q", content) - } -} - -func TestEditRejectsFinalSymlink(t *testing.T) { - dir := t.TempDir() - target := filepath.Join(dir, "target.txt") - if err := os.WriteFile(target, []byte("old\n"), 0o644); err != nil { - t.Fatal(err) - } - link := filepath.Join(dir, "link.txt") - if err := os.Symlink("target.txt", link); err != nil { - t.Skipf("symlinks unavailable: %v", err) - } - - _, err := (&Edit{}).Execute(context.Background(), agent.ToolContext{WorkingDir: dir}, map[string]any{ - "path": "link.txt", - "old_text": "old", - "new_text": "new", - }) - if err == nil { - t.Fatal("expected final symlink edit to fail") - } - if !strings.Contains(err.Error(), "is a symlink") { - t.Fatalf("err = %v", err) - } - content, err := os.ReadFile(target) - if err != nil { - t.Fatal(err) - } - if string(content) != "old\n" { - t.Fatalf("target content changed to %q", content) - } - info, err := os.Lstat(link) - if err != nil { - t.Fatal(err) - } - if info.Mode()&os.ModeSymlink == 0 { - t.Fatalf("link mode = %v, want symlink", info.Mode()) - } -} - -func TestReadRejectsParentOutsideCurrentWorkingDir(t *testing.T) { - root := t.TempDir() - subdir := filepath.Join(root, "sub") - if err := os.Mkdir(subdir, 0o755); err != nil { - t.Fatal(err) - } - if err := os.WriteFile(filepath.Join(root, "note.txt"), []byte("hello"), 0o644); err != nil { - t.Fatal(err) - } - - _, err := (&Read{}).Execute(context.Background(), agent.ToolContext{WorkingDir: subdir}, map[string]any{ - "path": "../note.txt", - }) - if err == nil { - t.Fatal("expected parent path to fail") - } - if !strings.Contains(err.Error(), "path escapes working directory") { - t.Fatalf("err = %v", err) - } -} - -func TestReadRequiresApproval(t *testing.T) { - if !agent.ToolRequiresApproval((&Read{}), map[string]any{"path": "note.txt"}) { - t.Fatal("read should require approval") - } -} - -func TestReadDefaultsToEntireFile(t *testing.T) { - dir := t.TempDir() - content := "one\ntwo\nthree\n" - if err := os.WriteFile(filepath.Join(dir, "note.txt"), []byte(content), 0o644); err != nil { - t.Fatal(err) - } - - result, err := (&Read{}).Execute(context.Background(), agent.ToolContext{WorkingDir: dir}, map[string]any{ - "path": "note.txt", - }) - if err != nil { - t.Fatal(err) - } - if result.Content != content { - t.Fatalf("content = %q", result.Content) - } -} - -func TestReadAllowsAbsolutePath(t *testing.T) { - dir := t.TempDir() - path := filepath.Join(dir, "note.txt") - content := "one\ntwo\nthree\n" - if err := os.WriteFile(path, []byte(content), 0o644); err != nil { - t.Fatal(err) - } - - result, err := (&Read{}).Execute(context.Background(), agent.ToolContext{WorkingDir: t.TempDir()}, map[string]any{ - "path": path, - }) - if err != nil { - t.Fatal(err) - } - if result.Content != content { - t.Fatalf("content = %q", result.Content) - } -} - -func TestReadRejectsAbsoluteSymlink(t *testing.T) { - dir := t.TempDir() - target := filepath.Join(dir, "target.txt") - if err := os.WriteFile(target, []byte("hello\n"), 0o644); err != nil { - t.Fatal(err) - } - link := filepath.Join(dir, "alias") - if err := os.Symlink(target, link); err != nil { - t.Skipf("symlinks unavailable: %v", err) - } - - _, err := (&Read{}).Execute(context.Background(), agent.ToolContext{WorkingDir: t.TempDir()}, map[string]any{ - "path": link, - }) - if err == nil { - t.Fatal("expected absolute symlink to be rejected") - } - if !strings.Contains(err.Error(), "symlink") { - t.Fatalf("err = %v, want symlink rejection", err) - } -} - -func TestReadStartEnd(t *testing.T) { - dir := t.TempDir() - if err := os.WriteFile(filepath.Join(dir, "note.txt"), []byte("one\ntwo\nthree\nfour\n"), 0o644); err != nil { - t.Fatal(err) - } - - result, err := (&Read{}).Execute(context.Background(), agent.ToolContext{WorkingDir: dir}, map[string]any{ - "path": "note.txt", - "start": 2, - "end": 3, - }) - if err != nil { - t.Fatal(err) - } - if result.Content != "two\nthree\n" { - t.Fatalf("content = %q", result.Content) - } -} - -func TestReadStartOnly(t *testing.T) { - dir := t.TempDir() - if err := os.WriteFile(filepath.Join(dir, "note.txt"), []byte("one\ntwo\nthree\nfour\n"), 0o644); err != nil { - t.Fatal(err) - } - - result, err := (&Read{}).Execute(context.Background(), agent.ToolContext{WorkingDir: dir}, map[string]any{ - "path": "note.txt", - "start": 3, - }) - if err != nil { - t.Fatal(err) - } - if result.Content != "three\nfour\n" { - t.Fatalf("content = %q", result.Content) - } -} - -func TestReadEndOnly(t *testing.T) { - dir := t.TempDir() - if err := os.WriteFile(filepath.Join(dir, "note.txt"), []byte("one\ntwo\nthree\nfour\n"), 0o644); err != nil { - t.Fatal(err) - } - - result, err := (&Read{}).Execute(context.Background(), agent.ToolContext{WorkingDir: dir}, map[string]any{ - "path": "note.txt", - "end": 2, - }) - if err != nil { - t.Fatal(err) - } - if result.Content != "one\ntwo\n" { - t.Fatalf("content = %q", result.Content) - } -} - -func TestReadSelectionRejectsHugeSingleLine(t *testing.T) { - dir := t.TempDir() - if err := os.WriteFile(filepath.Join(dir, "note.txt"), []byte(strings.Repeat("x", maxReadBytes+1)), 0o644); err != nil { - t.Fatal(err) - } - - _, err := (&Read{}).Execute(context.Background(), agent.ToolContext{WorkingDir: dir}, map[string]any{ - "path": "note.txt", - "start": 1, - "end": 1, - }) - if err == nil { - t.Fatal("expected huge selected line to fail") - } - if !strings.Contains(err.Error(), "selected content is too large") { - t.Fatalf("err = %v", err) - } -} - -func TestReadAllWithinLimitRejectsGrowingRead(t *testing.T) { - reader := io.MultiReader( - strings.NewReader(strings.Repeat("x", maxReadBytes)), - strings.NewReader("x"), - ) - - _, err := readAllWithinLimit(reader, maxReadBytes) - if err == nil { - t.Fatal("expected over-limit read to fail") - } - if !strings.Contains(err.Error(), "content is too large") { - t.Fatalf("err = %v", err) - } -} - -func TestReadRejectsInvalidRange(t *testing.T) { - dir := t.TempDir() - if err := os.WriteFile(filepath.Join(dir, "note.txt"), []byte("one\ntwo\n"), 0o644); err != nil { - t.Fatal(err) - } - - _, err := (&Read{}).Execute(context.Background(), agent.ToolContext{WorkingDir: dir}, map[string]any{ - "path": "note.txt", - "start": 4, - "end": 2, - }) - if err == nil { - t.Fatal("expected invalid range to fail") - } - if !strings.Contains(err.Error(), "end must") { - t.Fatalf("err = %v", err) - } -} diff --git a/agent/tools/file_unix_test.go b/agent/tools/file_unix_test.go deleted file mode 100644 index a0973abc203..00000000000 --- a/agent/tools/file_unix_test.go +++ /dev/null @@ -1,121 +0,0 @@ -//go:build !windows - -package tools - -import ( - "context" - "os" - "path/filepath" - "strings" - "syscall" - "testing" - "time" - - "github.com/ollama/ollama/agent" -) - -func TestOpenRegularFileRejectsFIFO(t *testing.T) { - dir := t.TempDir() - path := filepath.Join(dir, "pipe") - if err := syscall.Mkfifo(path, 0o600); err != nil { - t.Skipf("mkfifo unavailable: %v", err) - } - - done := make(chan error, 1) - go func() { - file, _, err := openRegularFile(dir, "pipe", false) - if file != nil { - file.Close() - } - done <- err - }() - - select { - case err := <-done: - if err == nil { - t.Fatal("expected FIFO to be rejected") - } - if !strings.Contains(err.Error(), "not a regular file") { - t.Fatalf("err = %v", err) - } - case <-time.After(time.Second): - t.Fatal("openRegularFile blocked on FIFO") - } -} - -func TestEditPreservesModeDespiteUmask(t *testing.T) { - oldUmask := syscall.Umask(0o077) - defer syscall.Umask(oldUmask) - - dir := t.TempDir() - path := filepath.Join(dir, "note.txt") - if err := os.WriteFile(path, []byte("hello\n"), 0o666); err != nil { - t.Fatal(err) - } - if err := os.Chmod(path, 0o666); err != nil { - t.Fatal(err) - } - - _, err := (&Edit{}).Execute(context.Background(), agent.ToolContext{WorkingDir: dir}, map[string]any{ - "path": "note.txt", - "old_text": "hello", - "new_text": "hi", - }) - if err != nil { - t.Fatal(err) - } - - info, err := os.Stat(path) - if err != nil { - t.Fatal(err) - } - if got := info.Mode().Perm(); got != 0o666 { - t.Fatalf("mode = %#o, want 0666", got) - } -} - -func TestReadRejectsSymlinkEscapingWorkingDir(t *testing.T) { - root := t.TempDir() - secret := filepath.Join(t.TempDir(), "secret.txt") - if err := os.WriteFile(secret, []byte("top secret\n"), 0o600); err != nil { - t.Fatal(err) - } - link := filepath.Join(root, "notes") - if err := os.Symlink(secret, link); err != nil { - t.Fatal(err) - } - - _, err := (&Read{}).Execute(context.Background(), agent.ToolContext{WorkingDir: root}, map[string]any{ - "path": "notes", - }) - if err == nil { - t.Fatal("expected symlink escaping working dir to be rejected") - } - if !strings.Contains(err.Error(), "symlink") { - t.Fatalf("err = %v, want symlink rejection", err) - } -} - -func TestReadRejectsSymlinkInsideWorkingDirToOutside(t *testing.T) { - root := t.TempDir() - target := filepath.Join(root, "real.txt") - if err := os.WriteFile(target, []byte("hello\n"), 0o644); err != nil { - t.Fatal(err) - } - // A symlink to a sibling file still resolves inside the root; Read must - // reject it regardless, consistent with Edit's rejectFinalSymlink. - link := filepath.Join(root, "alias") - if err := os.Symlink(target, link); err != nil { - t.Fatal(err) - } - - _, err := (&Read{}).Execute(context.Background(), agent.ToolContext{WorkingDir: root}, map[string]any{ - "path": "alias", - }) - if err == nil { - t.Fatal("expected symlink to be rejected even when target is inside root") - } - if !strings.Contains(err.Error(), "symlink") { - t.Fatalf("err = %v, want symlink rejection", err) - } -} diff --git a/agent/tools/skill.go b/agent/tools/skill.go deleted file mode 100644 index 37620b6b239..00000000000 --- a/agent/tools/skill.go +++ /dev/null @@ -1,41 +0,0 @@ -package tools - -import ( - "context" - "errors" - - "github.com/ollama/ollama/agent" - "github.com/ollama/ollama/api" -) - -// Skill is the model-facing adapter for the core agent skill catalog. -// Model-initiated loads require approval because a skill's instructions can -// influence the rest of the run. Explicit user activation is handled by the -// session's synthetic skill call and bypasses this adapter. -type Skill struct{ Catalog *agent.SkillCatalog } - -func (t *Skill) Name() string { return "skill" } - -func (t *Skill) Description() string { - return "Load a named Ollama skill and return its instructions." -} - -func (t *Skill) Schema() api.ToolFunction { - props := api.NewToolPropertiesMap() - props.Set("name", api.ToolProperty{Type: api.PropertyType{"string"}, Description: "Name of the skill to load."}) - return api.ToolFunction{Name: t.Name(), Description: t.Description(), Parameters: api.ToolFunctionParameters{Type: "object", Properties: props, Required: []string{"name"}}} -} - -func (t *Skill) RequiresApproval(map[string]any) bool { return true } - -func (t *Skill) Execute(_ context.Context, _ agent.ToolContext, args map[string]any) (agent.ToolResult, error) { - name, ok := args["name"].(string) - if !ok { - return agent.ToolResult{}, errors.New("name parameter is required") - } - skill, err := t.Catalog.Load(name) - if err != nil { - return agent.ToolResult{}, err - } - return agent.ToolResult{Content: skill.Content()}, nil -} diff --git a/agent/tools/skill_test.go b/agent/tools/skill_test.go deleted file mode 100644 index 3835325abc0..00000000000 --- a/agent/tools/skill_test.go +++ /dev/null @@ -1,163 +0,0 @@ -package tools - -import ( - "context" - "os" - "path/filepath" - "strings" - "testing" - - "github.com/ollama/ollama/agent" - "github.com/ollama/ollama/api" -) - -func TestSkillLoadsCoreCatalogWithApproval(t *testing.T) { - catalog := testSkillCatalog(t) - tool := &Skill{Catalog: catalog} - if !agent.ToolRequiresApproval(tool, map[string]any{"name": "release-notes"}) { - t.Fatal("model-initiated skill loading should require approval") - } - result, err := tool.Execute(context.Background(), agent.ToolContext{}, map[string]any{"name": "release-notes"}) - if err != nil || !strings.Contains(result.Content, "Use concise bullets.") { - t.Fatalf("tool result = %#v, %v", result, err) - } -} - -func TestModelSkillLoadRequiresApproval(t *testing.T) { - for _, tt := range []struct { - name string - approval agent.Approval - prompt bool - wantCalls int - wantPrompts int - wantResult string - }{ - {name: "rejected", approval: agent.Approval{Reason: "Skill loading denied."}, prompt: true, wantCalls: 1, wantPrompts: 1, wantResult: "Skill loading denied."}, - {name: "approved", approval: agent.Approval{Allow: true}, prompt: true, wantCalls: 2, wantPrompts: 1, wantResult: "Use concise bullets."}, - {name: "headless denied", wantCalls: 1, wantResult: "Tool execution requires approval"}, - } { - t.Run(tt.name, func(t *testing.T) { - catalog := testSkillCatalog(t) - args := api.NewToolCallFunctionArguments() - args.Set("name", "release-notes") - client := &skillTestClient{responses: [][]api.ChatResponse{ - {{Message: api.Message{Role: "assistant", ToolCalls: []api.ToolCall{{ - ID: "call_skill_1", - Function: api.ToolCallFunction{Name: "skill", Arguments: args}, - }}}}}, - {{Message: api.Message{Role: "assistant", Content: "done"}}}, - }} - var prompter *skillApprovalPrompter - var approvalPrompter agent.ApprovalPrompter - if tt.prompt { - prompter = &skillApprovalPrompter{result: tt.approval} - approvalPrompter = prompter - } - registry := &agent.Registry{} - registry.Register(&Skill{Catalog: catalog}) - - result, err := (&agent.Session{ - Client: client, - Tools: registry, - ApprovalPrompter: approvalPrompter, - }).Run(context.Background(), agent.RunOptions{ - Model: "test", - NewMessages: []api.Message{{Role: "user", Content: "load the release-notes skill"}}, - }) - if err != nil { - t.Fatal(err) - } - if tt.prompt { - if got := len(prompter.requests); got != tt.wantPrompts { - t.Fatalf("approval prompts = %d, want %d", got, tt.wantPrompts) - } - request := prompter.requests[0] - if len(request.Calls) != 1 || request.Calls[0].ToolName != "skill" || request.Calls[0].ApprovalScope != "skill" || request.Calls[0].Args["name"] != "release-notes" { - t.Fatalf("approval request = %#v", request) - } - } - if got := client.calls; got != tt.wantCalls { - t.Fatalf("model calls = %d, want %d", got, tt.wantCalls) - } - var toolResult string - for _, message := range result.Messages { - if message.Role == "tool" && message.ToolCallID == "call_skill_1" { - toolResult = message.Content - break - } - } - if !strings.Contains(toolResult, tt.wantResult) { - t.Fatalf("skill tool result = %q, want it to contain %q", toolResult, tt.wantResult) - } - }) - } -} - -func TestExplicitSkillActivationBypassesApproval(t *testing.T) { - catalog := testSkillCatalog(t) - client := &skillTestClient{responses: [][]api.ChatResponse{{{Message: api.Message{Role: "assistant", Content: "done"}}}}} - prompter := &skillApprovalPrompter{result: agent.Approval{}} - result, err := (&agent.Session{ - Client: client, - Skills: catalog, - ApprovalPrompter: prompter, - }).Run(context.Background(), agent.RunOptions{ - Model: "test", - NewMessages: []api.Message{{Role: "user", Content: "draft release notes"}}, - SkillName: "release-notes", - }) - if err != nil { - t.Fatal(err) - } - if len(prompter.requests) != 0 { - t.Fatalf("explicit activation prompted for approval: %#v", prompter.requests) - } - if len(result.Messages) != 4 || result.Messages[2].ToolName != "skill" || !strings.Contains(result.Messages[2].Content, "Use concise bullets.") { - t.Fatalf("synthetic skill activation = %#v", result.Messages) - } -} - -func testSkillCatalog(t *testing.T) *agent.SkillCatalog { - t.Helper() - dir := t.TempDir() - path := filepath.Join(dir, "release-notes") - if err := os.Mkdir(path, 0o755); err != nil { - t.Fatal(err) - } - if err := os.WriteFile(filepath.Join(path, "SKILL.md"), []byte("---\nname: release-notes\ndescription: Draft release notes.\n---\nUse concise bullets."), 0o644); err != nil { - t.Fatal(err) - } - catalog, err := agent.DiscoverSkills(dir) - if err != nil { - t.Fatal(err) - } - return catalog -} - -type skillTestClient struct { - responses [][]api.ChatResponse - calls int -} - -func (c *skillTestClient) Chat(_ context.Context, _ *api.ChatRequest, fn api.ChatResponseFunc) error { - if c.calls >= len(c.responses) { - return nil - } - for _, response := range c.responses[c.calls] { - if err := fn(response); err != nil { - return err - } - } - c.calls++ - return nil -} - -type skillApprovalPrompter struct { - requests []agent.ApprovalRequest - result agent.Approval -} - -func (p *skillApprovalPrompter) PromptApproval(_ context.Context, request agent.ApprovalRequest) (agent.Approval, error) { - p.requests = append(p.requests, request) - return p.result, nil -} diff --git a/agent/tools/web.go b/agent/tools/web.go deleted file mode 100644 index fe6849ad61b..00000000000 --- a/agent/tools/web.go +++ /dev/null @@ -1,186 +0,0 @@ -package tools - -import ( - "context" - "errors" - "fmt" - "net/url" - "strings" - "time" - - "github.com/ollama/ollama/agent" - "github.com/ollama/ollama/api" - internalcloud "github.com/ollama/ollama/internal/cloud" -) - -const ( - maxWebFetchContentRunes = 60_000 - webSearchTimeout = 15 * time.Second - webFetchTimeout = 30 * time.Second -) - -var ErrWebAuthRequired = errors.New("Not authenticated. Run `ollama signin` and try again.") - -type WebSearch struct{} - -func (w *WebSearch) Name() string { - return "web_search" -} - -func (w *WebSearch) Description() string { - return "Search the web for current information that may not be in the model's training data." -} - -func (w *WebSearch) Schema() api.ToolFunction { - props := api.NewToolPropertiesMap() - props.Set("query", api.ToolProperty{ - Type: api.PropertyType{"string"}, - Description: "The search query to look up on the web.", - }) - return api.ToolFunction{ - Name: w.Name(), - Description: w.Description(), - Parameters: api.ToolFunctionParameters{ - Type: "object", - Properties: props, - Required: []string{"query"}, - }, - } -} - -func (w *WebSearch) RequiresApproval(map[string]any) bool { - return true -} - -func (w *WebSearch) Execute(ctx context.Context, _ agent.ToolContext, args map[string]any) (agent.ToolResult, error) { - // TODO: use shared agent.RequiredStringArg for the "query" parameter (see agent package cleanup plan). - if internalcloud.Disabled() { - return agent.ToolResult{}, errors.New(internalcloud.DisabledError("web search is unavailable")) - } - query, ok := args["query"].(string) - if !ok || strings.TrimSpace(query) == "" { - return agent.ToolResult{}, fmt.Errorf("query parameter is required") - } - - client, err := api.ClientFromEnvironment() - if err != nil { - return agent.ToolResult{}, err - } - - ctx, cancel := context.WithTimeout(ctx, webSearchTimeout) - defer cancel() - - searchResp, err := client.WebSearchExperimental(ctx, &api.WebSearchRequest{Query: query, MaxResults: 5}) - if err != nil { - var authErr api.AuthorizationError - if errors.As(err, &authErr) { - return agent.ToolResult{}, ErrWebAuthRequired - } - return agent.ToolResult{}, err - } - if len(searchResp.Results) == 0 { - return agent.ToolResult{Content: "No results found for query: " + query}, nil - } - - var sb strings.Builder - sb.WriteString(fmt.Sprintf("Search results for: %s\n\n", query)) - for i, result := range searchResp.Results { - sb.WriteString(fmt.Sprintf("%d. %s\n", i+1, result.Title)) - sb.WriteString(fmt.Sprintf(" URL: %s\n", result.URL)) - if result.Content != "" { - content := []rune(result.Content) - if len(content) > 300 { - content = append(content[:300], []rune("...")...) - } - sb.WriteString(fmt.Sprintf(" %s\n", string(content))) - } - sb.WriteByte('\n') - } - return agent.ToolResult{Content: sb.String()}, nil -} - -type WebFetch struct{} - -func (w *WebFetch) Name() string { - return "web_fetch" -} - -func (w *WebFetch) Description() string { - return "Fetch and extract text content from a web page." -} - -func (w *WebFetch) Schema() api.ToolFunction { - props := api.NewToolPropertiesMap() - props.Set("url", api.ToolProperty{ - Type: api.PropertyType{"string"}, - Description: "The URL to fetch and extract content from.", - }) - return api.ToolFunction{ - Name: w.Name(), - Description: w.Description(), - Parameters: api.ToolFunctionParameters{ - Type: "object", - Properties: props, - Required: []string{"url"}, - }, - } -} - -func (w *WebFetch) RequiresApproval(map[string]any) bool { - return true -} - -func (w *WebFetch) Execute(ctx context.Context, _ agent.ToolContext, args map[string]any) (agent.ToolResult, error) { - // TODO: use shared agent.RequiredStringArg for the "url" parameter (see agent package cleanup plan). - if internalcloud.Disabled() { - return agent.ToolResult{}, errors.New(internalcloud.DisabledError("web fetch is unavailable")) - } - urlStr, ok := args["url"].(string) - if !ok || strings.TrimSpace(urlStr) == "" { - return agent.ToolResult{}, fmt.Errorf("url parameter is required") - } - parsed, err := url.Parse(urlStr) - if err != nil { - return agent.ToolResult{}, fmt.Errorf("invalid URL: %w", err) - } - if scheme := strings.ToLower(parsed.Scheme); scheme != "http" && scheme != "https" { - return agent.ToolResult{}, fmt.Errorf("unsupported URL scheme %q: only http and https are allowed", parsed.Scheme) - } - - client, err := api.ClientFromEnvironment() - if err != nil { - return agent.ToolResult{}, err - } - - ctx, cancel := context.WithTimeout(ctx, webFetchTimeout) - defer cancel() - - fetchResp, err := client.WebFetchExperimental(ctx, &api.WebFetchRequest{URL: urlStr}) - if err != nil { - var authErr api.AuthorizationError - if errors.As(err, &authErr) { - return agent.ToolResult{}, ErrWebAuthRequired - } - return agent.ToolResult{}, err - } - - var sb strings.Builder - if fetchResp.Title != "" { - sb.WriteString(fmt.Sprintf("Title: %s\n\n", fetchResp.Title)) - } - if fetchResp.Content != "" { - sb.WriteString("Content:\n") - sb.WriteString(truncateWebFetchContent(fetchResp.Content)) - } else { - sb.WriteString("No content could be extracted from the page.") - } - return agent.ToolResult{Content: sb.String()}, nil -} - -func truncateWebFetchContent(content string) string { - return agent.Truncate(content, agent.TruncateConfig{ - MaxRunes: maxWebFetchContentRunes, - Label: "tool output", - Hint: "Use a narrower request or search query if more detail is needed.", - }) -} diff --git a/agent/tools/web_test.go b/agent/tools/web_test.go deleted file mode 100644 index a9327820fe3..00000000000 --- a/agent/tools/web_test.go +++ /dev/null @@ -1,214 +0,0 @@ -package tools - -import ( - "encoding/json" - "errors" - "net/http" - "net/http/httptest" - "os" - "path/filepath" - "strings" - "testing" - - coreagent "github.com/ollama/ollama/agent" - "github.com/ollama/ollama/api" - "github.com/ollama/ollama/envconfig" - internalcloud "github.com/ollama/ollama/internal/cloud" -) - -func TestWebToolsRequireApproval(t *testing.T) { - if !coreagent.ToolRequiresApproval((&WebSearch{}), map[string]any{"query": "ollama"}) { - t.Fatal("web search should require approval") - } - if !coreagent.ToolRequiresApproval((&WebFetch{}), map[string]any{"url": "https://ollama.com"}) { - t.Fatal("web fetch should require approval") - } -} - -var webToolCases = []struct { - name string - tool coreagent.Tool - args map[string]any - path string - operation string -}{ - {"search", &WebSearch{}, map[string]any{"query": "ollama"}, "/api/experimental/web_search", "web search is unavailable"}, - {"fetch", &WebFetch{}, map[string]any{"url": "https://ollama.com"}, "/api/experimental/web_fetch", "web fetch is unavailable"}, -} - -// enableWebToolsForTest isolates web tool tests from the runner's cloud -// policy. In particular, Windows can inherit both OLLAMA_NO_CLOUD and a -// server.json from USERPROFILE. -func enableWebToolsForTest(t *testing.T) { - t.Helper() - - // Register before t.Setenv so the cache is refreshed after t.Setenv has - // restored the runner's environment during cleanup. - t.Cleanup(envconfig.ReloadServerConfig) - - home := t.TempDir() - t.Setenv("HOME", home) - t.Setenv("USERPROFILE", home) - t.Setenv("OLLAMA_NO_CLOUD", "") - envconfig.ReloadServerConfig() -} - -// runWebTool executes tool against a stub server that responds to every -// request with status and body, returning the resulting error. -func runWebTool(t *testing.T, tool coreagent.Tool, args map[string]any, path string, status int, body string) error { - t.Helper() - ts := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { - if r.URL.Path != path { - t.Fatalf("path = %q, want %q", r.URL.Path, path) - } - w.Header().Set("Content-Type", "application/json") - w.WriteHeader(status) - _, _ = w.Write([]byte(body)) - })) - t.Cleanup(ts.Close) - t.Setenv("OLLAMA_HOST", ts.URL) - _, err := tool.Execute(t.Context(), coreagent.ToolContext{}, args) - return err -} - -func TestWebToolsReportAuthenticationError(t *testing.T) { - enableWebToolsForTest(t) - - for _, tt := range webToolCases { - t.Run(tt.name, func(t *testing.T) { - err := runWebTool(t, tt.tool, tt.args, tt.path, http.StatusUnauthorized, - `{"error":"unauthorized","signin_url":"https://ollama.com/signin"}`) - if !errors.Is(err, ErrWebAuthRequired) { - t.Fatalf("error = %v, want %v", err, ErrWebAuthRequired) - } - }) - } -} - -func TestWebToolsPreserveNonAuthenticationErrors(t *testing.T) { - enableWebToolsForTest(t) - - for _, tt := range webToolCases { - t.Run(tt.name, func(t *testing.T) { - err := runWebTool(t, tt.tool, tt.args, tt.path, http.StatusTooManyRequests, - `{"error":"web search quota exceeded"}`) - if err == nil { - t.Fatal("expected error") - } - if !strings.Contains(err.Error(), "web search quota exceeded") { - t.Fatalf("error = %q, want original error message", err) - } - }) - } -} - -func TestWebToolsIgnoreInheritedCloudPolicy(t *testing.T) { - // This cleanup is registered before the test environment, so it restores - // the server config cache after t.Setenv restores the runner's values. - t.Cleanup(envconfig.ReloadServerConfig) - - home := t.TempDir() - configPath := filepath.Join(home, ".ollama", "server.json") - if err := os.MkdirAll(filepath.Dir(configPath), 0o755); err != nil { - t.Fatal(err) - } - if err := os.WriteFile(configPath, []byte(`{"disable_ollama_cloud":true}`), 0o644); err != nil { - t.Fatal(err) - } - t.Setenv("HOME", home) - t.Setenv("USERPROFILE", home) - t.Setenv("OLLAMA_NO_CLOUD", "1") - envconfig.ReloadServerConfig() - - enableWebToolsForTest(t) - err := runWebTool(t, &WebSearch{}, map[string]any{"query": "ollama"}, "/api/experimental/web_search", http.StatusUnauthorized, - `{"error":"unauthorized","signin_url":"https://ollama.com/signin"}`) - if !errors.Is(err, ErrWebAuthRequired) { - t.Fatalf("error = %v, want %v", err, ErrWebAuthRequired) - } -} - -func TestWebFetchRejectsUnsupportedScheme(t *testing.T) { - enableWebToolsForTest(t) - - tests := []struct { - name string - url string - wantErr bool - }{ - {name: "file scheme", url: "file:///etc/passwd", wantErr: true}, - {name: "data scheme", url: "data:text/plain,secret", wantErr: true}, - {name: "ftp scheme", url: "ftp://example.com/secret", wantErr: true}, - {name: "http allowed", url: "http://example.com", wantErr: false}, - {name: "https allowed", url: "https://example.com", wantErr: false}, - } - for _, tt := range tests { - t.Run(tt.name, func(t *testing.T) { - _, err := (&WebFetch{}).Execute(t.Context(), coreagent.ToolContext{}, map[string]any{"url": tt.url}) - if tt.wantErr && err == nil { - t.Fatal("expected unsupported scheme to be rejected") - } - // For allowed schemes we expect an error only from the missing - // server/auth path, not from scheme validation. The http/https - // cases reach the client and may fail on connection/auth; we only - // assert that the error is NOT a scheme error. - if !tt.wantErr && err != nil && strings.Contains(err.Error(), "unsupported URL scheme") { - t.Fatalf("http/https rejected as unsupported: %v", err) - } - }) - } -} - -func TestWebFetchBoundsContentBeforeReturning(t *testing.T) { - enableWebToolsForTest(t) - - ts := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { - if r.URL.Path != "/api/experimental/web_fetch" { - t.Fatalf("path = %q, want /api/experimental/web_fetch", r.URL.Path) - } - var req api.WebFetchRequest - if err := json.NewDecoder(r.Body).Decode(&req); err != nil { - t.Fatal(err) - } - if req.URL != "https://ollama.com" { - t.Fatalf("request URL = %q, want https://ollama.com", req.URL) - } - if err := json.NewEncoder(w).Encode(api.WebFetchResponse{ - Title: "Ollama", - Content: strings.Repeat("x", maxWebFetchContentRunes+25), - }); err != nil { - t.Fatal(err) - } - })) - defer ts.Close() - t.Setenv("OLLAMA_HOST", ts.URL) - - result, err := (&WebFetch{}).Execute(t.Context(), coreagent.ToolContext{}, map[string]any{ - "url": "https://ollama.com", - }) - if err != nil { - t.Fatal(err) - } - if !strings.Contains(result.Content, "[tool output truncated: showing first ~") || - !strings.Contains(result.Content, "omitted ~7 tokens") || - !strings.Contains(result.Content, "Use a narrower request or search query") { - t.Fatalf("content missing truncation marker: %q", result.Content) - } - if count := strings.Count(result.Content, "x"); count != maxWebFetchContentRunes { - t.Fatalf("captured content count = %d, want %d", count, maxWebFetchContentRunes) - } -} - -func TestWebToolsRejectWhenCloudDisabled(t *testing.T) { - t.Setenv("OLLAMA_NO_CLOUD", "1") - - for _, tt := range webToolCases { - t.Run(tt.name, func(t *testing.T) { - _, err := tt.tool.Execute(t.Context(), coreagent.ToolContext{}, tt.args) - want := internalcloud.DisabledError(tt.operation) - if err == nil || err.Error() != want { - t.Fatalf("error = %v, want %q", err, want) - } - }) - } -} diff --git a/api/client.go b/api/client.go index 46e948d0920..10913308aae 100644 --- a/api/client.go +++ b/api/client.go @@ -450,6 +450,19 @@ func (c *Client) CreateBlob(ctx context.Context, digest string, r io.Reader) err return c.do(ctx, http.MethodPost, fmt.Sprintf("/api/blobs/%s", digest), r, nil) } +// HeadBlob checks if a blob exists on the server. It returns false for a 404 +// and an error for any unexpected response. +func (c *Client) HeadBlob(ctx context.Context, digest string) (bool, error) { + if err := c.do(ctx, http.MethodHead, fmt.Sprintf("/api/blobs/%s", digest), nil, nil); err != nil { + var statusErr StatusError + if errors.As(err, &statusErr) && statusErr.StatusCode == http.StatusNotFound { + return false, nil + } + return false, err + } + return true, nil +} + // Version returns the Ollama server version as a string. func (c *Client) Version(ctx context.Context) (string, error) { var version struct { diff --git a/api/client_test.go b/api/client_test.go index 5fd89a5afc1..b73983ffde6 100644 --- a/api/client_test.go +++ b/api/client_test.go @@ -352,6 +352,75 @@ func TestClientDo(t *testing.T) { } } +func TestHeadBlob(t *testing.T) { + testCases := []struct { + name string + statusCode int + body string + wantExists bool + wantErr string + }{ + { + name: "exists", + statusCode: http.StatusOK, + wantExists: true, + }, + { + name: "missing", + statusCode: http.StatusNotFound, + body: `{"error":"missing"}`, + }, + { + name: "server error", + statusCode: http.StatusInternalServerError, + body: `{"error":"stat failed"}`, + wantErr: "stat failed", + }, + } + + for _, tc := range testCases { + t.Run(tc.name, func(t *testing.T) { + client := NewClient( + &url.URL{Scheme: "http", Host: "example.com"}, + &http.Client{Transport: roundTripFunc(func(req *http.Request) (*http.Response, error) { + if req.Method != http.MethodHead { + t.Fatalf("method = %q, want HEAD", req.Method) + } + if req.URL.Path != "/api/blobs/sha256:abc" { + t.Fatalf("path = %q, want /api/blobs/sha256:abc", req.URL.Path) + } + + status := http.StatusText(tc.statusCode) + return &http.Response{ + StatusCode: tc.statusCode, + Status: fmt.Sprintf("%d %s", tc.statusCode, status), + Body: io.NopCloser(strings.NewReader(tc.body)), + Header: make(http.Header), + Request: req, + }, nil + })}, + ) + + exists, err := client.HeadBlob(t.Context(), "sha256:abc") + if tc.wantErr != "" { + if err == nil { + t.Fatalf("got nil, want error %q", tc.wantErr) + } + if err.Error() != tc.wantErr { + t.Fatalf("error = %q, want %q", err.Error(), tc.wantErr) + } + return + } + if err != nil { + t.Fatalf("HeadBlob() error = %v", err) + } + if exists != tc.wantExists { + t.Fatalf("HeadBlob() exists = %v, want %v", exists, tc.wantExists) + } + }) + } +} + func TestClientWebSearchExperimentalUsesLocalRoute(t *testing.T) { var gotPath string var gotMethod string diff --git a/api/types.go b/api/types.go index 7ea9022be92..c83f3f2228e 100644 --- a/api/types.go +++ b/api/types.go @@ -576,7 +576,7 @@ type Options struct { TopK int `json:"top_k,omitempty"` TopP float32 `json:"top_p,omitempty"` MinP float32 `json:"min_p,omitempty"` - TypicalP float32 `json:"typical_p,omitempty"` + TypicalP float32 `json:"typical_p,omitempty"` // Deprecated: rejected on new requests and models; still honored from existing model parameters RepeatLastN int `json:"repeat_last_n,omitempty"` Temperature float32 `json:"temperature,omitempty"` RepeatPenalty float32 `json:"repeat_penalty,omitempty"` @@ -657,10 +657,10 @@ type CreateRequest struct { // Stream specifies whether the response is streaming; it is true by default. Stream *bool `json:"stream,omitempty"` - // Quantize is the quantization format for the model; leave blank to not change the quantization level. + // Quantize is the quantization format to apply when importing safetensors weights. Quantize string `json:"quantize,omitempty"` - // DraftQuantize is the quantization format for the draft model. + // DraftQuantize is the quantization format to apply when importing safetensors draft weights. DraftQuantize string `json:"draft_quantize,omitempty"` // From is the name of the model or file to use as the source. @@ -669,13 +669,15 @@ type CreateRequest struct { // RemoteHost is the URL of the upstream ollama API for the model (if any). RemoteHost string `json:"remote_host,omitempty"` - // Files is a map of files include when creating the model. + // Files maps source file names to their SHA-256 digests. Files map[string]string `json:"files,omitempty"` - // DraftFiles is a map of draft model files to include when creating the model. + // DraftFiles maps draft source file names to their SHA-256 digests. DraftFiles map[string]string `json:"draft_files,omitempty"` // Adapters is a map of LoRA adapters to include when creating the model. + // + // Deprecated: LoRA adapters are no longer supported. Adapters map[string]string `json:"adapters,omitempty"` // Template is the template used when constructing a request to the model. diff --git a/app/cmd/app/webview.go b/app/cmd/app/webview.go index 0427601a74e..89da905075d 100644 --- a/app/cmd/app/webview.go +++ b/app/cmd/app/webview.go @@ -238,6 +238,13 @@ func (w *Webview) Run(path string) unsafe.Pointer { return } + if runtime.GOOS == "darwin" { + // Keep the current frame through the handoff. SetSize also + // recenters the macOS window and would jump before Apps paints. + setOnboardingWindowStyle(wv.Window(), false) + return + } + width, height := defaultWindowWidth, defaultWindowHeight if w.Store != nil { storedWidth, storedHeight, err := w.Store.WindowSize() diff --git a/app/ui/app/src/components/CodexDesktopModelsSettings.tsx b/app/ui/app/src/components/CodexDesktopModelsSettings.tsx index 5a4da77121d..678d0e67b7e 100644 --- a/app/ui/app/src/components/CodexDesktopModelsSettings.tsx +++ b/app/ui/app/src/components/CodexDesktopModelsSettings.tsx @@ -566,7 +566,7 @@ export const CodexDesktopModelsSettings = forwardRef<
Choose ChatGPT models -
+
{selected.map((model) => ( { - it("renders a disconnected ChatGPT toggle", () => { - const html = renderToStaticMarkup( - , - ); - - expect(html).toContain(">ChatGPT (Desktop)

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

- ChatGPT (Desktop) +

+ ChatGPT

{description}

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

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

- {item.name} -

-

- {item.description} -

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