From 9c04e04e6b26655aee03dff90639ad03ef4e0367 Mon Sep 17 00:00:00 2001 From: ParthSareen Date: Thu, 18 Jun 2026 16:24:43 -0700 Subject: [PATCH] agent: address review feedback --- agent/approval.go | 54 +++++++- agent/approval_test.go | 54 ++++++++ agent/chatstore/sqlite.go | 7 +- agent/compactor.go | 6 +- agent/compactor_test.go | 26 ++++ agent/events.go | 9 ++ agent/session.go | 65 +++++++-- agent/session_test.go | 70 ++++++++++ agent/skills/import.go | 46 +++++-- agent/skills/skills.go | 91 ++++++++++--- agent/skills/skills_test.go | 80 +++++++++++ agent/tools/bash.go | 5 +- agent/tools/bash_test.go | 18 +++ agent/tools/bash_unix.go | 4 + agent/tools/bash_windows.go | 64 ++++++++- agent/tools/file.go | 177 +++++++++++++++++-------- agent/tools/file_test.go | 49 +++++++ cmd/cmd.go | 4 + cmd/cmd_launcher_test.go | 20 +++ cmd/internal/filedata/filedata.go | 14 +- cmd/internal/filedata/filedata_test.go | 16 +++ cmd/tui/chat.go | 32 +++-- cmd/tui/chat_events.go | 8 ++ cmd/tui/chat_input.go | 27 +++- cmd/tui/chat_input_test.go | 21 +++ cmd/tui/chat_multimodal_test.go | 26 ++++ cmd/tui/chat_render.go | 18 ++- cmd/tui/chat_render_test.go | 43 ++++++ cmd/tui/tui.go | 2 +- progress/progress.go | 20 ++- progress/progress_test.go | 42 ++++++ 31 files changed, 985 insertions(+), 133 deletions(-) create mode 100644 progress/progress_test.go diff --git a/agent/approval.go b/agent/approval.go index a10d85ebe..42894c82e 100644 --- a/agent/approval.go +++ b/agent/approval.go @@ -438,6 +438,8 @@ func (c *bashClassifier) addCall(args []string) { c.addReason("runs destructive git clean") c.high = true } + case "find": + c.addFindReasons(args) case "chmod", "chown": if hasAnyFlag(args[1:], "R", "recursive") { c.addReason("changes permissions or ownership recursively") @@ -459,6 +461,27 @@ func (c *bashClassifier) addCall(args []string) { } } +func (c *bashClassifier) addFindReasons(args []string) { + for i := 1; i < len(args); i++ { + switch args[i] { + case "-delete": + c.addReason("deletes files via find") + c.high = true + case "-exec", "-execdir": + c.addReason("executes commands via find") + c.high = true + end := i + 1 + for end < len(args) && args[end] != ";" && args[end] != `\;` && args[end] != "+" { + end++ + } + if end > i+1 { + c.addCall(args[i+1 : end]) + } + i = end + } + } +} + func hasAnyFlag(args []string, flags ...string) bool { for _, arg := range args { if arg == "--" { @@ -481,14 +504,39 @@ func hasAnyFlag(args []string, flags ...string) bool { } func isGitResetHard(args []string) bool { - return len(args) >= 3 && args[1] == "reset" && slices.Contains(args[2:], "--hard") + subcommand := gitSubcommandArgs(args) + return len(subcommand) >= 2 && subcommand[0] == "reset" && slices.Contains(subcommand[1:], "--hard") } func isGitCleanDestructive(args []string) bool { - if len(args) < 2 || args[1] != "clean" { + subcommand := gitSubcommandArgs(args) + if len(subcommand) < 1 || subcommand[0] != "clean" { return false } - return hasAnyFlag(args[2:], "f", "force") && (hasAnyFlag(args[2:], "d") || hasAnyFlag(args[2:], "x", "X")) + return hasAnyFlag(subcommand[1:], "f", "force") && (hasAnyFlag(subcommand[1:], "d") || hasAnyFlag(subcommand[1:], "x", "X")) +} + +func gitSubcommandArgs(args []string) []string { + for i := 1; i < len(args); i++ { + arg := args[i] + if arg == "--" { + return args[i+1:] + } + switch { + case arg == "-C", arg == "-c", arg == "--git-dir", arg == "--work-tree": + i++ + continue + case strings.HasPrefix(arg, "-C"), strings.HasPrefix(arg, "-c"): + continue + case strings.HasPrefix(arg, "--git-dir="), strings.HasPrefix(arg, "--work-tree="): + continue + case strings.HasPrefix(arg, "-"): + continue + default: + return args[i:] + } + } + return nil } type bashTokenKind int diff --git a/agent/approval_test.go b/agent/approval_test.go index 12c337ae3..817872728 100644 --- a/agent/approval_test.go +++ b/agent/approval_test.go @@ -185,6 +185,60 @@ func TestBashApprovalClassifiesDynamicShellEvasions(t *testing.T) { } } +func TestBashApprovalClassifiesDestructiveGitWithGlobalOptions(t *testing.T) { + tests := []struct { + name string + cmd string + reason string + }{ + {name: "git reset hard after cwd", cmd: "git -C /tmp reset --hard", reason: "runs destructive git reset"}, + {name: "git reset hard after config", cmd: "git -c core.autocrlf=false reset --hard", reason: "runs destructive git reset"}, + {name: "git clean after work tree", cmd: "git --work-tree=/tmp clean -fdx", reason: "runs destructive git clean"}, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + evaluation := DefaultApprovalPolicy{}.EvaluateApproval(context.Background(), ApprovalRequest{ + ToolName: "bash", + Args: map[string]any{"command": tt.cmd}, + }) + if evaluation.Risk != ApprovalRiskHigh { + t.Fatalf("risk = %q, want high", evaluation.Risk) + } + if reasons := strings.Join(evaluation.Reasons, " "); !strings.Contains(reasons, tt.reason) { + t.Fatalf("reasons = %#v, want %q", evaluation.Reasons, tt.reason) + } + }) + } +} + +func TestBashApprovalClassifiesFindMutations(t *testing.T) { + tests := []struct { + name string + cmd string + reason string + }{ + {name: "delete", cmd: "find . -name '*.tmp' -delete", reason: "deletes files via find"}, + {name: "exec", cmd: `find . -name '*.tmp' -exec rm -rf {} \;`, reason: "executes commands via find"}, + {name: "exec nested destructive command", cmd: `find . -name '*.tmp' -exec rm -rf {} \;`, reason: "removes files destructively"}, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + evaluation := DefaultApprovalPolicy{}.EvaluateApproval(context.Background(), ApprovalRequest{ + ToolName: "bash", + Args: map[string]any{"command": tt.cmd}, + }) + if evaluation.Risk != ApprovalRiskHigh { + t.Fatalf("risk = %q, want high", evaluation.Risk) + } + if reasons := strings.Join(evaluation.Reasons, " "); !strings.Contains(reasons, tt.reason) { + t.Fatalf("reasons = %#v, want %q", evaluation.Reasons, tt.reason) + } + }) + } +} + func TestWebApprovalAllowsWithoutPrompt(t *testing.T) { tests := []struct { name string diff --git a/agent/chatstore/sqlite.go b/agent/chatstore/sqlite.go index 0172dd522..800d8153d 100644 --- a/agent/chatstore/sqlite.go +++ b/agent/chatstore/sqlite.go @@ -107,7 +107,6 @@ func (s *Store) init(ctx context.Context) error { content TEXT NOT NULL DEFAULT '', thinking TEXT NOT NULL DEFAULT '', images TEXT NOT NULL DEFAULT '[]', - stream BOOLEAN NOT NULL DEFAULT 0, model_name TEXT, created_at TIMESTAMP NOT NULL DEFAULT CURRENT_TIMESTAMP, updated_at TIMESTAMP NOT NULL DEFAULT CURRENT_TIMESTAMP, @@ -779,9 +778,9 @@ func insertMessage(ctx context.Context, tx *sql.Tx, chatID string, msg api.Messa return 0, err } result, err := tx.ExecContext(ctx, ` - INSERT INTO messages (chat_id, role, content, thinking, images, stream, tool_name, tool_call_id, model_name, created_at, updated_at) - VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?) - `, chatID, msg.Role, msg.Content, msg.Thinking, imagesJSON, false, msg.ToolName, msg.ToolCallID, modelName, now, now) + INSERT INTO messages (chat_id, role, content, thinking, images, tool_name, tool_call_id, model_name, created_at, updated_at) + VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?) + `, chatID, msg.Role, msg.Content, msg.Thinking, imagesJSON, msg.ToolName, msg.ToolCallID, modelName, now, now) if err != nil { return 0, fmt.Errorf("insert message: %w", err) } diff --git a/agent/compactor.go b/agent/compactor.go index d1bdb4a03..12966f20f 100644 --- a/agent/compactor.go +++ b/agent/compactor.go @@ -424,7 +424,11 @@ func fitCompactionMessagesToBudget(previousSummary string, messages []api.Messag if nextRunes >= currentRunes { nextRunes = currentRunes / 2 } - fitted[idx].Content = truncateToolResultContentTo(fitted[idx].Content, nextRunes) + nextContent := truncateToolResultContentTo(fitted[idx].Content, nextRunes) + if nextContent == fitted[idx].Content && len([]rune(nextContent)) > max(0, nextRunes) { + nextContent = forceTruncateToolResultContentTo(fitted[idx].Content, nextRunes) + } + fitted[idx].Content = nextContent } return fitted } diff --git a/agent/compactor_test.go b/agent/compactor_test.go index 592176d16..5ccd15a74 100644 --- a/agent/compactor_test.go +++ b/agent/compactor_test.go @@ -439,6 +439,32 @@ func TestCompactionPromptFitsBudgetByTruncatingLargeToolOutput(t *testing.T) { } } +func TestCompactionPromptRetruncatesAlreadyTruncatedToolOutput(t *testing.T) { + alreadyTruncated := strings.Repeat("x", 7000) + "\n\n[tool output truncated: omitted 99999 characters]\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: omitted ") { + t.Fatalf("truncation marker missing from compaction prompt: %q", body) + } +} + func TestCompactionSummaryTextStripsPrefix(t *testing.T) { content := compactionSummaryMessage("worked on branch changes") if got := compactionSummaryText(content); got != "worked on branch changes" { diff --git a/agent/events.go b/agent/events.go index 49623b709..ff0d278c6 100644 --- a/agent/events.go +++ b/agent/events.go @@ -1,6 +1,7 @@ package agent import ( + "context" "time" "github.com/ollama/ollama/api" @@ -64,3 +65,11 @@ func emit(sink EventSink, event Event) error { } return sink.Emit(event) } + +func emitIgnoringCanceled(ctx context.Context, sink EventSink, event Event) error { + err := emit(sink, event) + if err != nil && ctx != nil && ctx.Err() != nil { + return nil + } + return err +} diff --git a/agent/session.go b/agent/session.go index e7021742e..bc606d7a2 100644 --- a/agent/session.go +++ b/agent/session.go @@ -169,7 +169,7 @@ func (s *Session) Run(ctx context.Context, opts RunOptions) (*RunResult, error) if canceled { if len(pendingToolCalls) > 0 { - skipped, skipErr := s.skipToolCalls(context.WithoutCancel(ctx), runID, opts, pendingToolCalls, "Tool execution skipped because the run was canceled.") + skipped, skipErr := s.skipToolCalls(ctx, runID, opts, pendingToolCalls, "Tool execution skipped because the run was canceled.") if skipErr != nil { emit(s.Events, Event{Type: EventError, RunID: runID, ChatID: opts.ChatID, Error: skipErr.Error()}) return nil, skipErr @@ -354,14 +354,23 @@ func (s *Session) executeToolCalls(ctx context.Context, runID string, opts RunOp toolMessages := make([]api.Message, 0, len(calls)) projectedMessages := append([]api.Message(nil), messages...) + persistCtx := context.WithoutCancel(ctx) for i, call := range calls { toolName := call.Function.Name args := call.Function.Arguments.ToMap() + 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 nil, false, skipErr + } + toolMessages = append(toolMessages, skipped...) + return toolMessages, true, nil + } tool, ok := s.Tools.Get(toolName) if !ok { content := fmt.Sprintf("Error: unknown tool: %s", toolName) msg := s.toolMessageForContext(toolName, call.ID, content, opts, projectedMessages) - if err := s.appendToolMessage(ctx, opts.ChatID, msg); err != nil { + if err := s.appendToolMessage(persistCtx, opts.ChatID, msg); err != nil { return nil, false, err } toolMessages = append(toolMessages, msg) @@ -382,6 +391,14 @@ func (s *Session) executeToolCalls(ctx context.Context, runID string, opts RunOp if toolNeedsApproval(ctx, approval, tool, approvalRequest) { result, err := approval.Approve(ctx, approvalRequest) if err != nil { + 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 nil, false, skipErr + } + toolMessages = append(toolMessages, skipped...) + return toolMessages, true, nil + } return nil, false, err } if result.Decision == ApprovalDeny { @@ -390,7 +407,7 @@ func (s *Session) executeToolCalls(ctx context.Context, runID string, opts RunOp content = "Tool execution denied." } msg := s.toolMessageForContext(toolName, call.ID, content, opts, projectedMessages) - if err := s.appendToolMessage(ctx, opts.ChatID, msg); err != nil { + if err := s.appendToolMessage(persistCtx, opts.ChatID, msg); err != nil { return nil, false, err } toolMessages = append(toolMessages, msg) @@ -404,7 +421,7 @@ func (s *Session) executeToolCalls(ctx context.Context, runID string, opts RunOp skippedArgs := skipped.Function.Arguments.ToMap() skippedContent := "Tool execution skipped because a previous tool call in this assistant message was denied." skippedMsg := s.toolMessageForContext(skippedToolName, skipped.ID, skippedContent, opts, projectedMessages) - if err := s.appendToolMessage(ctx, opts.ChatID, skippedMsg); err != nil { + if err := s.appendToolMessage(persistCtx, opts.ChatID, skippedMsg); err != nil { return nil, false, err } toolMessages = append(toolMessages, skippedMsg) @@ -427,15 +444,23 @@ func (s *Session) executeToolCalls(ctx context.Context, runID string, opts RunOp if err != nil { content := fmt.Sprintf("Error: %v", err) msg := s.toolMessageForContext(toolName, call.ID, content, opts, projectedMessages) - if appendErr := s.appendToolMessage(ctx, opts.ChatID, msg); appendErr != nil { + if appendErr := s.appendToolMessage(persistCtx, opts.ChatID, msg); appendErr != nil { return nil, false, appendErr } toolMessages = append(toolMessages, msg) projectedMessages = append(projectedMessages, msg) content = msg.Content - if emitErr := emit(s.Events, Event{Type: EventToolFinished, RunID: runID, ChatID: opts.ChatID, ToolCallID: call.ID, ToolName: toolName, Args: args, Content: content, Error: err.Error(), FinishedAt: time.Now()}); emitErr != nil { + if emitErr := emitIgnoringCanceled(ctx, s.Events, Event{Type: EventToolFinished, RunID: runID, ChatID: opts.ChatID, ToolCallID: call.ID, ToolName: toolName, Args: args, Content: content, Error: err.Error(), FinishedAt: time.Now()}); emitErr != nil { return nil, false, 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 nil, false, skipErr + } + toolMessages = append(toolMessages, skipped...) + return toolMessages, true, nil + } continue } @@ -443,31 +468,43 @@ func (s *Session) executeToolCalls(ctx context.Context, runID string, opts RunOp content := result.Content msg := s.toolMessageForContext(toolName, call.ID, content, opts, projectedMessages) - if err := s.appendToolMessage(ctx, opts.ChatID, msg); err != nil { + if err := s.appendToolMessage(persistCtx, opts.ChatID, msg); err != nil { return nil, false, err } toolMessages = append(toolMessages, msg) projectedMessages = append(projectedMessages, msg) content = msg.Content - if err := emit(s.Events, Event{Type: EventToolFinished, RunID: runID, ChatID: opts.ChatID, ToolCallID: call.ID, ToolName: toolName, WorkingDir: s.WorkingDir, Args: args, Content: content, FinishedAt: time.Now()}); err != nil { + if err := emitIgnoringCanceled(ctx, s.Events, Event{Type: EventToolFinished, RunID: runID, ChatID: opts.ChatID, ToolCallID: call.ID, ToolName: toolName, WorkingDir: s.WorkingDir, Args: args, Content: content, FinishedAt: time.Now()}); err != nil { return nil, false, 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 nil, false, skipErr + } + toolMessages = append(toolMessages, skipped...) + return toolMessages, true, nil + } } return toolMessages, false, nil } func (s *Session) skipToolCalls(ctx context.Context, runID string, opts RunOptions, calls []api.ToolCall, content string) ([]api.Message, error) { toolMessages := make([]api.Message, 0, len(calls)) + appendCtx := ctx + if ctx != nil && ctx.Err() != nil { + appendCtx = context.WithoutCancel(ctx) + } for _, call := range calls { toolName := call.Function.Name args := call.Function.Arguments.ToMap() msg := toolMessage(toolName, call.ID, content) - if err := s.appendToolMessage(ctx, opts.ChatID, msg); err != nil { + if err := s.appendToolMessage(appendCtx, opts.ChatID, msg); err != nil { return nil, err } toolMessages = append(toolMessages, msg) - if emitErr := emit(s.Events, Event{Type: EventToolFinished, RunID: runID, ChatID: opts.ChatID, ToolCallID: call.ID, ToolName: toolName, Args: args, Content: msg.Content, Error: msg.Content, FinishedAt: time.Now()}); emitErr != nil { + if emitErr := emitIgnoringCanceled(ctx, s.Events, Event{Type: EventToolFinished, RunID: runID, ChatID: opts.ChatID, ToolCallID: call.ID, ToolName: toolName, Args: args, Content: msg.Content, Error: msg.Content, FinishedAt: time.Now()}); emitErr != nil { return nil, emitErr } } @@ -744,6 +781,14 @@ func truncateToolResultContentTo(content string, maxRunes int) string { if strings.Contains(content, "[tool output truncated: omitted ") { return content } + return truncateToolResultContentToLimit(content, maxRunes) +} + +func forceTruncateToolResultContentTo(content string, maxRunes int) string { + return truncateToolResultContentToLimit(content, maxRunes) +} + +func truncateToolResultContentToLimit(content string, maxRunes int) string { if maxRunes <= 0 { return fmt.Sprintf("[tool output truncated: omitted %d characters]", len([]rune(content))) } diff --git a/agent/session_test.go b/agent/session_test.go index 0170cd5a5..ec357fda1 100644 --- a/agent/session_test.go +++ b/agent/session_test.go @@ -99,6 +99,10 @@ type cwdTestTool struct{} type largeTool struct{} +type cancelingTool struct { + cancel context.CancelFunc +} + type cancelAfterToolCallClient struct { cancel context.CancelFunc } @@ -229,6 +233,30 @@ func (largeTool) Execute(context.Context, ToolContext, map[string]any) (ToolResu return ToolResult{Content: strings.Repeat("x", maxToolResultRunes+100)}, 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" } @@ -637,6 +665,48 @@ func TestSessionCancellationAfterToolCallAppendsSkippedToolMessage(t *testing.T) } } +func TestSessionCancellationDuringToolExecutionAppendsToolMessage(t *testing.T) { + ctx, cancel := context.WithCancel(context.Background()) + store := &contextAwareStore{} + registry := NewRegistry() + 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, + Store: store, + Tools: registry, + } + + result, err := session.Run(ctx, RunOptions{ + ChatID: "chat-1", + Model: "model", + NewMessages: []api.Message{{Role: "user", Content: "cancel during tool"}}, + UseTools: true, + }) + if err != nil { + t.Fatal(err) + } + if len(result.Messages) != 3 { + t.Fatalf("messages = %#v", result.Messages) + } + if len(store.messages) != len(result.Messages) { + t.Fatalf("stored messages = %#v, want %d messages", store.messages, len(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) + } +} + func TestSessionToolLoopAllowsRoundsUnderDefaultCap(t *testing.T) { responses := make([][]api.ChatResponse, 0, 26) for i := range 25 { diff --git a/agent/skills/import.go b/agent/skills/import.go index 986226cb9..da282a95e 100644 --- a/agent/skills/import.go +++ b/agent/skills/import.go @@ -18,6 +18,12 @@ type ImportResult struct { Error string } +type skillDirCandidate struct { + Dir string + Skipped bool + Error string +} + func Import(source string, force bool) ([]ImportResult, error) { dest, err := DefaultDir() if err != nil { @@ -45,9 +51,16 @@ func ImportToDir(source, dest string, force bool) ([]ImportResult, error) { results = append(results, ImportResult{Source: source, From: root, Skipped: true, Error: err.Error()}) continue } - for _, dir := range candidates { - skill, err := ReadMetadata(filepath.Join(dir, SkillFile)) - result := ImportResult{Source: source, From: dir} + for _, candidate := range candidates { + result := ImportResult{Source: source, From: candidate.Dir} + if candidate.Skipped { + result.Skipped = true + result.Error = candidate.Error + results = append(results, result) + continue + } + + skill, err := ReadMetadata(filepath.Join(candidate.Dir, SkillFile)) if err != nil { result.Skipped = true result.Error = err.Error() @@ -57,13 +70,15 @@ func ImportToDir(source, dest string, force bool) ([]ImportResult, error) { result.Skill = skill result.To = filepath.Join(dest, skill.Name) - err = copyDir(dir, result.To, force) + copyResult, err := copyDir(candidate.Dir, result.To, force) if errors.Is(err, os.ErrExist) { result.Skipped = true result.Error = "already exists" } else if err != nil { result.Skipped = true result.Error = err.Error() + } else if len(copyResult.Skipped) > 0 { + result.Error = "skipped symlinks: " + strings.Join(copyResult.Skipped, ", ") } results = append(results, result) } @@ -104,22 +119,35 @@ func SourceDirs(source string) ([]string, error) { return nil, fmt.Errorf("unknown skill source %q (use claude, codex, pi, agents, or all)", source) } -func skillDirs(root string) ([]string, error) { +func skillDirs(root string) ([]skillDirCandidate, error) { entries, err := os.ReadDir(root) if err != nil { return nil, err } - var dirs []string + var dirs []skillDirCandidate for _, entry := range entries { - if !entry.IsDir() || strings.HasPrefix(entry.Name(), ".") { + if strings.HasPrefix(entry.Name(), ".") { continue } dir := filepath.Join(root, entry.Name()) + if entry.Type()&os.ModeSymlink != 0 { + dirs = append(dirs, skillDirCandidate{ + Dir: dir, + Skipped: true, + Error: "symlinked skill directories are not supported", + }) + continue + } + if !entry.IsDir() { + continue + } if _, err := os.Stat(filepath.Join(dir, SkillFile)); err == nil { - dirs = append(dirs, dir) + dirs = append(dirs, skillDirCandidate{Dir: dir}) } } - slices.Sort(dirs) + slices.SortFunc(dirs, func(a, b skillDirCandidate) int { + return strings.Compare(a.Dir, b.Dir) + }) return dirs, nil } diff --git a/agent/skills/skills.go b/agent/skills/skills.go index 166863bad..2e17c174f 100644 --- a/agent/skills/skills.go +++ b/agent/skills/skills.go @@ -14,7 +14,8 @@ import ( ) const ( - SkillFile = "SKILL.md" + SkillFile = "SKILL.md" + maxSkillFileBytes = 1 << 20 ) var validName = regexp.MustCompile(`^[a-z0-9][a-z0-9-]{0,63}$`) @@ -92,7 +93,7 @@ func Load(dir string) (*Catalog, error) { } func ReadMetadata(path string) (Skill, error) { - data, err := os.ReadFile(path) + data, err := readSkillFile(path) if err != nil { return Skill{}, err } @@ -170,13 +171,30 @@ func (s Skill) Read() (string, error) { if s.File == "" { return "", fmt.Errorf("skill %q has no %s path", s.Name, SkillFile) } - data, err := os.ReadFile(s.File) + data, err := readSkillFile(s.File) if err != nil { return "", err } return string(data), nil } +func readSkillFile(path string) ([]byte, error) { + info, err := os.Lstat(path) + if err != nil { + return nil, err + } + if info.Mode()&os.ModeSymlink != 0 { + return nil, fmt.Errorf("%s must not be a symlink", SkillFile) + } + if !info.Mode().IsRegular() { + return nil, fmt.Errorf("%s must be a regular file", SkillFile) + } + if info.Size() > maxSkillFileBytes { + return nil, fmt.Errorf("%s exceeds %d bytes", SkillFile, maxSkillFileBytes) + } + return os.ReadFile(path) +} + func NormalizeName(name string) string { name = strings.TrimSpace(name) name = strings.TrimPrefix(name, "/") @@ -218,19 +236,34 @@ func validateMetadata(meta frontmatter) error { return nil } -func copyDir(src, dst string, force bool) error { - if force { - if err := os.RemoveAll(dst); err != nil { - return err - } - } - if _, err := os.Stat(dst); err == nil { - return fs.ErrExist +type copyDirResult struct { + Skipped []string +} + +func copyDir(src, dst string, force bool) (copyDirResult, error) { + if _, err := os.Stat(dst); err == nil && !force { + return copyDirResult{}, fs.ErrExist } else if err != nil && !errors.Is(err, os.ErrNotExist) { - return err + return copyDirResult{}, err } - return filepath.WalkDir(src, func(path string, d fs.DirEntry, err error) error { + parent := filepath.Dir(dst) + if err := os.MkdirAll(parent, 0o755); err != nil { + return copyDirResult{}, err + } + tmp, err := os.MkdirTemp(parent, "."+filepath.Base(dst)+".tmp-*") + if err != nil { + return copyDirResult{}, err + } + moved := false + defer func() { + if !moved { + _ = os.RemoveAll(tmp) + } + }() + + var result copyDirResult + if err := filepath.WalkDir(src, func(path string, d fs.DirEntry, err error) error { if err != nil { return err } @@ -238,13 +271,21 @@ func copyDir(src, dst string, force bool) error { if err != nil { return err } - target := filepath.Join(dst, rel) + if rel == "." { + return nil + } + target := filepath.Join(tmp, rel) if d.Type()&os.ModeSymlink != 0 { + result.Skipped = append(result.Skipped, rel) return nil } if d.IsDir() { - return os.MkdirAll(target, 0o755) + info, err := d.Info() + if err != nil { + return err + } + return os.MkdirAll(target, info.Mode().Perm()) } info, err := d.Info() @@ -256,5 +297,23 @@ func copyDir(src, dst string, force bool) error { return err } return os.WriteFile(target, data, info.Mode().Perm()) - }) + }); err != nil { + return result, err + } + + if force { + if err := os.RemoveAll(dst); err != nil { + return result, err + } + } + if _, err := os.Stat(dst); err == nil { + return result, fs.ErrExist + } else if err != nil && !errors.Is(err, os.ErrNotExist) { + return result, err + } + if err := os.Rename(tmp, dst); err != nil { + return result, err + } + moved = true + return result, nil } diff --git a/agent/skills/skills_test.go b/agent/skills/skills_test.go index fbde52aba..7e2604ca3 100644 --- a/agent/skills/skills_test.go +++ b/agent/skills/skills_test.go @@ -3,6 +3,7 @@ package skills import ( "os" "path/filepath" + "runtime" "strings" "testing" ) @@ -65,6 +66,85 @@ func TestImportToDirCopiesCanonicalSkill(t *testing.T) { } } +func TestReadMetadataRejectsSymlink(t *testing.T) { + if runtime.GOOS == "windows" { + t.Skip("symlink permissions vary on Windows") + } + dir := t.TempDir() + real := filepath.Join(dir, "real.md") + if err := os.WriteFile(real, []byte("---\nname: go-code\ndescription: Write idiomatic Go code.\n---\n"), 0o644); err != nil { + t.Fatal(err) + } + link := filepath.Join(dir, SkillFile) + if err := os.Symlink(real, link); err != nil { + t.Fatal(err) + } + if _, err := ReadMetadata(link); err == nil || !strings.Contains(err.Error(), "must not be a symlink") { + t.Fatalf("ReadMetadata error = %v, want symlink rejection", err) + } +} + +func TestImportToDirReportsSymlinkedSkillDirectory(t *testing.T) { + if runtime.GOOS == "windows" { + t.Skip("symlink permissions vary on Windows") + } + home := t.TempDir() + t.Setenv("HOME", home) + t.Setenv("USERPROFILE", home) + srcRoot := filepath.Join(home, ".claude", "skills") + real := filepath.Join(home, "elsewhere", "go-code") + writeSkill(t, real, "go-code", "Write idiomatic Go code.") + if err := os.MkdirAll(srcRoot, 0o755); err != nil { + t.Fatal(err) + } + if err := os.Symlink(real, filepath.Join(srcRoot, "go-code")); err != nil { + t.Fatal(err) + } + + results, err := ImportToDir("claude", filepath.Join(home, ".ollama", "skills"), false) + if err != nil { + t.Fatal(err) + } + if len(results) != 1 || !results[0].Skipped { + t.Fatalf("results = %#v, want one skipped symlink directory", results) + } + if !strings.Contains(results[0].Error, "symlinked skill directories") { + t.Fatalf("error = %q, want symlink directory warning", results[0].Error) + } +} + +func TestImportToDirReportsSkippedSymlinkEntries(t *testing.T) { + if runtime.GOOS == "windows" { + t.Skip("symlink permissions vary on Windows") + } + home := t.TempDir() + t.Setenv("HOME", home) + t.Setenv("USERPROFILE", home) + src := filepath.Join(home, ".claude", "skills", "go-code") + writeSkill(t, src, "go-code", "Write idiomatic Go code.") + target := filepath.Join(home, "outside.md") + if err := os.WriteFile(target, []byte("outside"), 0o644); err != nil { + t.Fatal(err) + } + if err := os.Symlink(target, filepath.Join(src, "outside.md")); err != nil { + t.Fatal(err) + } + + results, err := ImportToDir("claude", filepath.Join(home, ".ollama", "skills"), false) + if err != nil { + t.Fatal(err) + } + if len(results) != 1 || results[0].Skipped { + t.Fatalf("results = %#v, want one imported skill", results) + } + if !strings.Contains(results[0].Error, "skipped symlinks: outside.md") { + t.Fatalf("error = %q, want skipped symlink warning", results[0].Error) + } + if _, err := os.Lstat(filepath.Join(home, ".ollama", "skills", "go-code", "outside.md")); !os.IsNotExist(err) { + t.Fatalf("copied symlink err = %v, want missing symlink", err) + } +} + func writeSkill(t *testing.T, dir, name, description string) { t.Helper() if err := os.MkdirAll(dir, 0o755); err != nil { diff --git a/agent/tools/bash.go b/agent/tools/bash.go index 7e80acc20..768c47059 100644 --- a/agent/tools/bash.go +++ b/agent/tools/bash.go @@ -87,7 +87,7 @@ func (b *Bash) Execute(ctx context.Context, toolCtx agent.ToolContext, args map[ cmd.Stdout = &stdout cmd.Stderr = &stderr - err = cmd.Run() + err = runBashCommand(cmd) finalWorkingDir := readFinalWorkingDir(cwdPath) var sb strings.Builder @@ -106,6 +106,9 @@ func (b *Bash) Execute(ctx context.Context, toolCtx agent.ToolContext, args map[ if ctx.Err() == context.DeadlineExceeded { return agent.ToolResult{Content: sb.String() + "\n\nError: command timed out after " + bashTimeout.String(), WorkingDir: finalWorkingDir}, nil } + if ctx.Err() == context.Canceled { + return agent.ToolResult{Content: sb.String() + "\n\nError: command was canceled", WorkingDir: finalWorkingDir}, nil + } if exitErr, ok := err.(*exec.ExitError); ok { return agent.ToolResult{Content: sb.String() + fmt.Sprintf("\n\nExit code: %d", exitErr.ExitCode()), WorkingDir: finalWorkingDir}, nil } diff --git a/agent/tools/bash_test.go b/agent/tools/bash_test.go index 97b86c6c0..aa9710235 100644 --- a/agent/tools/bash_test.go +++ b/agent/tools/bash_test.go @@ -54,6 +54,24 @@ func TestBashBoundsOutputWhileRunning(t *testing.T) { } } +func TestBashReportsCanceledCommand(t *testing.T) { + ctx, cancel := context.WithCancel(context.Background()) + cancel() + + result, err := NewBash().Execute(ctx, agent.ToolContext{WorkingDir: t.TempDir()}, map[string]any{ + "command": "sleep 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 TestReadFinalWorkingDirRejectsInvalidPaths(t *testing.T) { dir := t.TempDir() cwdFile := filepath.Join(dir, "cwd") diff --git a/agent/tools/bash_unix.go b/agent/tools/bash_unix.go index c0264930f..08dab5b8e 100644 --- a/agent/tools/bash_unix.go +++ b/agent/tools/bash_unix.go @@ -11,6 +11,10 @@ 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 diff --git a/agent/tools/bash_windows.go b/agent/tools/bash_windows.go index 8da84a9b4..af6ce2c35 100644 --- a/agent/tools/bash_windows.go +++ b/agent/tools/bash_windows.go @@ -2,14 +2,76 @@ package tools -import "os/exec" +import ( + "os/exec" + "sync" + "unsafe" + + "golang.org/x/sys/windows" +) + +var bashJobHandles sync.Map func configureBashCommand(*exec.Cmd) {} +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/file.go b/agent/tools/file.go index 89eebeee8..12fba3374 100644 --- a/agent/tools/file.go +++ b/agent/tools/file.go @@ -71,18 +71,11 @@ func (r *Read) Execute(ctx context.Context, toolCtx agent.ToolContext, args map[ return agent.ToolResult{}, fmt.Errorf("path parameter is required") } - resolved, err := resolvePath(toolCtx.WorkingDir, path) + file, info, err := openRegularFile(toolCtx.WorkingDir, path) if err != nil { return agent.ToolResult{}, err } - - info, err := os.Stat(resolved) - if err != nil { - return agent.ToolResult{}, err - } - if info.IsDir() { - return agent.ToolResult{}, fmt.Errorf("%s is a directory", path) - } + defer file.Close() selection, err := readSelectionFromArgs(args) if err != nil { @@ -100,10 +93,10 @@ func (r *Read) Execute(ctx context.Context, toolCtx agent.ToolContext, args map[ var content string if selection.enabled { - content, err = readLineSelection(resolved, selection) + content, err = readLineSelection(file, selection) } else { var contentBytes []byte - contentBytes, err = os.ReadFile(resolved) + contentBytes, err = io.ReadAll(file) content = string(contentBytes) } if err != nil { @@ -177,29 +170,26 @@ func (e *Edit) Execute(ctx context.Context, toolCtx agent.ToolContext, args map[ replaceAll, _ := args["replace_all"].(bool) - resolved, err := resolvePath(toolCtx.WorkingDir, path) + file, info, err := openRegularFile(toolCtx.WorkingDir, path) if err != nil { return agent.ToolResult{}, err } - - info, err := os.Stat(resolved) - if err != nil { - return agent.ToolResult{}, err - } - if info.IsDir() { - return agent.ToolResult{}, fmt.Errorf("%s is a directory", path) - } 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 := os.ReadFile(resolved) + contentBytes, err := io.ReadAll(file) + if closeErr := file.Close(); err == nil && closeErr != nil { + err = closeErr + } if err != nil { return agent.ToolResult{}, err } @@ -222,14 +212,121 @@ func (e *Edit) Execute(ctx context.Context, toolCtx agent.ToolContext, args map[ return agent.ToolResult{}, fmt.Errorf("edited content is too large (%d bytes)", len(updated)) } - if err := os.WriteFile(resolved, []byte(updated), info.Mode().Perm()); err != nil { + if err := writeFileAtomic(toolCtx.WorkingDir, path, []byte(updated), info.Mode().Perm()); err != nil { return agent.ToolResult{}, err } return agent.ToolResult{Content: fmt.Sprintf("Updated %s (%d replacement%s).", path, matches, plural(matches))}, nil } -func resolvePath(workingDir, path string) (string, error) { +func cleanRelativePath(path string) (string, error) { + path = strings.TrimSpace(path) + if path == "" { + return "", fmt.Errorf("path parameter is required") + } + if filepath.IsAbs(path) { + return "", fmt.Errorf("absolute paths are not allowed") + } + cleaned := filepath.Clean(path) + if cleaned == "." || cleaned == ".." || strings.HasPrefix(cleaned, ".."+string(os.PathSeparator)) { + return "", fmt.Errorf("path escapes working directory") + } + return cleaned, nil +} + +func openRegularFile(workingDir, path string) (*os.File, os.FileInfo, error) { + rel, err := cleanRelativePath(path) + if err != nil { + return nil, nil, err + } + root, err := openWorkingRoot(workingDir) + if err != nil { + return nil, nil, err + } + defer root.Close() + + 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 info.IsDir() { + file.Close() + return nil, nil, fmt.Errorf("%s is a directory", path) + } + return file, info, 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() + + 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) + } + 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 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 workingDirAbs(workingDir string) (string, error) { base := workingDir if base == "" { var err error @@ -238,29 +335,7 @@ func resolvePath(workingDir, path string) (string, error) { return "", err } } - - if filepath.IsAbs(path) { - return "", fmt.Errorf("absolute paths are not allowed") - } - - baseAbs, err := canonicalPath(base) - if err != nil { - return "", err - } - resolved := filepath.Clean(filepath.Join(baseAbs, path)) - resolvedForCheck := resolved - if canonical, err := canonicalPath(resolved); err == nil { - resolvedForCheck = canonical - } - rel, err := filepath.Rel(baseAbs, resolvedForCheck) - if err != nil { - return "", err - } - if rel == ".." || strings.HasPrefix(rel, ".."+string(os.PathSeparator)) { - return "", fmt.Errorf("path escapes working directory") - } - - return resolved, nil + return canonicalPath(base) } func canonicalPath(path string) (string, error) { @@ -346,7 +421,7 @@ func readSelectionFromArgs(args map[string]any) (readSelection, error) { } func readLineCountArg(args map[string]any) (int, bool, error) { - for _, key := range []string{"line_count", "num_lines", "lines"} { + for _, key := range []string{"line_count", "num_lines"} { value, ok, err := intReadArg(args, key) if err != nil || ok { if ok && value < 1 { @@ -404,13 +479,7 @@ func parseLineRange(value string) (int, int, error) { return start, end, nil } -func readLineSelection(path string, selection readSelection) (string, error) { - file, err := os.Open(path) - if err != nil { - return "", err - } - defer file.Close() - +func readLineSelection(file *os.File, selection readSelection) (string, error) { reader := bufio.NewReader(file) var b strings.Builder for lineNo := 1; ; lineNo++ { diff --git a/agent/tools/file_test.go b/agent/tools/file_test.go index 362c80f8f..433446d6c 100644 --- a/agent/tools/file_test.go +++ b/agent/tools/file_test.go @@ -73,6 +73,37 @@ func TestEditRejectsEscapingPath(t *testing.T) { } } +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 := NewEdit().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 TestReadRejectsParentOutsideCurrentWorkingDir(t *testing.T) { root := t.TempDir() subdir := filepath.Join(root, "sub") @@ -130,6 +161,24 @@ func TestReadLineRange(t *testing.T) { } } +func TestReadLinesAliasAsRange(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 := NewRead().Execute(context.Background(), agent.ToolContext{WorkingDir: dir}, map[string]any{ + "path": "note.txt", + "lines": "2-3", + }) + if err != nil { + t.Fatal(err) + } + if result.Content != "two\nthree\n" { + t.Fatalf("content = %q", result.Content) + } +} + func TestReadLineCountFromStartLine(t *testing.T) { dir := t.TempDir() if err := os.WriteFile(filepath.Join(dir, "note.txt"), []byte("one\ntwo\nthree\nfour\n"), 0o644); err != nil { diff --git a/cmd/cmd.go b/cmd/cmd.go index 207efcb61..64d827547 100644 --- a/cmd/cmd.go +++ b/cmd/cmd.go @@ -1948,6 +1948,7 @@ var ( agentOnboardingPrompt = tui.RunAgentSignInOnboarding agentOnboardingSignIn = runAgentOnboardingSignIn agentOnboardingSignedInStatus = runAgentOnboardingSignedInStatus + errAgentOnboardingNotSignedIn = errors.New("not signed in") ) func runRootResume(rootCmd, runCmd *cobra.Command, args []string) error { @@ -2294,6 +2295,9 @@ func runAgentOnboardingSignIn(ctx context.Context) error { if err == nil && user != nil && user.Name != "" { return nil } + if err == nil { + return errAgentOnboardingNotSignedIn + } var authErr api.AuthorizationError if !errors.As(err, &authErr) || authErr.SigninURL == "" { diff --git a/cmd/cmd_launcher_test.go b/cmd/cmd_launcher_test.go index 562ae8747..33dfa8efa 100644 --- a/cmd/cmd_launcher_test.go +++ b/cmd/cmd_launcher_test.go @@ -3,6 +3,8 @@ package cmd import ( "context" "errors" + "net/http" + "net/http/httptest" "testing" "github.com/spf13/cobra" @@ -226,6 +228,24 @@ func TestMaybeRunAgentOnboarding(t *testing.T) { }) } +func TestRunAgentOnboardingSignInEmptyWhoamiDoesNotSilentlySucceed(t *testing.T) { + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + if r.URL.Path != "/api/me" { + http.NotFound(w, r) + return + } + w.Header().Set("Content-Type", "application/json") + _, _ = w.Write([]byte(`{}`)) + })) + defer server.Close() + t.Setenv("OLLAMA_HOST", server.URL) + + err := runAgentOnboardingSignIn(context.Background()) + if !errors.Is(err, errAgentOnboardingNotSignedIn) { + t.Fatalf("error = %v, want errAgentOnboardingNotSignedIn", err) + } +} + func TestRunInteractiveTUI_RunModelActionsUseResolveRunModel(t *testing.T) { tests := []struct { name string diff --git a/cmd/internal/filedata/filedata.go b/cmd/internal/filedata/filedata.go index 523e08e96..9b09b06bf 100644 --- a/cmd/internal/filedata/filedata.go +++ b/cmd/internal/filedata/filedata.go @@ -145,14 +145,12 @@ func normalizeFileURL(u *url.URL) string { if unescaped, err := url.PathUnescape(path); err == nil { path = unescaped } - if path == "" && u.Host != "" { - host := u.Host - if unescaped, err := url.PathUnescape(host); err == nil { - host = unescaped - } - if len(host) >= 2 && host[1] == ':' && isASCIIAlpha(host[0]) { - return filepath.Clean(host) - } + host := u.Host + if unescaped, err := url.PathUnescape(host); err == nil { + host = unescaped + } + if len(host) >= 2 && host[1] == ':' && isASCIIAlpha(host[0]) { + return filepath.Clean(filepath.FromSlash(host + path)) } if len(path) >= 4 && path[0] == '/' && path[2] == ':' && isASCIIAlpha(path[1]) { path = path[1:] diff --git a/cmd/internal/filedata/filedata_test.go b/cmd/internal/filedata/filedata_test.go index 8cbf37fba..98e1fc729 100644 --- a/cmd/internal/filedata/filedata_test.go +++ b/cmd/internal/filedata/filedata_test.go @@ -12,3 +12,19 @@ func TestNormalizePathMalformedWindowsFileURL(t *testing.T) { t.Fatalf("path = %q, want %q", got, want) } } + +func TestNormalizePathTwoSlashWindowsFileURL(t *testing.T) { + got := NormalizePath(`file://C:/Users/jdoe/Pictures/img.png`) + want := filepath.Clean(`C:/Users/jdoe/Pictures/img.png`) + if got != want { + t.Fatalf("path = %q, want %q", got, want) + } +} + +func TestNormalizePathLocalhostWindowsFileURL(t *testing.T) { + got := NormalizePath(`file://localhost/C:/Users/jdoe/Pictures/img.png`) + want := filepath.Clean(`C:/Users/jdoe/Pictures/img.png`) + if got != want { + t.Fatalf("path = %q, want %q", got, want) + } +} diff --git a/cmd/tui/chat.go b/cmd/tui/chat.go index f13a6e30a..e8116d69c 100644 --- a/cmd/tui/chat.go +++ b/cmd/tui/chat.go @@ -130,6 +130,7 @@ type chatModel struct { height int status string spinner int + tickActive bool complete int systemPromptDisabled bool quitting bool @@ -216,11 +217,13 @@ func (m chatModel) Update(msg tea.Msg) (tea.Model, tea.Cmd) { return m, tea.ClearScreen case chatTickMsg: + m.tickActive = false if !m.running && !m.compacting { return m, nil } m.spinner++ - return m, chatTickCmd() + cmd := m.scheduleTick() + return m, cmd case chatAgentMsg: m.applyAgentEvent(msg.event) @@ -634,14 +637,17 @@ func (m chatModel) View() string { available = 0 } - transcriptLines := m.visibleTranscriptLines(width, available) - if len(transcriptLines) == 0 { + allTranscriptLines := m.transcriptLines(width) + transcriptLines := m.visibleTranscriptLinesForLines(allTranscriptLines, available) + if len(allTranscriptLines) == 0 { if available > 0 { transcriptLines = []string{chatMetaStyle.Render(m.emptyChatHint())} } } - for len(transcriptLines) < available { - transcriptLines = append(transcriptLines, "") + if len(allTranscriptLines) > available { + for len(transcriptLines) < available { + transcriptLines = append(transcriptLines, "") + } } lines := append(headerLines, transcriptLines...) @@ -760,9 +766,13 @@ func (m *chatModel) startManualCompaction() (tea.Model, tea.Cmd) { go func() { defer close(events) result, err := compactor.MaybeCompact(runCtx, req) - events <- chatCompactDoneMsg{result: result, err: err} + select { + case events <- chatCompactDoneMsg{result: result, err: err}: + case <-runCtx.Done(): + } }() - return *m, tea.Batch(waitForChatMsg(events), chatTickCmd()) + tickCmd := m.scheduleTick() + return *m, tea.Batch(waitForChatMsg(events), tickCmd) } func (m chatModel) finishManualCompaction(msg chatCompactDoneMsg) (tea.Model, tea.Cmd) { @@ -911,10 +921,14 @@ func (m *chatModel) startRunWithMessages(displayInput string, newMessages []api. go func() { defer close(events) result, err := session.Run(runCtx, opts) - events <- chatRunDoneMsg{result: result, err: err} + select { + case events <- chatRunDoneMsg{result: result, err: err}: + case <-runCtx.Done(): + } }() - return *m, tea.Batch(waitForChatMsg(events), chatTickCmd()) + tickCmd := m.scheduleTick() + return *m, tea.Batch(waitForChatMsg(events), tickCmd) } func (m *chatModel) startNextQueued() tea.Cmd { diff --git a/cmd/tui/chat_events.go b/cmd/tui/chat_events.go index c3f9163fa..06496c9d5 100644 --- a/cmd/tui/chat_events.go +++ b/cmd/tui/chat_events.go @@ -226,6 +226,14 @@ func waitForChatMsg(ch <-chan tea.Msg) tea.Cmd { } } +func (m *chatModel) scheduleTick() tea.Cmd { + if m.tickActive { + return nil + } + m.tickActive = true + return chatTickCmd() +} + func chatTickCmd() tea.Cmd { return tea.Tick(120*time.Millisecond, func(time.Time) tea.Msg { return chatTickMsg{} diff --git a/cmd/tui/chat_input.go b/cmd/tui/chat_input.go index 40d4645c2..819a66c2e 100644 --- a/cmd/tui/chat_input.go +++ b/cmd/tui/chat_input.go @@ -445,7 +445,7 @@ func (m *chatModel) deleteInputWordBackward() { if cursor <= 0 { return } - start, end, ok := m.placeholderRangeForBackspace(cursor) + start, end, ok := m.placeholderRangeForWordDelete(cursor) if !ok { start, end = previousInputWordStart(m.input, cursor), cursor } @@ -477,6 +477,25 @@ func (m chatModel) placeholderRangeForBackspace(cursor int) (int, int, bool) { return 0, 0, false } +func (m chatModel) placeholderRangeForWordDelete(cursor int) (int, int, bool) { + cursor = clamp(cursor, 0, len(m.input)) + end := cursor + for end > 0 && unicode.IsSpace(m.input[end-1]) { + end-- + } + input := string(m.input) + for _, placeholder := range m.inputPlaceholders() { + if placeholder == "" { + continue + } + start, placeholderEnd, ok := inputPlaceholderRuneRange(input, placeholder) + if ok && end > start && end <= placeholderEnd { + return start, cursor, true + } + } + return 0, 0, false +} + func (m chatModel) inputPlaceholders() []string { placeholders := make([]string, 0, len(m.inputAttachments)+len(m.inputPastedTexts)) for _, attachment := range m.inputAttachments { @@ -1253,7 +1272,11 @@ func (m *chatModel) importSkills(args []string) string { continue } imported++ - lines = append(lines, "imported "+name) + line := "imported " + name + if result.Error != "" { + line += " (" + result.Error + ")" + } + lines = append(lines, line) } lines = append(lines, "", fmt.Sprintf("%d imported, %d skipped", imported, skipped)) return strings.Join(lines, "\n") diff --git a/cmd/tui/chat_input_test.go b/cmd/tui/chat_input_test.go index 37e28c146..86a544b10 100644 --- a/cmd/tui/chat_input_test.go +++ b/cmd/tui/chat_input_test.go @@ -161,6 +161,27 @@ func TestChatBackspaceDeletesWholePastedTextPlaceholder(t *testing.T) { } } +func TestChatWordBackspaceDeletesWholePastedTextPlaceholder(t *testing.T) { + pasted := strings.Repeat("long paste line\n", pastedTextPlaceholderMinLines) + m := chatModel{input: []rune("use ")} + m.inputCursor = len(m.input) + m.inputCursorSet = true + + updated, _ := m.Update(tea.KeyMsg{Type: tea.KeyRunes, Runes: []rune(pasted), Paste: true}) + m = updated.(chatModel) + updated, _ = m.Update(tea.KeyMsg{Type: tea.KeySpace}) + m = updated.(chatModel) + updated, _ = m.Update(tea.KeyMsg{Type: tea.KeyBackspace, Alt: true}) + m = updated.(chatModel) + + if got := string(m.input); got != "use " { + t.Fatalf("input after word backspace = %q, want pasted text placeholder removed", got) + } + if got := len(m.inputPastedTexts); got != 0 { + t.Fatalf("pasted texts after word backspace = %d, want 0", got) + } +} + func TestChatEditorResultReplacesInput(t *testing.T) { m := chatModel{ input: []rune("old"), diff --git a/cmd/tui/chat_multimodal_test.go b/cmd/tui/chat_multimodal_test.go index d4b0a4184..9d9fb7dfb 100644 --- a/cmd/tui/chat_multimodal_test.go +++ b/cmd/tui/chat_multimodal_test.go @@ -221,6 +221,32 @@ func TestChatDeletingImagePlaceholderRemovesAttachment(t *testing.T) { } } +func TestChatWordDeletingImagePlaceholderRemovesAttachment(t *testing.T) { + fp := writeTestPNG(t) + m := chatModel{ + ctx: context.Background(), + opts: ChatOptions{ + Model: "test", + Client: chatTestClient{}, + MultiModal: true, + }, + } + + updated, _ := m.Update(tea.KeyMsg{Type: tea.KeyRunes, Runes: []rune("describe " + fp), Paste: true}) + m = updated.(chatModel) + updated, _ = m.Update(tea.KeyMsg{Type: tea.KeySpace}) + m = updated.(chatModel) + updated, _ = m.Update(tea.KeyMsg{Type: tea.KeyBackspace, Alt: true}) + m = updated.(chatModel) + + if got := string(m.input); got != "describe " { + t.Fatalf("input after word backspace = %q, want image placeholder removed", got) + } + if got := len(m.inputAttachments); got != 0 { + t.Fatalf("input attachments after word backspace = %d, want 0", got) + } +} + func writeTestPNG(t *testing.T) string { t.Helper() dir := t.TempDir() diff --git a/cmd/tui/chat_render.go b/cmd/tui/chat_render.go index 34c8c90a0..b393f0554 100644 --- a/cmd/tui/chat_render.go +++ b/cmd/tui/chat_render.go @@ -227,15 +227,19 @@ func (m chatModel) transcriptLines(width int) []string { } func (m chatModel) visibleTranscriptLines(width, available int) []string { + return m.visibleTranscriptLinesForLines(m.transcriptLines(width), available) +} + +func (m chatModel) visibleTranscriptLinesForLines(lines []string, available int) []string { if available <= 0 { return nil } - lines := m.transcriptLines(width) + start := 0 if len(lines) > available { - start := m.visibleTranscriptStartLineForLines(len(lines), available) + start = m.visibleTranscriptStartLineForLines(len(lines), available) lines = lines[start : start+available] } - lines = m.applyTranscriptSelection(lines, width, available) + lines = m.applyTranscriptSelection(lines, start) return lines } @@ -268,10 +272,11 @@ func (m chatModel) transcriptHeight() int { if height <= 0 { height = 24 } + lineCount := len(m.transcriptLines(width)) baseHeaderHeight := 2 baseHeight := max(0, height-baseHeaderHeight-len(m.bottomLines(width, height-baseHeaderHeight))) - if len(m.transcriptLines(width)) <= baseHeight { - return baseHeight + if lineCount <= baseHeight { + return lineCount } statusHeaderHeight := 3 return max(0, height-statusHeaderHeight-len(m.bottomLines(width, height-statusHeaderHeight))) @@ -375,12 +380,11 @@ func normalizedSelectionRangeFor(selection chatSelection) (chatSelectionPoint, c return start, end, true } -func (m chatModel) applyTranscriptSelection(lines []string, width, available int) []string { +func (m chatModel) applyTranscriptSelection(lines []string, offset int) []string { start, end, ok := m.normalizedSelectionRange() if !ok || len(lines) == 0 { return lines } - offset := m.visibleTranscriptStartLine(width, available) out := slices.Clone(lines) for i, line := range out { lineIndex := offset + i diff --git a/cmd/tui/chat_render_test.go b/cmd/tui/chat_render_test.go index 518525e41..b7d551a8c 100644 --- a/cmd/tui/chat_render_test.go +++ b/cmd/tui/chat_render_test.go @@ -377,6 +377,49 @@ func TestChatViewPadsToTerminalHeight(t *testing.T) { } } +func TestChatInputFloatsUntilTranscriptOverflows(t *testing.T) { + short := chatModel{ + input: []rune("next"), + width: 60, + height: 14, + entries: []chatEntry{ + {role: "assistant", content: "hello"}, + }, + } + shortInputLine := renderedInputLine(short.View()) + if shortInputLine < 0 { + t.Fatalf("short view missing input:\n%s", stripANSI(short.View())) + } + + long := chatModel{ + input: []rune("next"), + width: short.width, + height: short.height, + } + for i := range 20 { + long.entries = append(long.entries, chatEntry{role: "assistant", content: fmt.Sprintf("line %02d", i)}) + } + longInputLine := renderedInputLine(long.View()) + if longInputLine < 0 { + t.Fatalf("long view missing input:\n%s", stripANSI(long.View())) + } + if shortInputLine >= longInputLine { + t.Fatalf("input line did not move down as transcript grew: short=%d long=%d\nshort:\n%s\nlong:\n%s", shortInputLine, longInputLine, stripANSI(short.View()), stripANSI(long.View())) + } + if longInputLine < long.height-4 { + t.Fatalf("overflowing transcript should pin input near bottom, input line=%d height=%d\n%s", longInputLine, long.height, stripANSI(long.View())) + } +} + +func renderedInputLine(view string) int { + for i, line := range strings.Split(stripANSI(view), "\n") { + if strings.Contains(line, "> next") { + return i + } + } + return -1 +} + func TestChatScrollsTranscript(t *testing.T) { m := chatModel{ width: 80, diff --git a/cmd/tui/tui.go b/cmd/tui/tui.go index 7dae825fd..5293b2f17 100644 --- a/cmd/tui/tui.go +++ b/cmd/tui/tui.go @@ -48,7 +48,7 @@ type menuItem struct { const pinnedIntegrationCount = 4 var runModelMenuItem = menuItem{ - title: "Agent", + title: "Chat and Code", description: "Ollama's built-in agent to chat, code, and do work", isRunModel: true, } diff --git a/progress/progress.go b/progress/progress.go index d12145f4e..328dba650 100644 --- a/progress/progress.go +++ b/progress/progress.go @@ -28,6 +28,7 @@ type Progress struct { pos int done chan struct{} + exited chan struct{} stopOnce sync.Once ticker *time.Ticker states []State @@ -38,6 +39,7 @@ func NewProgress(w io.Writer) *Progress { p := &Progress{ w: bufio.NewWriter(w), done: make(chan struct{}), + exited: make(chan struct{}), ticker: ticker, } go p.start(ticker) @@ -64,6 +66,7 @@ func (p *Progress) stop() bool { if ticker != nil { ticker.Stop() + <-p.exited p.render() return true } @@ -74,6 +77,8 @@ func (p *Progress) stop() bool { func (p *Progress) Stop() bool { stopped := p.stop() if stopped { + p.mu.Lock() + defer p.mu.Unlock() fmt.Fprint(p.w, "\n") p.w.Flush() } @@ -81,17 +86,17 @@ func (p *Progress) Stop() bool { } func (p *Progress) StopAndClear() bool { - defer p.w.Flush() - - fmt.Fprint(p.w, "\033[?25l") - defer fmt.Fprint(p.w, "\033[?25h") - stopped := p.stop() if stopped { - // clear all progress lines p.mu.Lock() + defer p.mu.Unlock() + defer p.w.Flush() + + fmt.Fprint(p.w, "\033[?25l") + defer fmt.Fprint(p.w, "\033[?25h") + + // clear all progress lines pos := p.pos - p.mu.Unlock() for i := range pos { if i > 0 { fmt.Fprint(p.w, "\033[A") @@ -147,6 +152,7 @@ func (p *Progress) render() { } func (p *Progress) start(ticker *time.Ticker) { + defer close(p.exited) for { select { case <-p.done: diff --git a/progress/progress_test.go b/progress/progress_test.go new file mode 100644 index 000000000..823005124 --- /dev/null +++ b/progress/progress_test.go @@ -0,0 +1,42 @@ +package progress + +import ( + "bytes" + "sync" + "testing" +) + +type progressTestState string + +func (s progressTestState) String() string { + return string(s) +} + +func TestProgressStopIsIdempotent(t *testing.T) { + var buf bytes.Buffer + p := NewProgress(&buf) + p.Add("one", progressTestState("working")) + + if !p.Stop() { + t.Fatal("first stop should report stopped") + } + if p.Stop() { + t.Fatal("second stop should report already stopped") + } +} + +func TestProgressConcurrentStopAndClear(t *testing.T) { + var buf bytes.Buffer + p := NewProgress(&buf) + p.Add("one", progressTestState("working")) + + var wg sync.WaitGroup + for range 8 { + wg.Add(1) + go func() { + defer wg.Done() + p.StopAndClear() + }() + } + wg.Wait() +}